1use crate::stream::Stream;
2use std::fmt;
3use std::io::{self, copy, empty, Cursor, Read, Write};
4
5#[cfg(feature = "charset")]
6use crate::response::DEFAULT_CHARACTER_SET;
7#[cfg(feature = "charset")]
8use encoding_rs::Encoding;
9
10/// The different kinds of bodies to send.
11///
12/// *Internal API*
13pub(crate) enum Payload<'a> {
14 Empty,
15 Text(&'a str, String),
16 Reader(Box<dyn Read + 'a>),
17 Bytes(&'a [u8]),
18}
19
20impl fmt::Debug for Payload<'_> {
21 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
22 match self {
23 Payload::Empty => write!(f, "Empty"),
24 Payload::Text(t: &&str, _) => write!(f, "{}", t),
25 Payload::Reader(_) => write!(f, "Reader"),
26 Payload::Bytes(v: &&[u8]) => write!(f, "{:?}", v),
27 }
28 }
29}
30
31#[allow(clippy::derivable_impls)]
32impl Default for Payload<'_> {
33 fn default() -> Self {
34 Payload::Empty
35 }
36}
37
38/// The size of the body.
39///
40/// *Internal API*
41#[derive(Debug)]
42pub(crate) enum BodySize {
43 Empty,
44 Unknown,
45 Known(u64),
46}
47
48/// Payloads are turned into this type where we can hold both a size and the reader.
49///
50/// *Internal API*
51pub(crate) struct SizedReader<'a> {
52 pub size: BodySize,
53 pub reader: Box<dyn Read + 'a>,
54}
55
56impl fmt::Debug for SizedReader<'_> {
57 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
58 write!(f, "SizedReader[size={:?},reader]", self.size)
59 }
60}
61
62impl<'a> SizedReader<'a> {
63 fn new(size: BodySize, reader: Box<dyn Read + 'a>) -> Self {
64 SizedReader { size, reader }
65 }
66}
67
68impl<'a> Payload<'a> {
69 pub fn into_read(self) -> SizedReader<'a> {
70 match self {
71 Payload::Empty => SizedReader::new(BodySize::Empty, Box::new(empty())),
72 Payload::Text(text, _charset) => {
73 #[cfg(feature = "charset")]
74 let bytes = {
75 let encoding = Encoding::for_label(_charset.as_bytes())
76 .or_else(|| Encoding::for_label(DEFAULT_CHARACTER_SET.as_bytes()))
77 .unwrap();
78 encoding.encode(text).0
79 };
80 #[cfg(not(feature = "charset"))]
81 let bytes = text.as_bytes();
82 let len = bytes.len();
83 let cursor = Cursor::new(bytes);
84 SizedReader::new(BodySize::Known(len as u64), Box::new(cursor))
85 }
86 Payload::Reader(read) => SizedReader::new(BodySize::Unknown, read),
87 Payload::Bytes(bytes) => {
88 let len = bytes.len();
89 let cursor = Cursor::new(bytes);
90 SizedReader::new(BodySize::Known(len as u64), Box::new(cursor))
91 }
92 }
93 }
94}
95
96const CHUNK_MAX_SIZE: usize = 0x4000; // Maximum size of a TLS fragment
97const CHUNK_HEADER_MAX_SIZE: usize = 6; // four hex digits plus "\r\n"
98const CHUNK_FOOTER_SIZE: usize = 2; // "\r\n"
99const CHUNK_MAX_PAYLOAD_SIZE: usize = CHUNK_MAX_SIZE - CHUNK_HEADER_MAX_SIZE - CHUNK_FOOTER_SIZE;
100
101// copy_chunks() improves over chunked_transfer's Encoder + io::copy with the
102// following performance optimizations:
103// 1) It avoid copying memory.
104// 2) chunked_transfer's Encoder issues 4 separate write() per chunk. This is costly
105// overhead. Instead, we do a single write() per chunk.
106// The measured benefit on a Linux machine is a 50% reduction in CPU usage on a https connection.
107fn copy_chunked<R: Read, W: Write>(reader: &mut R, writer: &mut W) -> io::Result<u64> {
108 // The chunk layout is:
109 // header:header_max_size | payload:max_payload_size | footer:footer_size
110 let mut chunk = Vec::with_capacity(CHUNK_MAX_SIZE);
111 let mut written = 0;
112 loop {
113 // We first read the payload
114 chunk.resize(CHUNK_HEADER_MAX_SIZE, 0);
115 let payload_size = reader
116 .take(CHUNK_MAX_PAYLOAD_SIZE as u64)
117 .read_to_end(&mut chunk)?;
118
119 // Then write the header
120 let header_str = format!("{:x}\r\n", payload_size);
121 let header = header_str.as_bytes();
122 assert!(header.len() <= CHUNK_HEADER_MAX_SIZE);
123 let start_index = CHUNK_HEADER_MAX_SIZE - header.len();
124 (&mut chunk[start_index..]).write_all(header).unwrap();
125
126 // And add the footer
127 chunk.extend_from_slice(b"\r\n");
128
129 // Finally Write the chunk
130 writer.write_all(&chunk[start_index..])?;
131 written += payload_size as u64;
132
133 // On EOF, we wrote a 0 sized chunk. This is what the chunked encoding protocol requires.
134 if payload_size == 0 {
135 return Ok(written);
136 }
137 }
138}
139
140/// Helper to send a body, either as chunked or not.
141pub(crate) fn send_body(
142 mut body: SizedReader,
143 do_chunk: bool,
144 stream: &mut Stream,
145) -> io::Result<()> {
146 if do_chunk {
147 copy_chunked(&mut body.reader, writer:stream)?;
148 } else {
149 copy(&mut body.reader, writer:stream)?;
150 };
151
152 Ok(())
153}
154
155#[cfg(test)]
156mod tests {
157 use super::*;
158
159 #[test]
160 fn test_copy_chunked() {
161 let mut source = Vec::<u8>::new();
162 source.resize(CHUNK_MAX_PAYLOAD_SIZE, 33);
163 source.extend_from_slice(b"hello world");
164
165 let mut dest = Vec::<u8>::new();
166 copy_chunked(&mut &source[..], &mut dest).unwrap();
167
168 let mut dest_expected = Vec::<u8>::new();
169 dest_expected.extend_from_slice(format!("{:x}\r\n", CHUNK_MAX_PAYLOAD_SIZE).as_bytes());
170 dest_expected.resize(dest_expected.len() + CHUNK_MAX_PAYLOAD_SIZE, 33);
171 dest_expected.extend_from_slice(b"\r\n");
172
173 dest_expected.extend_from_slice(b"b\r\nhello world\r\n");
174 dest_expected.extend_from_slice(b"0\r\n\r\n");
175
176 assert_eq!(dest, dest_expected);
177 }
178}
179