use crate::{
header::{HeaderKey, HeaderValue},
Headers, StatusCode,
};
#[derive(Debug)]
enum ProtocolData<'a> {
Ref(&'a str),
Owned(String),
}
impl<'a> AsRef<str> for ProtocolData<'a> {
fn as_ref(&self) -> &str {
match self {
Self::Ref(tmp) => tmp,
Self::Owned(tmp) => &tmp,
}
}
}
impl<'a> PartialEq for ProtocolData<'a> {
fn eq(&self, other: &ProtocolData<'_>) -> bool {
self.as_ref().eq(other.as_ref())
}
}
#[derive(Debug, PartialEq)]
pub struct Response<'a> {
status_code: StatusCode,
protocol: ProtocolData<'a>,
headers: Headers<'a>,
body: Vec<u8>,
}
impl<'a> Response<'a> {
pub fn new(
protocol: &'a str,
status_code: StatusCode,
headers: Headers<'a>,
body: Vec<u8>,
) -> Self {
Self {
status_code,
protocol: ProtocolData::Ref(protocol),
headers,
body,
}
}
pub(crate) fn new_owned(
protocol: String,
status_code: StatusCode,
headers: Headers<'a>,
body: Vec<u8>,
) -> Self {
Self {
status_code,
protocol: ProtocolData::Owned(protocol),
headers,
body,
}
}
pub fn serialize(&self) -> (Vec<u8>, &[u8]) {
let protocol = self.protocol.as_ref();
let status_code = self.status_code.serialize();
let capacity = protocol.len() + 1 + status_code.len() + 4;
let mut result = Vec::with_capacity(capacity);
result.extend_from_slice(protocol.as_bytes());
result.push(b' ');
result.extend_from_slice(status_code.as_bytes());
result.extend_from_slice("\r\n".as_bytes());
self.headers.serialize(&mut result);
result.extend_from_slice("\r\n".as_bytes());
(result, &self.body)
}
pub fn protocol(&self) -> &str {
self.protocol.as_ref()
}
pub fn status_code(&self) -> &StatusCode {
&self.status_code
}
pub fn headers(&self) -> &Headers<'a> {
&self.headers
}
pub fn body(&self) -> &[u8] {
&self.body
}
pub fn add_header<'b, K, V>(&mut self, key: K, value: V)
where
'b: 'a,
K: Into<HeaderKey<'a>>,
V: Into<HeaderValue<'a>>,
{
self.headers.set(key, value);
}
pub fn set_body(&mut self, n_body: Vec<u8>) {
self.body = n_body;
self.add_header("Content-Length", self.body.len());
}
pub fn is_chunked(&self) -> bool {
match self.headers.get("Transfer-Encoding") {
None => false,
Some(value) => value.eq_ignore_case(&HeaderValue::StrRef("Chunked")),
}
}
pub fn to_owned<'refed, 'owned>(&'refed self) -> Response<'owned> {
Response::new_owned(
self.protocol.as_ref().to_owned(),
self.status_code.clone(),
self.headers.to_owned(),
self.body.clone(),
)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn serialize_valid() {
let mut headers = Headers::new();
headers.set("test-1", "value-1");
let req = Response::new(
"HTTP/1.1",
StatusCode::OK,
headers,
"body".as_bytes().to_vec(),
);
let raw_resp_header = "HTTP/1.1 200 OK\r\ntest-1: value-1\r\n\r\n";
let resp_header = raw_resp_header.as_bytes().to_vec();
let resp_body = "body".as_bytes();
assert_eq!(req.serialize(), (resp_header, resp_body));
}
#[test]
fn serialize_valid_no_body() {
let mut headers = Headers::new();
headers.set("test-1", "value-1");
let req = Response::new("HTTP/1.1", StatusCode::OK, headers, "".as_bytes().to_vec());
let raw_resp_header = "HTTP/1.1 200 OK\r\ntest-1: value-1\r\n\r\n";
let resp_header = raw_resp_header.as_bytes().to_vec();
let resp_body = "".as_bytes();
assert_eq!(req.serialize(), (resp_header, resp_body));
}
#[test]
fn is_chunked_not_set() {
let mut headers = Headers::new();
headers.set("test-1", "value-1");
let resp = Response::new("HTTP/1.1", StatusCode::OK, headers, "".as_bytes().to_vec());
assert_eq!(false, resp.is_chunked());
}
#[test]
fn is_chunked_set() {
let mut headers = Headers::new();
headers.set("Transfer-Encoding", "Chunked");
let resp = Response::new("HTTP/1.1", StatusCode::OK, headers, "".as_bytes().to_vec());
assert_eq!(true, resp.is_chunked());
}
#[test]
fn is_chunked_set_differently() {
let mut headers = Headers::new();
headers.set("Transfer-Encoding", "compress");
let resp = Response::new("HTTP/1.1", StatusCode::OK, headers, "".as_bytes().to_vec());
assert_eq!(false, resp.is_chunked());
}
#[test]
fn to_owned() {
let resp = Response::new("HTTP/1.1", StatusCode::OK, Headers::new(), Vec::new());
let cloned = resp.to_owned();
drop(resp);
assert_eq!(&StatusCode::OK, cloned.status_code())
}
}