sark_core/http/codec/decode/
head.rs

1use http::{HeaderName, HeaderValue, StatusCode};
2
3use crate::error::{Error, Result};
4use crate::http::codec::Header;
5use crate::http::codec::decode::HeaderScan;
6use crate::http::codec::decode::chunked::BodyDecoder;
7use crate::http::{Body, Response};
8
9const MAX_HEADERS: usize = 100;
10
11#[derive(Clone, Copy, Debug, PartialEq, Eq)]
12pub enum DecodeMode {
13    Response,
14    Head,
15}
16
17impl DecodeMode {
18    const fn is_head(self) -> bool {
19        matches!(self, Self::Head)
20    }
21}
22
23#[derive(Clone, Copy, Debug, PartialEq, Eq)]
24pub enum BodyKind {
25    NoBody,
26    ContentLength(usize),
27    Chunked,
28    UntilEof,
29}
30
31pub struct DecodedHead {
32    pub status: StatusCode,
33    pub headers: Vec<(HeaderName, HeaderValue)>,
34    pub header_len: usize,
35    pub body_kind: BodyKind,
36}
37
38impl DecodedHead {
39    pub fn into_response<B, I>(self, body: B, trailers: I) -> Response
40    where
41        B: Into<Body<'static>>,
42        I: IntoIterator<Item = (HeaderName, HeaderValue)>,
43    {
44        let mut response = Response::new(self.status);
45        for (name, value) in self.headers {
46            response.headers_mut().insert(name, value);
47        }
48        response.headers_mut().extend_trailers(trailers);
49        response.set_body(body);
50        response
51    }
52}
53
54pub struct ResponseDecoder {
55    mode: DecodeMode,
56}
57
58impl ResponseDecoder {
59    pub const fn new(mode: DecodeMode) -> Self {
60        Self { mode }
61    }
62
63    fn status_has_no_body(status: StatusCode) -> bool {
64        let code = status.as_u16();
65        code < 200 || code == 204 || code == 304
66    }
67
68    fn parse(&self, buf: &[u8]) -> Result<Option<DecodedHead>> {
69        let mut raw = [httparse::EMPTY_HEADER; MAX_HEADERS];
70        let mut parsed = httparse::Response::new(&mut raw);
71
72        match parsed.parse(buf)? {
73            httparse::Status::Partial => Ok(None),
74            httparse::Status::Complete(header_len) => {
75                let status_code = parsed
76                    .code
77                    .ok_or_else(|| Error::BadRequest("Missing status code".into()))?;
78                let status = StatusCode::from_u16(status_code)
79                    .map_err(|_| Error::BadRequest("Invalid status code".into()))?;
80
81                let mut headers = Vec::with_capacity(parsed.headers.len());
82                let mut content_length = None;
83                let mut is_chunked = false;
84                let mut has_transfer_encoding = false;
85
86                for h in parsed.headers.iter() {
87                    let name = HeaderName::from_bytes(h.name.as_bytes()).map_err(|_| {
88                        Error::BadRequest(format!("Invalid header name: {}", h.name).into())
89                    })?;
90                    let value = HeaderValue::from_bytes(h.value).map_err(|_| {
91                        Error::BadRequest(format!("Invalid header value for: {}", h.name).into())
92                    })?;
93
94                    if name == "content-length" {
95                        let len = Header::content_length(value.as_bytes())?;
96                        if let Some(existing) = content_length
97                            && existing != len
98                        {
99                            return Err(Error::BadRequest("Conflicting Content-Length".into()));
100                        }
101                        content_length = Some(len);
102                    }
103                    if name == "transfer-encoding" {
104                        let te = HeaderScan::transfer_encoding(value.as_bytes())?;
105                        has_transfer_encoding = has_transfer_encoding || te.has_transfer_encoding;
106                        if te.is_chunked {
107                            is_chunked = true;
108                        }
109                    }
110                    headers.push((name, value));
111                }
112
113                if has_transfer_encoding && content_length.is_some() {
114                    return Err(Error::BadRequest(
115                        "Content-Length with Transfer-Encoding is not allowed".into(),
116                    ));
117                }
118
119                let body_kind = if self.mode.is_head() || Self::status_has_no_body(status) {
120                    BodyKind::NoBody
121                } else if is_chunked {
122                    BodyKind::Chunked
123                } else if let Some(len) = content_length {
124                    BodyKind::ContentLength(len)
125                } else {
126                    BodyKind::UntilEof
127                };
128
129                Ok(Some(DecodedHead {
130                    status,
131                    headers,
132                    header_len,
133                    body_kind,
134                }))
135            }
136        }
137    }
138
139    pub fn head(&self, buf: &[u8]) -> Result<Option<DecodedHead>> {
140        self.parse(buf)
141    }
142
143    pub fn response(&self, buf: &[u8]) -> Result<Option<Response>> {
144        let head = match self.parse(buf)? {
145            Some(h) => h,
146            None => return Ok(None),
147        };
148        let body_data = &buf[head.header_len..];
149
150        match head.body_kind {
151            BodyKind::NoBody => Ok(Some(head.into_response(&[][..], []))),
152            BodyKind::Chunked => {
153                match BodyDecoder::decode_all(body_data, BodyDecoder::DEFAULT_MAX_BODY)? {
154                    Some(result) => Ok(Some(head.into_response(result.body, result.trailers))),
155                    None => Ok(None),
156                }
157            }
158            BodyKind::ContentLength(len) if body_data.len() < len => Ok(None),
159            BodyKind::ContentLength(len) => Ok(Some(head.into_response(&body_data[..len], []))),
160            BodyKind::UntilEof => Ok(None),
161        }
162    }
163
164    pub fn response_after_eof(&self, buf: &[u8]) -> Result<Response> {
165        let head = self
166            .parse(buf)?
167            .ok_or_else(|| Error::BadRequest("Incomplete HTTP response".into()))?;
168        let body_data = &buf[head.header_len..];
169
170        match head.body_kind {
171            BodyKind::Chunked => {
172                let result = BodyDecoder::decode_all(body_data, BodyDecoder::DEFAULT_MAX_BODY)?
173                    .ok_or_else(|| Error::BadRequest("Incomplete chunked response".into()))?;
174                Ok(head.into_response(result.body, result.trailers))
175            }
176            BodyKind::ContentLength(len) if body_data.len() < len => {
177                Err(Error::BadRequest("Incomplete HTTP response body".into()))
178            }
179            BodyKind::ContentLength(len) => Ok(head.into_response(&body_data[..len], [])),
180            BodyKind::NoBody => Ok(head.into_response(&[][..], [])),
181            BodyKind::UntilEof => Ok(head.into_response(body_data, [])),
182        }
183    }
184}
sark_core/http/codec/decode/head.rs

sark_core/http/codec/decode/
head.rs
