From 5f4eace9050fd7dc41e6f02da44c4aae43782c14 Mon Sep 17 00:00:00 2001 From: Dominik Csapak Date: Tue, 14 Jul 2020 13:09:49 +0200 Subject: [PATCH] proxmox/tools/byte_buffer: improve ByteBuffer interface 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 Signed-off-by: Wolfgang Bumiller --- 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 abfa6c8a..7d4fce4e 100644 --- a/proxmox/src/tools/byte_buffer.rs +++ b/proxmox/src/tools/byte_buffer.rs @@ -9,14 +9,13 @@ //! fn code(input: &mut T) -> std::io::Result> { //! 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 3877e4e9..04173bb1 100644 --- a/proxmox/src/tools/websocket.rs +++ b/proxmox/src/tools/websocket.rs @@ -485,7 +485,7 @@ impl AsyncRead for WebSocketReader 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 AsyncRead for WebSocketReader }; 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 AsyncRead for WebSocketReader } } - 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;