[pbs-devel] [PATCH proxmox 1/9] proxmox/tools/byte_buffer: improve ByteBuffer interface

[Dominik Csapak <d.csapak@proxmox.com>]

by implementing Deref and DerefMut, renaming consume to 'remove_data'
adapt the usage inside of websocket (we did not use it anywhere else for now)

Signed-off-by: Dominik Csapak <d.csapak@proxmox.com>
---
 proxmox/src/tools/byte_buffer.rs | 91 ++++++++++++++++++++------------
 proxmox/src/tools/websocket.rs   | 11 ++--
 2 files changed, 64 insertions(+), 38 deletions(-)

diff --git a/proxmox/src/tools/byte_buffer.rs b/proxmox/src/tools/byte_buffer.rs
index abfa6c8..7d4fce4 100644
--- a/proxmox/src/tools/byte_buffer.rs
+++ b/proxmox/src/tools/byte_buffer.rs
@@ -9,14 +9,13 @@
 //! fn code<T: Read + ?Sized>(input: &mut T) -> std::io::Result<Box<[u8]>> {
 //!     let mut buffer = ByteBuffer::new();
 //!     let amount = buffer.read_from(input)?;
-//!     let data = buffer.consume(amount);
+//!     let data = buffer.remove_data(amount);
 //!     assert_eq!(data.len(), amount);
 //!     Ok(data)
 //! }
 //! # code(&mut &b"testdata"[..]).expect("byte buffer test failed");
 //! ```
 
-use std::cmp::min;
 use std::io::{Read, Result};
 
 use crate::tools::vec;
@@ -49,23 +48,18 @@ impl ByteBuffer {
         }
     }
 
-    /// Returns the length of the data in the Buffer
-    pub fn data_size(&self) -> usize {
-        self.data_size
-    }
-
     pub fn free_size(&self) -> usize {
         self.capacity - self.data_size
     }
 
-    pub fn is_empty(&self) -> bool {
-        self.data_size == 0
-    }
-
     pub fn is_full(&self) -> bool {
         self.data_size >= self.capacity
     }
 
+    pub fn clear(&mut self) {
+        self.data_size = 0
+    }
+
     /// Sets the length of the data. Useful if data was manually added
     /// with a mutable slice (e.g. from [get_free_mut_slice](#method.get_free_mut_slice)).
     ///
@@ -92,19 +86,6 @@ impl ByteBuffer {
         self.data_size += size;
     }
 
-    /// Gets an immutable reference to the data in the buffer
-    /// Example:
-    /// ```
-    /// # use proxmox::tools::byte_buffer::ByteBuffer;
-    /// let mut buf = ByteBuffer::new();
-    /// buf.get_free_mut_slice()[..2].copy_from_slice(&[1u8, 2u8]);
-    /// buf.add_size(2);
-    /// assert_eq!(buf.get_data_slice(), &[1u8, 2u8]);
-    /// ```
-    pub fn get_data_slice(&self) -> &[u8] {
-        &self.buf[..self.data_size]
-    }
-
     /// Returns a mutable reference to the free section of the
     /// Buffer. There are no guarantees about the content of the
     /// free part of the Buffer (may even be uninitialized).
@@ -113,8 +94,8 @@ impl ByteBuffer {
         &mut self.buf[self.data_size..self.capacity]
     }
 
-    /// Consumes up to max_amount of data from the front
-    /// of the buffer. If there was less than max_amount present,
+    /// Removes up to max_amount of data from the front
+    /// of the buffer and returns. If there was less than max_amount present,
     /// it will return as much data as there was in the buffer.
     /// The rest of the data will be moved to the front, and
     /// the data size will be updated accordingly.
