From 83722c24aa859f884444d27106dc02e60c437c75 Mon Sep 17 00:00:00 2001
From: Daniel Abramov <dabramov@snapview.de>
Date: Tue, 6 Jul 2021 16:33:44 +0200
Subject: [PATCH] Clean up the read portion of the buffer

---
 src/buffer.rs | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/src/buffer.rs b/src/buffer.rs
index 74aef57..b738594 100644
--- a/src/buffer.rs
+++ b/src/buffer.rs
@@ -45,9 +45,7 @@ impl<const CHUNK_SIZE: usize> ReadBuffer<CHUNK_SIZE> {
     pub fn into_vec(mut self) -> Vec<u8> {
         // Current implementation of `tungstenite-rs` expects that the `into_vec()` drains
         // the data from the container that has already been read by the cursor.
-        let pos = self.storage.position() as usize;
-        self.storage.get_mut().drain(0..pos).count();
-        self.storage.set_position(0);
+        self.clean_up();
 
         // Now we can safely return the internal container.
         self.storage.into_inner()
@@ -55,10 +53,18 @@ impl<const CHUNK_SIZE: usize> ReadBuffer<CHUNK_SIZE> {
 
     /// Read next portion of data from the given input stream.
     pub fn read_from<S: Read>(&mut self, stream: &mut S) -> IoResult<usize> {
+        self.clean_up();
         let size = stream.read(&mut self.chunk)?;
         self.storage.get_mut().extend_from_slice(&self.chunk[..size]);
         Ok(size)
     }
+
+    /// Cleans ups the part of the vector that has been already read by the cursor.
+    fn clean_up(&mut self) {
+        let pos = self.storage.position() as usize;
+        self.storage.get_mut().drain(0..pos).count();
+        self.storage.set_position(0);
+    }
 }
 
 impl<const CHUNK_SIZE: usize> Buf for ReadBuffer<CHUNK_SIZE> {
@@ -99,10 +105,12 @@ mod tests {
 
         buf.advance(2);
         assert_eq!(buf.chunk(), b"ll");
+        assert_eq!(buf.storage.get_mut(), b"Hell");
 
         let size = buf.read_from(&mut inp).unwrap();
         assert_eq!(size, 4);
         assert_eq!(buf.chunk(), b"llo Wo");
+        assert_eq!(buf.storage.get_mut(), b"llo Wo");
 
         let size = buf.read_from(&mut inp).unwrap();
         assert_eq!(size, 4);