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 let skip_ct_te = match status {
47 StatusCode::SWITCHING_PROTOCOLS => true,
48 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 self.replace_headers(headers);
65
66 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 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 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 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 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 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 if self.is_head_method() {
171 try_remove_body(buf, skip_ct_te, size, &mut encoding);
172 } 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 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}