wireshift_core/buffer/

mod.rs

//! Type-state buffer pool.

use std::marker::PhantomData;

use crate::error::{Error, Result};

pub(crate) mod pool;
pub use pool::BufferPool;

pub(crate) const REGISTERED_BUFFER_CAPACITY: usize = 256 * 1024;
pub(crate) const REGISTERED_BUFFER_COUNT: usize = 32;

/// Marker type for buffers owned by user code.
#[derive(Debug)]
pub struct Owned;

/// Marker type for buffers currently owned by the backend.
#[derive(Debug)]
pub struct Submitted;

/// Marker type for buffers returned from the backend with data available.
#[derive(Debug)]
pub struct Completed;

/// Marker type for buffers part of a special pool.
#[derive(Debug)]
pub struct Pool;

#[derive(Debug)]
pub(crate) struct PooledBuffer {
    pub(crate) data: Vec<u8>,
    pub(crate) reusable: bool,
}

impl PooledBuffer {
    pub(crate) fn new(data: Vec<u8>, reusable: bool) -> Self {
        Self { data, reusable }
    }

    pub(crate) fn capacity(&self) -> usize {
        self.data.capacity()
    }

    pub(crate) fn zero_contents(&mut self) {
        self.data.fill(0);
    }

    pub(crate) fn deallocate(self) {
        drop(self.data);
    }
}

#[cfg(test)]
mod tests;

pub(crate) fn allocate_reusable_buffer(capacity: usize) -> Result<PooledBuffer> {
    if capacity > 1024 * 1024 * 1024 {
        return Err(Error::validation(
            "requested buffer capacity is too large",
            "request a buffer size less than 1 GiB",
        ));
    }
    // For simplicity and safety, we use a standard Vec.
    // If strict 4096 alignment is required for O_DIRECT,
    // we would need a safe aligned allocation library.
    // For now, we prioritize the "No Unsafe" law.
    let mut data = vec![0; capacity];
    data.shrink_to_fit();
    Ok(PooledBuffer::new(data, true))
}

pub(crate) fn allocate_overflow_buffer(capacity: usize) -> PooledBuffer {
    let mut data = vec![0; capacity];
    data.shrink_to_fit();
    PooledBuffer::new(data, false)
}

/// A buffer with type-state tracking ownership.
#[derive(Debug)]
pub struct Buffer<State> {
    pool: std::sync::Arc<pool::BufferPoolInner>,
    data: Option<PooledBuffer>,
    filled: usize,
    state: PhantomData<State>,
}

impl<State> Buffer<State> {
    pub(crate) fn new(
        pool: std::sync::Arc<pool::BufferPoolInner>,
        data: PooledBuffer,
        filled: usize,
    ) -> Self {
        Self {
            pool,
            data: Some(data),
            filled,
            state: PhantomData,
        }
    }

    fn data_ref(&self) -> &[u8] {
        self.data.as_ref().map_or(&[], |data| data.data.as_slice())
    }

    fn data_mut(&mut self) -> &mut [u8] {
        self.data
            .as_mut()
            .map_or(&mut [], |data| data.data.as_mut_slice())
    }

    /// Returns the total capacity.
    #[must_use]
    pub fn capacity(&self) -> usize {
        self.data.as_ref().map_or(0, PooledBuffer::capacity)
    }

    /// Returns the number of initialized bytes.
    #[must_use]
    pub fn filled_len(&self) -> usize {
        self.filled
    }
}

impl Buffer<Owned> {
    /// Returns the full immutable slice.
    #[must_use]
    pub fn as_slice(&self) -> &[u8] {
        self.data_ref()
    }

    /// Returns the full mutable slice for writing.
    pub fn as_mut_slice(&mut self) -> &mut [u8] {
        self.data_mut()
    }

    /// Marks bytes as initialized.
    pub fn set_filled_len(mut self, len: usize) -> Result<Self> {
        if len > self.capacity() {
            return Err(Error::validation(
                "filled length exceeds buffer capacity",
                "ensure data written does not exceed buffer size",
            ));
        }
        self.filled = len;
        Ok(self)
    }

