1 //! A buffer for reading data from the network. 2 //! 3 //! The `ReadBuffer` is a buffer of bytes similar to a first-in, first-out queue. 4 //! It is filled by reading from a stream supporting `Read` and is then 5 //! accessible as a cursor for reading bytes. 6 7 use std::io::{Cursor, Read, Result as IoResult}; 8 9 use bytes::Buf; 10 11 /// A FIFO buffer for reading packets from the network. 12 #[derive(Debug)] 13 pub struct ReadBuffer<const CHUNK_SIZE: usize> { 14 storage: Cursor<Vec<u8>>, 15 chunk: Box<[u8; CHUNK_SIZE]>, 16 } 17 18 impl<const CHUNK_SIZE: usize> ReadBuffer<CHUNK_SIZE> { 19 /// Create a new empty input buffer. new() -> Self20 pub fn new() -> Self { 21 Self::with_capacity(CHUNK_SIZE) 22 } 23 24 /// Create a new empty input buffer with a given `capacity`. with_capacity(capacity: usize) -> Self25 pub fn with_capacity(capacity: usize) -> Self { 26 Self::from_partially_read(Vec::with_capacity(capacity)) 27 } 28 29 /// Create a input buffer filled with previously read data. from_partially_read(part: Vec<u8>) -> Self30 pub fn from_partially_read(part: Vec<u8>) -> Self { 31 Self { storage: Cursor::new(part), chunk: Box::new([0; CHUNK_SIZE]) } 32 } 33 34 /// Get a cursor to the data storage. as_cursor(&self) -> &Cursor<Vec<u8>>35 pub fn as_cursor(&self) -> &Cursor<Vec<u8>> { 36 &self.storage 37 } 38 39 /// Get a cursor to the mutable data storage. as_cursor_mut(&mut self) -> &mut Cursor<Vec<u8>>40 pub fn as_cursor_mut(&mut self) -> &mut Cursor<Vec<u8>> { 41 &mut self.storage 42 } 43 44 /// Consume the `ReadBuffer` and get the internal storage. into_vec(mut self) -> Vec<u8>45 pub fn into_vec(mut self) -> Vec<u8> { 46 // Current implementation of `tungstenite-rs` expects that the `into_vec()` drains 47 // the data from the container that has already been read by the cursor. 48 self.clean_up(); 49 50 // Now we can safely return the internal container. 51 self.storage.into_inner() 52 } 53 54 /// Read next portion of data from the given input stream. read_from<S: Read>(&mut self, stream: &mut S) -> IoResult<usize>55 pub fn read_from<S: Read>(&mut self, stream: &mut S) -> IoResult<usize> { 56 self.clean_up(); 57 let size = stream.read(&mut *self.chunk)?; 58 self.storage.get_mut().extend_from_slice(&self.chunk[..size]); 59 Ok(size) 60 } 61 62 /// Cleans ups the part of the vector that has been already read by the cursor. clean_up(&mut self)63 fn clean_up(&mut self) { 64 let pos = self.storage.position() as usize; 65 self.storage.get_mut().drain(0..pos).count(); 66 self.storage.set_position(0); 67 } 68 } 69 70 impl<const CHUNK_SIZE: usize> Buf for ReadBuffer<CHUNK_SIZE> { remaining(&self) -> usize71 fn remaining(&self) -> usize { 72 Buf::remaining(self.as_cursor()) 73 } 74 chunk(&self) -> &[u8]75 fn chunk(&self) -> &[u8] { 76 Buf::chunk(self.as_cursor()) 77 } 78 advance(&mut self, cnt: usize)79 fn advance(&mut self, cnt: usize) { 80 Buf::advance(self.as_cursor_mut(), cnt) 81 } 82 } 83 84 impl<const CHUNK_SIZE: usize> Default for ReadBuffer<CHUNK_SIZE> { default() -> Self85 fn default() -> Self { 86 Self::new() 87 } 88 } 89 90 #[cfg(test)] 91 mod tests { 92 use super::*; 93 94 #[test] simple_reading()95 fn simple_reading() { 96 let mut input = Cursor::new(b"Hello World!".to_vec()); 97 let mut buffer = ReadBuffer::<4096>::new(); 98 let size = buffer.read_from(&mut input).unwrap(); 99 assert_eq!(size, 12); 100 assert_eq!(buffer.chunk(), b"Hello World!"); 101 } 102 103 #[test] reading_in_chunks()104 fn reading_in_chunks() { 105 let mut inp = Cursor::new(b"Hello World!".to_vec()); 106 let mut buf = ReadBuffer::<4>::new(); 107 108 let size = buf.read_from(&mut inp).unwrap(); 109 assert_eq!(size, 4); 110 assert_eq!(buf.chunk(), b"Hell"); 111 112 buf.advance(2); 113 assert_eq!(buf.chunk(), b"ll"); 114 assert_eq!(buf.storage.get_mut(), b"Hell"); 115 116 let size = buf.read_from(&mut inp).unwrap(); 117 assert_eq!(size, 4); 118 assert_eq!(buf.chunk(), b"llo Wo"); 119 assert_eq!(buf.storage.get_mut(), b"llo Wo"); 120 121 let size = buf.read_from(&mut inp).unwrap(); 122 assert_eq!(size, 4); 123 assert_eq!(buf.chunk(), b"llo World!"); 124 } 125 } 126