xitca_http/h1/proto/
encode.rs

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