    /// Transitions the buffer into submitted state.
    #[must_use]
    pub fn into_submitted(mut self) -> Buffer<Submitted> {
        Buffer {
            pool: self.pool.clone(),
            data: self.data.take(),
            filled: self.filled,
            state: PhantomData,
        }
    }
}

impl Buffer<Submitted> {
    /// Creates a submitted buffer from caller-owned memory.
    ///
    /// This bypasses the buffer pool entirely, useful for operations like `Write`
    /// where the user provides an arbitrary vector.
    #[allow(clippy::needless_pass_by_value)]
    pub fn from_vec(data: Vec<u8>) -> Result<Self> {
        if data.is_empty() {
            return Err(Error::validation(
                "submitted buffer must be non-empty",
                "allocate at least one byte before submitting a read buffer",
            ));
        }
        let capacity = data.len();
        let pool = std::sync::Arc::new(pool::BufferPoolInner {
            capacity,
            max_buffers: 1,
            overflow_on_exhaustion: false,
            free: crossbeam_queue::ArrayQueue::new(1),
        });
        Ok(Buffer::new(pool, PooledBuffer::new(data, false), 0))
    }

    /// Provides a mutable view for backend code.
    #[doc(hidden)]
    pub fn backend_mut(&mut self) -> &mut [u8] {
        self.data_mut()
    }

    /// Provides an immutable view for backend code.
    #[doc(hidden)]
    #[must_use]
    pub fn backend_ref(&self) -> &[u8] {
        self.data_ref()
    }

    /// Transitions the buffer into completed state.
    pub fn into_completed(mut self, filled: usize) -> Result<Buffer<Completed>> {
        if filled > self.capacity() {
            return Err(Error::validation(
                "completed length exceeds buffer capacity",
                "return at most the number of bytes the buffer can hold",
            ));
        }
        Ok(Buffer {
            pool: self.pool.clone(),
            data: self.data.take(),
            filled,
            state: PhantomData,
        })
    }
}

impl Buffer<Pool> {
    /// Provides a mutable view for backend code.
    #[doc(hidden)]
    pub fn as_mut_slice(&mut self) -> &mut [u8] {
        self.data_mut()
    }

    /// Transitions the buffer into completed state.
    pub fn into_completed(mut self, filled: usize) -> Result<Buffer<Completed>> {
        if filled > self.capacity() {
            return Err(Error::validation(
                "completed length exceeds buffer capacity",
                "return at most the number of bytes the buffer can hold",
            ));
        }
        Ok(Buffer {
            pool: self.pool.clone(),
            data: self.data.take(),
            filled,
            state: PhantomData,
        })
    }
}

impl Buffer<Completed> {
    /// Returns the slice of initialized data.
    #[must_use]
    pub fn filled(&self) -> &[u8] {
        &self.data_ref()[..self.filled]
    }

    /// Truncates the filled portion to `len` bytes.
    pub fn truncate(mut self, len: usize) -> Result<Self> {
        if len > self.filled {
            return Err(Error::validation(
                "truncate length exceeds currently filled bytes",
                "truncate to a length less than or equal to filled_len()",
            ));
        }
        self.filled = len;
        Ok(self)
    }

    /// Releases the buffer back to the user code explicitly resetting the filled length.
    #[must_use]
    pub fn release(mut self) -> Buffer<Owned> {
        let data = self
            .data
            .take()
            .unwrap_or_else(|| allocate_overflow_buffer(0));
        Buffer::new(self.pool.clone(), data, 0)
    }
}

impl<State> Drop for Buffer<State> {
    fn drop(&mut self) {
        if let Some(mut data) = self.data.take() {
            data.zero_contents();
            if data.reusable {
                if let Err(data) = self.pool.free.push(data) {
                    data.deallocate();
                }
            } else {
                data.deallocate();
            }
        }
    }
}