sark_core/http/codec/decode/
head.rs

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

sark_core/http/codec/decode/
head.rs
