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