Skip to main content

xitca_http/h1/proto/
encode.rs

1use tracing::{error, warn};
2
3use crate::{
4    body::{Body, SizeHint},
5    bytes::{Bytes, BytesMut},
6    date::DateTime,
7    http::{
8        StatusCode,
9        header::{CONNECTION, CONTENT_LENGTH, DATE, HeaderMap, SET_COOKIE, TE, TRANSFER_ENCODING, UPGRADE},
10        response::Parts,
11    },
12};
13
14use super::{buf_write::H1BufWrite, context::Context, error::ProtoError, header, trasnder_coding::TransferCoding};
15
16pub const CONTINUE: &[u8; 25] = b"HTTP/1.1 100 Continue\r\n\r\n";
17
18#[allow(clippy::declare_interior_mutable_const)]
19pub const CONTINUE_BYTES: Bytes = Bytes::from_static(CONTINUE);
20
21impl<D, const MAX_HEADERS: usize> Context<'_, D, MAX_HEADERS>
22where
23    D: DateTime,
24{
25    pub fn encode_head<B, W>(&mut self, parts: Parts, body: &B, buf: &mut W) -> Result<TransferCoding, ProtoError>
26    where
27        B: Body,
28        W: H1BufWrite,
29    {
30        buf.write_buf_head(|buf| self.encode_head_inner(parts, body, buf))
31    }
32
33    fn encode_head_inner<B>(&mut self, parts: Parts, body: &B, buf: &mut BytesMut) -> Result<TransferCoding, ProtoError>
34    where
35        B: Body,
36    {
37        let Parts {
38            mut headers,
39            mut extensions,
40            version: _,
41            status,
42            ..
43        } = parts;
44
45        // decide if content-length or transfer-encoding header would be skipped.
46        let skip_ct_te = match status {
47            StatusCode::SWITCHING_PROTOCOLS => true,
48            // Sending content-length or transfer-encoding header on 2xx response
49            // to CONNECT is forbidden in RFC 7231.
50            s if self.is_connect_method() && s.is_success() => true,
51            s if s.is_informational() => {
52                error!("response with 1xx status code not supported");
53                return Err(ProtoError::Status);
54            }
55            _ => false,
56        };
57
58        self.encode_version_status_reason(buf, status);
59
60        let size = body.size_hint();
61
62        self.encode_headers(&mut headers, size, buf, skip_ct_te).inspect(|_| {
63            // put header map back to cache.
64            self.replace_headers(headers);
65
66            // put extension back to cache;
67            extensions.clear();
68            self.replace_extensions(extensions);
69        })
70    }
71
72    #[inline]
73    fn encode_version_status_reason(&self, buf: &mut BytesMut, status: StatusCode) {
74        // encode version, status code and reason
75        match (status, self.is_http10()) {
76            (StatusCode::OK, false) => {
77                buf.extend_from_slice(b"HTTP/1.1 200 OK");
78                return;
79            }
80            (_, false) => buf.extend_from_slice(b"HTTP/1.1 "),
81            (_, true) => buf.extend_from_slice(b"HTTP/1.0 "),
82        };
83
84        // a reason MUST be written, as many parsers will expect it.
85        let reason = status.canonical_reason().unwrap_or("<none>").as_bytes();
86        let status = status.as_str().as_bytes();
87        buf.reserve(status.len() + reason.len() + 1);
88        buf.extend_from_slice(status);
89        buf.extend_from_slice(b" ");
90        buf.extend_from_slice(reason);
91    }
92}
93
94impl<D, const MAX_HEADERS: usize> Context<'_, D, MAX_HEADERS>
95where
96    D: DateTime,
97{
98    pub fn encode_headers(
99        &mut self,
100        headers: &mut HeaderMap,
101        size: SizeHint,
102        buf: &mut BytesMut,
103        mut skip_ct_te: bool,
104    ) -> Result<TransferCoding, ProtoError> {
105        let mut skip_date = false;
106
107        // use the shortest header name as default
108        let mut name = TE;
109
110        let mut encoding = TransferCoding::eof();
111
112        for (next_name, value) in headers.drain() {
113            let mut is_multi_value = next_name
114                .map(|next_name| {
115                    name = next_name;
116                    false
117                })
118                .unwrap_or(true);
119
120            match name {
121                CONNECTION => {
122                    if self.is_connection_closed() {
123                        // skip write header on close condition.
124                        // the header is checked again and written properly afterwards.
125                        continue;
126                    }
127                    self.try_set_close_from_header(&value)?;
128                }
129                UPGRADE => encoding = TransferCoding::upgrade(),
130                DATE => skip_date = true,
131                CONTENT_LENGTH => {
132                    debug_assert!(!skip_ct_te, "CONTENT_LENGTH header can not be set");
133                    let value = header::parse_content_length(&value)?;
134                    encoding = TransferCoding::length(value);
135                    skip_ct_te = true;
136                }
137                TRANSFER_ENCODING => {
138                    debug_assert!(!skip_ct_te, "TRANSFER_ENCODING header can not be set");
139                    for val in value.to_str().map_err(|_| ProtoError::HeaderValue)?.split(',') {
140                        let val = val.trim();
141                        if val.eq_ignore_ascii_case("chunked") {
142                            encoding = TransferCoding::encode_chunked();
143                            skip_ct_te = true;
144                        }
145                    }
146                }
147                // multiple header lines for set-cookie header is allowed
148                // https://www.rfc-editor.org/rfc/rfc6265#section-3
149                SET_COOKIE => is_multi_value = false,
150                _ => {}
151            }
152
153            let value = value.as_bytes();
154
155            if is_multi_value {
156                buf.reserve(value.len() + 2);
157                buf.extend_from_slice(b", ");
158                buf.extend_from_slice(value);
159            } else {
160                let name = name.as_str().as_bytes();
161                buf.reserve(name.len() + value.len() + 4);
162                buf.extend_from_slice(b"\r\n");
163                buf.extend_from_slice(name);
164                buf.extend_from_slice(b": ");
165                buf.extend_from_slice(value);
166            }
167        }
168
169        // special handling for head method request by removing potential unwanted response body.
170        if self.is_head_method() {
171            try_remove_body(buf, skip_ct_te, size, &mut encoding);
172        // encode transfer-encoding or content-length if header map didn't provide them.
173        } else if !skip_ct_te {
174            encoding = match size {
175                SizeHint::None => TransferCoding::eof(),
176                SizeHint::Unknown => {
177                    buf.extend_from_slice(CHUNKED_HEADER);
178                    TransferCoding::encode_chunked()
179                }
180                SizeHint::Exact(size) => {
181                    write_length_header(buf, size);
182                    TransferCoding::length(size)
183                }
184            };
185        }
186
187        if self.is_connection_closed() {
188            buf.extend_from_slice(CLOSE_HEADER);
189        }
190
191        // set date header if there is not any.
192        if !skip_date {
193            buf.reserve(D::DATE_SIZE_HINT + 12);
194            buf.extend_from_slice(b"\r\ndate: ");
195            self.date().with_date(|slice| buf.extend_from_slice(slice));
196        }
197
198        buf.extend_from_slice(b"\r\n\r\n");
199
200        Ok(encoding)
201    }
202}
203
204const CHUNKED_HEADER: &[u8; 28] = b"\r\ntransfer-encoding: chunked";
205const CLOSE_HEADER: &[u8; 19] = b"\r\nconnection: close";
206
207#[cold]
208#[inline(never)]
209fn try_remove_body(buf: &mut BytesMut, skip_ct_te: bool, size: SizeHint, encoding: &mut TransferCoding) {
210    *encoding = TransferCoding::eof();
211
212    match size {
213        SizeHint::None => return,
214        SizeHint::Unknown if !skip_ct_te => {
215            buf.extend_from_slice(CHUNKED_HEADER);
216        }
217        SizeHint::Exact(size) if !skip_ct_te => {
218            write_length_header(buf, size);
219        }
220        _ => {}
221    }
222
223    warn!("response to HEAD request should not bearing body. It will been dropped without polling.");
224}
225
226pub(crate) fn write_length_header(buf: &mut BytesMut, size: u64) {
227    let mut buffer = core::fmt::NumBuffer::new();
228
229    let buffer = size.format_into(&mut buffer).as_bytes();
230
231    buf.reserve(buffer.len() + 18);
232    buf.extend_from_slice(b"\r\ncontent-length: ");
233    buf.extend_from_slice(buffer);
234}
235
236#[cfg(test)]
237mod test {
238    use crate::{
239        body::{BoxBody, Full},
240        date::SystemTimeDateTimeHandler,
241        http::{HeaderValue, Response},
242    };
243
244    use super::*;
245
246    #[test]
247    fn append_header() {
248        let mut ctx = Context::<_, 64>::new(&SystemTimeDateTimeHandler);
249
250        let mut res = Response::new(BoxBody::new(Full::new(Bytes::new())));
251
252        res.headers_mut()
253            .insert(CONNECTION, HeaderValue::from_static("keep-alive"));
254        res.headers_mut()
255            .append(CONNECTION, HeaderValue::from_static("upgrade"));
256
257        let (parts, body) = res.into_parts();
258
259        let mut buf = BytesMut::new();
260        ctx.encode_head(parts, &body, &mut buf).unwrap();
261
262        let mut header = [httparse::EMPTY_HEADER; 8];
263        let mut res = httparse::Response::new(&mut header);
264
265        let httparse::Status::Complete(_) = res.parse(buf.as_ref()).unwrap() else {
266            panic!("failed to parse response")
267        };
268
269        for h in header {
270            if h.name == "connection" {
271                assert_eq!(h.value, b"keep-alive, upgrade");
272            }
273        }
274    }
275
276    #[test]
277    fn multi_set_cookie() {
278        let mut ctx = Context::<_, 64>::new(&SystemTimeDateTimeHandler);
279
280        let mut res = Response::new(BoxBody::new(Full::new(Bytes::new())));
281
282        res.headers_mut()
283            .insert(SET_COOKIE, HeaderValue::from_static("foo=foo"));
284        res.headers_mut()
285            .append(SET_COOKIE, HeaderValue::from_static("bar=bar"));
286
287        let (parts, body) = res.into_parts();
288
289        let mut buf = BytesMut::new();
290        ctx.encode_head(parts, &body, &mut buf).unwrap();
291
292        let mut header = [httparse::EMPTY_HEADER; 8];
293        let mut res = httparse::Response::new(&mut header);
294
295        let httparse::Status::Complete(_) = res.parse(buf.as_ref()).unwrap() else {
296            panic!("failed to parse response")
297        };
298
299        assert_eq!(header[0].name, "set-cookie");
300        assert_eq!(header[0].value, b"foo=foo");
301        assert_eq!(header[1].name, "set-cookie");
302        assert_eq!(header[1].value, b"bar=bar");
303    }
304}