1use std::io::Result as IoResult;
2use std::io::Write;
34/// Splits the incoming data into HTTP chunks.
5///
6/// # Example
7///
8/// ```
9/// use chunked_transfer::Encoder;
10/// use std::io::Write;
11///
12/// let mut decoded = "hello world";
13/// let mut encoded: Vec<u8> = vec![];
14///
15/// {
16/// let mut encoder = Encoder::with_chunks_size(&mut encoded, 5);
17/// encoder.write_all(decoded.as_bytes());
18/// }
19///
20/// assert_eq!(encoded, b"5\r\nhello\r\n5\r\n worl\r\n1\r\nd\r\n0\r\n\r\n");
21/// ```
22pub struct Encoder<W>
23where
24W: Write,
25{
26// where to send the result
27output: W,
2829// size of each chunk
30chunks_size: usize,
3132// data waiting to be sent is stored here
33 // This will always be at least 6 bytes long. The first 6 bytes
34 // are reserved for the chunk size and \r\n.
35buffer: Vec<u8>,
3637// Flushes the internal buffer after each write. This might be useful
38 // if data should be sent immediately to downstream consumers
39flush_after_write: bool,
40}
4142const MAX_CHUNK_SIZE: usize = std::u32::MAX as usize;
43// This accounts for four hex digits (enough to hold a u32) plus two bytes
44// for the \r\n
45const MAX_HEADER_SIZE: usize = 6;
4647impl<W> Encoder<W>
48where
49W: Write,
50{
51pub fn new(output: W) -> Encoder<W> {
52 Encoder::with_chunks_size(output, 8192)
53 }
5455/// Gets a reference to the underlying value in this encoder.
56pub fn get_ref(&self) -> &W {
57&self.output
58 }
5960/// Gets a mutable reference to the underlying value in this encoder.
61pub fn get_mut(&mut self) -> &mut W {
62&mut self.output
63 }
6465pub fn with_chunks_size(output: W, chunks: usize) -> Encoder<W> {
66let chunks_size = chunks.min(MAX_CHUNK_SIZE);
67let mut encoder = Encoder {
68 output,
69 chunks_size,
70 buffer: vec![0; MAX_HEADER_SIZE],
71 flush_after_write: false,
72 };
73 encoder.reset_buffer();
74 encoder
75 }
7677pub fn with_flush_after_write(output: W) -> Encoder<W> {
78let mut encoder = Encoder {
79 output,
80 chunks_size: 8192,
81 buffer: vec![0; MAX_HEADER_SIZE],
82 flush_after_write: true,
83 };
84 encoder.reset_buffer();
85 encoder
86 }
8788fn reset_buffer(&mut self) {
89// Reset buffer, still leaving space for the chunk size. That space
90 // will be populated once we know the size of the chunk.
91self.buffer.truncate(MAX_HEADER_SIZE);
92 }
9394fn is_buffer_empty(&self) -> bool {
95self.buffer.len() == MAX_HEADER_SIZE
96 }
9798fn buffer_len(&self) -> usize {
99self.buffer.len() - MAX_HEADER_SIZE
100 }
101102fn send(&mut self) -> IoResult<()> {
103// Never send an empty buffer, because that would be interpreted
104 // as the end of the stream, which we indicate explicitly on drop.
105if self.is_buffer_empty() {
106return Ok(());
107 }
108// Prepend the length and \r\n to the buffer.
109let prelude = format!("{:x}\r\n", self.buffer_len());
110let prelude = prelude.as_bytes();
111112// This should never happen because MAX_CHUNK_SIZE of u32::MAX
113 // can always be encoded in 4 hex bytes.
114assert!(
115 prelude.len() <= MAX_HEADER_SIZE,
116"invariant failed: prelude longer than MAX_HEADER_SIZE"
117);
118119// Copy the prelude into the buffer. For small chunks, this won't necessarily
120 // take up all the space that was reserved for the prelude.
121let offset = MAX_HEADER_SIZE - prelude.len();
122self.buffer[offset..MAX_HEADER_SIZE].clone_from_slice(prelude);
123124// Append the chunk-finishing \r\n to the buffer.
125self.buffer.write_all(b"\r\n")?;
126127self.output.write_all(&self.buffer[offset..])?;
128self.reset_buffer();
129130Ok(())
131 }
132}
133134impl<W> Write for Encoder<W>
135where
136W: Write,
137{
138fn write(&mut self, data: &[u8]) -> IoResult<usize> {
139let remaining_buffer_space = self.chunks_size - self.buffer_len();
140let bytes_to_buffer = std::cmp::min(remaining_buffer_space, data.len());
141self.buffer.extend_from_slice(&data[0..bytes_to_buffer]);
142let more_to_write: bool = bytes_to_buffer < data.len();
143if self.flush_after_write || more_to_write {
144self.send()?;
145 }
146147// If we didn't write the whole thing, keep working on it.
148if more_to_write {
149self.write_all(&data[bytes_to_buffer..])?;
150 }
151Ok(data.len())
152 }
153154fn flush(&mut self) -> IoResult<()> {
155self.send()
156 }
157}
158159impl<W> Drop for Encoder<W>
160where
161W: Write,
162{
163fn drop(&mut self) {
164self.flush().ok();
165write!(self.output, "0\r\n\r\n").ok();
166 }
167}
168169#[cfg(test)]
170mod test {
171use super::Encoder;
172use std::io;
173use std::io::Write;
174use std::str::from_utf8;
175176#[test]
177fn test() {
178let mut source = io::Cursor::new("hello world".to_string().into_bytes());
179let mut dest: Vec<u8> = vec![];
180181 {
182let mut encoder = Encoder::with_chunks_size(dest.by_ref(), 5);
183 io::copy(&mut source, &mut encoder).unwrap();
184assert!(!encoder.is_buffer_empty());
185 }
186187let output = from_utf8(&dest).unwrap();
188189assert_eq!(output, "5\r\nhello\r\n5\r\n worl\r\n1\r\nd\r\n0\r\n\r\n");
190 }
191#[test]
192fn flush_after_write() {
193let mut source = io::Cursor::new("hello world".to_string().into_bytes());
194let mut dest: Vec<u8> = vec![];
195196 {
197let mut encoder = Encoder::with_flush_after_write(dest.by_ref());
198 io::copy(&mut source, &mut encoder).unwrap();
199// The internal buffer has been flushed.
200assert!(encoder.is_buffer_empty());
201 }
202203let output = from_utf8(&dest).unwrap();
204205assert_eq!(output, "b\r\nhello world\r\n0\r\n\r\n");
206 }
207}