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