Skip to main content

sark_core/http/codec/decode/
head.rs

1use http::{HeaderName, HeaderValue, StatusCode};
2use o3::buffer::Owned;
3
4use crate::error::{Error, Result};
5use crate::http::Response;
6use crate::http::codec::Header;
7
8const MAX_HEADERS: usize = 100;
9
10#[derive(Clone, Copy, Debug, PartialEq, Eq)]
11pub enum DecodeMode {
12    Response,
13    Head,
14}
15
16impl DecodeMode {
17    const fn is_head(self) -> bool {
18        matches!(self, Self::Head)
19    }
20}
21
22pub(super) struct ParsedHead {
23    pub(super) status: StatusCode,
24    pub(super) headers: Vec<(HeaderName, HeaderValue)>,
25    pub(super) header_len: usize,
26    pub(super) content_length: Option<usize>,
27    pub(super) is_chunked: bool,
28}
29
30pub enum BodyKind {
31    NoBody,
32    ContentLength(usize),
33    Chunked,
34    UntilEof,
35}
36
37pub struct DecodedHead {
38    pub status: StatusCode,
39    pub headers: Vec<(HeaderName, HeaderValue)>,
40    pub header_len: usize,
41    pub body_kind: BodyKind,
42}
43
44impl crate::http::codec::Parse {
45    pub fn status_has_no_body(status: StatusCode) -> bool {
46        let code = status.as_u16();
47        code < 200 || code == 204 || code == 304
48    }
49
50    pub(super) fn parse(buf: &[u8]) -> Result<Option<ParsedHead>> {
51        let mut raw = [httparse::EMPTY_HEADER; MAX_HEADERS];
52        let mut parsed = httparse::Response::new(&mut raw);
53
54        match parsed.parse(buf)? {
55            httparse::Status::Partial => Ok(None),
56            httparse::Status::Complete(header_len) => {
57                let status_code = parsed
58                    .code
59                    .ok_or_else(|| Error::BadRequest("Missing status code".into()))?;
60                let status = StatusCode::from_u16(status_code)
61                    .map_err(|_| Error::BadRequest("Invalid status code".into()))?;
62
63                let mut headers = Vec::with_capacity(parsed.headers.len());
64                let mut content_length = None;
65                let mut is_chunked = false;
66                let mut has_transfer_encoding = false;
67
68                for h in parsed.headers.iter() {
69                    let name = HeaderName::from_bytes(h.name.as_bytes()).map_err(|_| {
70                        Error::BadRequest(format!("Invalid header name: {}", h.name).into())
71                    })?;
72                    let value = HeaderValue::from_bytes(h.value).map_err(|_| {
73                        Error::BadRequest(format!("Invalid header value for: {}", h.name).into())
74                    })?;
75
76                    if name == "content-length" {
77                        let len = Header::content_length(value.as_bytes())?;
78                        if let Some(existing) = content_length
79                            && existing != len
80                        {
81                            return Err(Error::BadRequest("Conflicting Content-Length".into()));
82                        }
83                        content_length = Some(len);
84                    }
85                    if name == "transfer-encoding" {
86                        let te = Self::parse_transfer_encoding_value(value.as_bytes())?;
87                        has_transfer_encoding = has_transfer_encoding || te.has_transfer_encoding;
88                        if te.is_chunked {
89                            is_chunked = true;
90                        }
91                    }
92                    headers.push((name, value));
93                }
94
95                if has_transfer_encoding && content_length.is_some() {
96                    return Err(Error::BadRequest(
97                        "Content-Length with Transfer-Encoding is not allowed".into(),
98                    ));
99                }
100
101                Ok(Some(ParsedHead {
102                    status,
103                    headers,
104                    header_len,
105                    content_length,
106                    is_chunked,
107                }))
108            }
109        }
110    }
111
112    pub(super) fn build_response(
113        head: ParsedHead,
114        body: &[u8],
115        trailers: &[(HeaderName, HeaderValue)],
116    ) -> Response {
117        let mut resp = Response::new(head.status);
118        for (name, value) in head.headers {
119            resp.headers_mut().insert(name, value);
120        }
121        for (name, value) in trailers {
122            resp.headers_mut().insert(name.clone(), value.clone());
123        }
124        if !body.is_empty() {
125            resp.set_body(Owned::from(body));
126        }
127        resp
128    }
129
130    pub fn head(buf: &[u8], mode: DecodeMode) -> Result<Option<DecodedHead>> {
131        let head = match Self::parse(buf)? {
132            Some(h) => h,
133            None => return Ok(None),
134        };
135
136        let body_kind = if mode.is_head() || Self::status_has_no_body(head.status) {
137            BodyKind::NoBody
138        } else if head.is_chunked {
139            BodyKind::Chunked
140        } else if let Some(len) = head.content_length {
141            BodyKind::ContentLength(len)
142        } else {
143            BodyKind::UntilEof
144        };
145
146        Ok(Some(DecodedHead {
147            status: head.status,
148            headers: head.headers,
149            header_len: head.header_len,
150            body_kind,
151        }))
152    }
153
154    pub fn response(buf: &[u8], mode: DecodeMode) -> Result<Option<Response>> {
155        let head = match Self::parse(buf)? {
156            Some(h) => h,
157            None => return Ok(None),
158        };
159
160        if mode.is_head() || Self::status_has_no_body(head.status) {
161            return Ok(Some(Self::build_response(head, &[], &[])));
162        }
163
164        let body_data = &buf[head.header_len..];
165
166        if head.is_chunked {
167            match Self::try_decode_chunked(body_data)? {
168                Some(result) => Ok(Some(Self::build_response(
169                    head,
170                    &result.body,
171                    &result.trailers,
172                ))),
173                None => Ok(None),
174            }
175        } else {
176            match head.content_length {
177                Some(len) if body_data.len() < len => Ok(None),
178                Some(len) => Ok(Some(Self::build_response(head, &body_data[..len], &[]))),
179                None => Ok(None),
180            }
181        }
182    }
183
184    pub fn response_after_eof(buf: &[u8]) -> Result<Response> {
185        let head = Self::parse(buf)?
186            .ok_or_else(|| Error::BadRequest("Incomplete HTTP response".into()))?;
187        let body_data = &buf[head.header_len..];
188
189        if head.is_chunked {
190            let result = Self::try_decode_chunked(body_data)?
191                .ok_or_else(|| Error::BadRequest("Incomplete chunked response".into()))?;
192            Ok(Self::build_response(head, &result.body, &result.trailers))
193        } else {
194            Ok(Self::build_response(head, body_data, &[]))
195        }
196    }
197}
198
199#[cfg(test)]
200mod tests {
201    #[test]
202    fn cl_and_te_coexistence_rejected() {
203        let raw = b"HTTP/1.1 200 OK\r\nContent-Length: 5\r\nTransfer-Encoding: chunked\r\n\r\n";
204        assert!(crate::http::codec::Parse::parse(raw).is_err());
205    }
206
207    #[test]
208    fn valid_single_content_length_accepted() {
209        let raw = b"HTTP/1.1 200 OK\r\nContent-Length: 5\r\n\r\n";
210        let head = crate::http::codec::Parse::parse(raw).unwrap().unwrap();
211        assert_eq!(head.content_length, Some(5));
212        assert!(!head.is_chunked);
213    }
214
215    #[test]
216    fn valid_te_chunked_alone_accepted() {
217        let raw = b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n";
218        let head = crate::http::codec::Parse::parse(raw).unwrap().unwrap();
219        assert!(head.is_chunked);
220        assert_eq!(head.content_length, None);
221    }
222
223    #[test]
224    fn response_content_length_plus_prefix_rejected() {
225        let raw = b"HTTP/1.1 200 OK\r\nContent-Length: +5\r\n\r\n";
226        assert!(crate::http::codec::Parse::parse(raw).is_err());
227    }
228}