use crate::error::{BufferError, Result};
use zeroize::Zeroize;
pub const BUF_MAX_INCR: usize = 1_000_000_000;
pub const BUF_MAX_SIZE: usize = 1_000_000_000;
#[derive(Clone, Zeroize)]
#[zeroize(drop)]
pub struct Buffer {
pub(crate) data: Vec<u8>,
pub(crate) pos: usize,
pub(crate) len: usize,
}
impl Buffer {
#[inline]
pub fn new(size: usize) -> Self {
assert!(
size <= BUF_MAX_SIZE,
"Buffer size {} exceeds maximum {}",
size,
BUF_MAX_SIZE
);
Self {
data: vec![0; size],
pos: 0,
len: 0,
}
}
#[inline]
pub fn with_capacity(capacity: usize) -> Self {
assert!(
capacity <= BUF_MAX_SIZE,
"Buffer capacity {} exceeds maximum {}",
capacity,
BUF_MAX_SIZE
);
Self {
data: Vec::with_capacity(capacity),
pos: 0,
len: 0,
}
}
pub fn from_vec(data: Vec<u8>) -> Self {
let len = data.len();
Self { data, pos: 0, len }
}
#[inline(always)]
pub fn capacity(&self) -> usize {
self.data.capacity()
}
#[inline(always)]
pub fn len(&self) -> usize {
self.len
}
#[inline(always)]
pub fn is_empty(&self) -> bool {
self.len == 0
}
#[inline(always)]
pub fn pos(&self) -> usize {
self.pos
}
#[inline(always)]
pub fn remaining(&self) -> usize {
self.len.saturating_sub(self.pos)
}
#[inline(always)]
pub fn has_remaining(&self, count: usize) -> bool {
self.remaining() >= count
}
#[inline]
pub fn set_pos(&mut self, pos: usize) -> Result<()> {
if pos > self.len {
return Err(BufferError::PositionOutOfBounds);
}
self.pos = pos;
Ok(())
}
pub fn set_len(&mut self, len: usize) -> Result<()> {
if len > BUF_MAX_SIZE {
return Err(BufferError::SizeTooBig);
}
if len > self.data.len() {
self.data.resize(len, 0);
}
self.len = len;
self.pos = self.pos.min(len);
Ok(())
}
pub fn incr_pos(&mut self, incr: usize) -> Result<()> {
if incr > BUF_MAX_INCR || self.pos + incr > self.len {
return Err(BufferError::IncrementTooLarge);
}
self.pos += incr;
Ok(())
}
pub fn decr_pos(&mut self, decr: usize) -> Result<()> {
if decr > self.pos {
return Err(BufferError::PositionOutOfBounds);
}
self.pos -= decr;
Ok(())
}
pub fn incr_len(&mut self, incr: usize) -> Result<()> {
if incr > BUF_MAX_INCR {
return Err(BufferError::IncrementTooLarge);
}
let new_len = self.len + incr;
if new_len > BUF_MAX_SIZE {
return Err(BufferError::SizeTooBig);
}
if new_len > self.data.len() {
self.data.resize(new_len, 0);
}
self.len = new_len;
Ok(())
}
pub fn incr_write_pos(&mut self, incr: usize) -> Result<()> {
if incr > BUF_MAX_INCR {
return Err(BufferError::IncrementTooLarge);
}
let new_pos = self.pos + incr;
if new_pos > BUF_MAX_SIZE {
return Err(BufferError::IncrementTooLarge);
}
if new_pos > self.data.len() {
self.data.resize(new_pos, 0);
}
self.pos = new_pos;
if self.pos > self.len {
self.len = self.pos;
}
Ok(())
}
#[inline]
pub fn reset(&mut self) {
self.pos = 0;
self.len = 0;
}
pub fn burn(&mut self) {
self.data.as_mut_slice().zeroize();
self.pos = 0;
self.len = 0;
}
pub fn burn_free(mut self) {
self.data.as_mut_slice().zeroize();
drop(self);
}
#[inline]
pub fn as_slice(&self) -> &[u8] {
&self.data[..self.len]
}
#[inline]
pub fn as_mut_slice(&mut self) -> &mut [u8] {
&mut self.data[..self.len]
}
pub fn resize(&mut self, new_size: usize) -> Result<()> {
if new_size > BUF_MAX_SIZE {
return Err(BufferError::SizeTooBig);
}
self.data.resize(new_size, 0);
self.len = self.len.min(new_size);
self.pos = self.pos.min(new_size);
Ok(())
}
#[inline]
pub fn reserve(&mut self, additional: usize) {
self.data.reserve(additional);
}
#[inline]
pub fn shrink_to_fit(&mut self) {
self.data.truncate(self.len);
self.data.shrink_to_fit();
}
pub fn into_vec(mut self) -> Vec<u8> {
self.data.truncate(self.len); let mut v = Vec::new();
std::mem::swap(&mut self.data, &mut v); v
}
}
impl AsRef<[u8]> for Buffer {
fn as_ref(&self) -> &[u8] {
self.as_slice()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_new() {
let buf = Buffer::new(1024);
assert_eq!(buf.capacity(), 1024);
assert_eq!(buf.len(), 0);
assert_eq!(buf.pos(), 0);
}
#[test]
fn test_with_capacity() {
let mut buf = Buffer::with_capacity(1024);
assert_eq!(buf.capacity(), 1024);
assert_eq!(buf.len(), 0);
buf.data.extend_from_slice(b"hello");
buf.len = 5;
assert_eq!(buf.len(), 5);
assert_eq!(buf.as_slice(), b"hello");
}
#[test]
fn test_from_vec() {
let data = vec![1, 2, 3, 4, 5];
let buf = Buffer::from_vec(data);
assert_eq!(buf.len(), 5);
assert_eq!(buf.as_slice(), &[1, 2, 3, 4, 5]);
}
#[test]
fn test_incr_len_grows_vec() {
let mut buf = Buffer::with_capacity(100);
buf.incr_len(50).unwrap();
assert_eq!(buf.len(), 50);
assert!(buf.data.len() >= 50);
}
#[test]
fn test_shrink_to_fit() {
let mut buf = Buffer::new(1024);
buf.len = 100;
buf.shrink_to_fit();
assert_eq!(buf.data.len(), 100);
assert_eq!(buf.len(), 100);
}
#[test]
fn test_reset() {
let mut buf = Buffer::new(1024);
buf.pos = 50;
buf.len = 100;
buf.reset();
assert_eq!(buf.pos(), 0);
assert_eq!(buf.len(), 0);
}
}