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