@@ -125,20 +106,50 @@ impl ByteBuffer {
     /// let mut buf = ByteBuffer::new();
     /// buf.get_free_mut_slice()[..2].copy_from_slice(&[1u8, 2u8]);
     /// buf.add_size(2);
-    /// assert_eq!(buf.data_size(), 2);
+    /// assert_eq!(buf.len(), 2);
     ///
-    /// let data = buf.consume(100);
+    /// let data = buf.remove_data(100);
     /// assert_eq!(&data[..], &[1u8, 2u8]);
     /// assert!(buf.is_empty());
     /// ```
-    pub fn consume(&mut self, max_amount: usize) -> Box<[u8]> {
-        let size = min(max_amount, self.data_size);
+    #[must_use]
+    pub fn remove_data(&mut self, max_amount: usize) -> Box<[u8]> {
+        let size = max_amount.min(self.data_size);
         let tmp: Box<[u8]> = self.buf[..size].into();
         self.buf.copy_within(size..self.capacity, 0);
         self.data_size -= size;
         tmp
     }
 
+    /// Removes up to max_amount of data from the front and returns
+    /// the amount of data removed. If there was less than max_amount present,
+    /// it will empty out the buffer and return the amount removed.
+    ///
+    /// Example:
+    /// ```
+    /// # use proxmox::tools::byte_buffer::ByteBuffer;
+    /// let mut buf = ByteBuffer::new();
+    /// buf.get_free_mut_slice()[..2].copy_from_slice(&[1u8, 2u8]);
+    /// buf.add_size(2);
+    /// assert_eq!(buf.len(), 2);
+    ///
+    /// let amount = buf.consume(1);
+    /// assert_eq!(amount, 1);
+    /// let amount = buf.consume(100);
+    /// assert_eq!(amount, 1);
+    /// assert!(buf.is_empty());
+    /// ```
+    pub fn consume(&mut self, max_amount: usize) -> usize {
+        let size = max_amount.min(self.data_size);
+        if size < max_amount {
+            self.clear()
+        } else {
+            self.buf.copy_within(size..self.capacity, 0);
+            self.data_size -= size;
+        }
+        size
+    }
+
     /// Takes a reader and reads into the back of the buffer (up to the
     /// free space in the buffer) and updates its size accordingly.
     ///
@@ -168,6 +179,20 @@ impl ByteBuffer {
     }
 }
 
+impl std::ops::Deref for ByteBuffer {
+    type Target = [u8];
+
+    fn deref(&self) -> &Self::Target {
+        &self.buf[..self.data_size]
+    }
+}
+
+impl std::ops::DerefMut for ByteBuffer {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.buf[..self.data_size]
+    }
+}
+
 #[cfg(test)]
 mod test {
     use crate::tools::byte_buffer::ByteBuffer;
@@ -181,7 +206,7 @@ mod test {
         }
         buffer.add_size(5);
 
-        let slice2 = buffer.get_data_slice();
+        let slice2 = &buffer[..];
 
         assert_eq!(slice2, &[0, 1, 2, 3, 4]);
     }
@@ -190,7 +215,7 @@ mod test {
     fn test2() {
         let mut buffer = ByteBuffer::with_capacity(1024);
         let size = buffer.read_from(&mut std::io::repeat(54)).unwrap();
-        assert_eq!(buffer.data_size(), size);
-        assert_eq!(buffer.get_data_slice()[0], 54);
+        assert_eq!(buffer.len(), size);
+        assert_eq!(buffer[0], 54);
     }
 }
diff --git a/proxmox/src/tools/websocket.rs b/proxmox/src/tools/websocket.rs
index 3877e4e..04173bb 100644
--- a/proxmox/src/tools/websocket.rs
+++ b/proxmox/src/tools/websocket.rs
@@ -485,7 +485,7 @@ impl<R: AsyncReadExt + Unpin + Send + 'static> AsyncRead for WebSocketReader<R>
                     let mut header = match this.header.take() {
                         Some(header) => header,
                         None => {
-                            let header = match FrameHeader::try_from_bytes(read_buffer.get_data_slice())? {
+                            let header = match FrameHeader::try_from_bytes(&read_buffer[..])? {
                                 Ok(header) => header,
                                 Err(_) => {
                                     this.state = ReaderState::NoData;
@@ -500,12 +500,12 @@ impl<R: AsyncReadExt + Unpin + Send + 'static> AsyncRead for WebSocketReader<R>
                     };
 
                     if header.is_control_frame() {
-                        if read_buffer.data_size() >= header.payload_len {
+                        if read_buffer.len() >= header.payload_len {
                             (this.callback)(
                                 header.frametype,
                                 mask_bytes(
                                     header.mask,
-                                    &mut read_buffer.consume(header.payload_len).into_vec(),
+                                    &mut read_buffer.remove_data(header.payload_len).into_vec(),
                                 ),
                             );
                             this.state =  if read_buffer.is_empty() {
@@ -523,8 +523,9 @@ impl<R: AsyncReadExt + Unpin + Send + 'static> AsyncRead for WebSocketReader<R>
                         }
                     }
 
-                    let len = min(buf.len() - offset, min(header.payload_len, read_buffer.data_size()));
-                    let mut data = read_buffer.consume(len).into_vec();
+                    let len = min(buf.len() - offset, min(header.payload_len, read_buffer.len()));
+
+                    let mut data = read_buffer.remove_data(len).into_vec();
                     buf[offset..offset+len].copy_from_slice(mask_bytes(header.mask, &mut data));
                     offset += len;
 
-- 
2.20.1