1// Copyright 2015 The tiny-http Contributors
2// Copyright 2015 The rust-chunked-transfer Contributors
3// Forked into ureq, 2022, from https://github.com/frewsxcv/rust-chunked-transfer
4//
5// Licensed under the Apache License, Version 2.0 (the "License");
6// you may not use this file except in compliance with the License.
7// You may obtain a copy of the License at
8//
9// https://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing, software
12// distributed under the License is distributed on an "AS IS" BASIS,
13// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14// See the License for the specific language governing permissions and
15// limitations under the License.
16
17use std::io::Result as IoResult;
18use std::io::Write;
19
20/// Splits the incoming data into HTTP chunks.
21///
22/// # Example
23///
24/// ```no_compile
25/// use chunked_transfer::Encoder;
26/// use std::io::Write;
27///
28/// let mut decoded = "hello world";
29/// let mut encoded: Vec<u8> = vec![];
30///
31/// {
32/// let mut encoder = Encoder::with_chunks_size(&mut encoded, 5);
33/// encoder.write_all(decoded.as_bytes());
34/// }
35///
36/// assert_eq!(encoded, b"5\r\nhello\r\n5\r\n worl\r\n1\r\nd\r\n0\r\n\r\n");
37/// ```
38pub struct Encoder<W>
39where
40 W: Write,
41{
42 // where to send the result
43 output: W,
44
45 // size of each chunk
46 chunks_size: usize,
47
48 // data waiting to be sent is stored here
49 // This will always be at least 6 bytes long. The first 6 bytes
50 // are reserved for the chunk size and \r\n.
51 buffer: Vec<u8>,
52
53 // Flushes the internal buffer after each write. This might be useful
54 // if data should be sent immediately to downstream consumers
55 flush_after_write: bool,
56}
57
58const MAX_CHUNK_SIZE: usize = std::u32::MAX as usize;
59// This accounts for four hex digits (enough to hold a u32) plus two bytes
60// for the \r\n
61const MAX_HEADER_SIZE: usize = 6;
62
63impl<W> Encoder<W>
64where
65 W: Write,
66{
67 pub fn new(output: W) -> Encoder<W> {
68 Encoder::with_chunks_size(output, 8192)
69 }
70
71 pub fn with_chunks_size(output: W, chunks: usize) -> Encoder<W> {
72 let chunks_size = chunks.min(MAX_CHUNK_SIZE);
73 let mut encoder = Encoder {
74 output,
75 chunks_size,
76 buffer: vec![0; MAX_HEADER_SIZE],
77 flush_after_write: false,
78 };
79 encoder.reset_buffer();
80 encoder
81 }
82
83 pub fn with_flush_after_write(output: W) -> Encoder<W> {
84 let mut encoder = Encoder {
85 output,
86 chunks_size: 8192,
87 buffer: vec![0; MAX_HEADER_SIZE],
88 flush_after_write: true,
89 };
90 encoder.reset_buffer();
91 encoder
92 }
93
94 fn reset_buffer(&mut self) {
95 // Reset buffer, still leaving space for the chunk size. That space
96 // will be populated once we know the size of the chunk.
97 self.buffer.truncate(MAX_HEADER_SIZE);
98 }
99
100 fn is_buffer_empty(&self) -> bool {
101 self.buffer.len() == MAX_HEADER_SIZE
102 }
103
104 fn buffer_len(&self) -> usize {
105 self.buffer.len() - MAX_HEADER_SIZE
106 }
107
108 fn send(&mut self) -> IoResult<()> {
109 // Never send an empty buffer, because that would be interpreted
110 // as the end of the stream, which we indicate explicitly on drop.
111 if self.is_buffer_empty() {
112 return Ok(());
113 }
114 // Prepend the length and \r\n to the buffer.
115 let prelude = format!("{:x}\r\n", self.buffer_len());
116 let prelude = prelude.as_bytes();
117
118 // This should never happen because MAX_CHUNK_SIZE of u32::MAX
119 // can always be encoded in 4 hex bytes.
120 assert!(
121 prelude.len() <= MAX_HEADER_SIZE,
122 "invariant failed: prelude longer than MAX_HEADER_SIZE"
123 );
124
125 // Copy the prelude into the buffer. For small chunks, this won't necessarily
126 // take up all the space that was reserved for the prelude.
127 let offset = MAX_HEADER_SIZE - prelude.len();
128 self.buffer[offset..MAX_HEADER_SIZE].clone_from_slice(prelude);
129
130 // Append the chunk-finishing \r\n to the buffer.
131 self.buffer.write_all(b"\r\n")?;
132
133 self.output.write_all(&self.buffer[offset..])?;
134 self.reset_buffer();
135
136 Ok(())
137 }
138}
139
140impl<W> Write for Encoder<W>
141where
142 W: Write,
143{
144 fn write(&mut self, data: &[u8]) -> IoResult<usize> {
145 let remaining_buffer_space: usize = self.chunks_size - self.buffer_len();
146 let bytes_to_buffer: usize = std::cmp::min(v1:remaining_buffer_space, v2:data.len());
147 self.buffer.extend_from_slice(&data[0..bytes_to_buffer]);
148 let more_to_write: bool = bytes_to_buffer < data.len();
149 if self.flush_after_write || more_to_write {
150 self.send()?;
151 }
152
153 // If we didn't write the whole thing, keep working on it.
154 if more_to_write {
155 self.write_all(&data[bytes_to_buffer..])?;
156 }
157 Ok(data.len())
158 }
159
160 fn flush(&mut self) -> IoResult<()> {
161 self.send()
162 }
163}
164
165impl<W> Drop for Encoder<W>
166where
167 W: Write,
168{
169 fn drop(&mut self) {
170 self.flush().ok();
171 write!(self.output, "0\r\n\r\n").ok();
172 }
173}
174
175#[cfg(test)]
176mod test {
177 use super::Encoder;
178 use std::io;
179 use std::io::Write;
180 use std::str::from_utf8;
181
182 #[test]
183 fn test() {
184 let mut source = io::Cursor::new("hello world".to_string().into_bytes());
185 let mut dest: Vec<u8> = vec![];
186
187 {
188 let mut encoder = Encoder::with_chunks_size(dest.by_ref(), 5);
189 io::copy(&mut source, &mut encoder).unwrap();
190 assert!(!encoder.is_buffer_empty());
191 }
192
193 let output = from_utf8(&dest).unwrap();
194
195 assert_eq!(output, "5\r\nhello\r\n5\r\n worl\r\n1\r\nd\r\n0\r\n\r\n");
196 }
197 #[test]
198 fn flush_after_write() {
199 let mut source = io::Cursor::new("hello world".to_string().into_bytes());
200 let mut dest: Vec<u8> = vec![];
201
202 {
203 let mut encoder = Encoder::with_flush_after_write(dest.by_ref());
204 io::copy(&mut source, &mut encoder).unwrap();
205 // The internal buffer has been flushed.
206 assert!(encoder.is_buffer_empty());
207 }
208
209 let output = from_utf8(&dest).unwrap();
210
211 assert_eq!(output, "b\r\nhello world\r\n0\r\n\r\n");
212 }
213}
214