lib/rust/wlan-common/src/buffer_writer.rs

// Copyright 2019 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use {
    failure::{Error, ensure, format_err},
    zerocopy::{self, LayoutVerified, Unaligned},
};

pub use zerocopy::ByteSliceMut as ByteSliceMut;

pub struct BufferWriter<B: ByteSliceMut> {
    buf: B,
    written_bytes: usize,
}

impl <B: ByteSliceMut> BufferWriter<B> {
    pub fn new(buf: B) -> Self {
        Self {written_bytes: 0, buf}
    }


    /// Writes a single given `byte` and returns a `BufferWriter` to write remaining bytes of the
    /// original buffer.
    pub fn write_byte(mut self, byte: u8) -> Result<Self, Error> {
        ensure!(self.buf.len() >= 1, "buffer too short");
        let (mut data, remaining) = self.buf.split_at(1);
        data[0] = byte;
        Ok(Self {
            buf: remaining,
            written_bytes: self.written_bytes + 1
        })
    }

    /// Writes the given `bytes` and returns a `BufferWriter` to write remaining bytes of the
    /// original buffer.
    pub fn write_bytes(mut self, bytes: &[u8]) -> Result<Self, Error> {
        ensure!(self.buf.len() >= bytes.len(), "buffer too short");
        let (mut data, remaining) = self.buf.split_at(bytes.len());
        data.copy_from_slice(bytes);
        Ok(Self {
            buf: remaining,
            written_bytes: self.written_bytes + bytes.len()
        })
    }

    /// Reserves and zeroes a typed chunk of bytes of size |mem::size_of::<T>|.
    /// Returns a tuple of the typed reserved chunk and a `BufferWriter` for writing remaining
    /// bytes of the original buffer.
    pub fn reserve_zeroed<T: Unaligned>(mut self) -> Result<(LayoutVerified<B, T>, Self), Error> {
        let (data, remaining) = LayoutVerified::new_unaligned_from_prefix_zeroed(self.buf)
            .ok_or(format_err!("buffer too short"))?;
        Ok((data, Self {
            buf: remaining,
            written_bytes: self.written_bytes + std::mem::size_of::<T>()
        }))
    }

    pub fn written_bytes(&self) -> usize {
        self.written_bytes
    }

    pub fn remaining_bytes(&self) -> usize {
        self.buf.len()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn buffer_too_short() {
        assert!(BufferWriter::new(&mut [0; 0][..]).write_byte(1).is_err());
        assert!(BufferWriter::new(&mut [0; 5][..]).write_bytes(&[0; 6]).is_err());
        assert!(BufferWriter::new(&mut [0; 5][..]).reserve_zeroed::<[u8; 6]>().is_err());
    }

    #[test]
    fn reserve_zeroed_for_type() {
        let mut buf = [1u8; 5];
        let (mut data, w) = BufferWriter::new(&mut buf[..])
            .reserve_zeroed::<[u8; 3]>().expect("failed writing buffer");
        data[0] = 42;
        // Don't write `data[1]`: BufferWriter should zero this byte.
        data[2] = 43;

        assert_eq!(3, w.written_bytes());
        assert_eq!(2, w.remaining_bytes());
        assert_eq!([42, 0, 43, 1, 1], buf);
    }

    #[test]
    fn write_bytes() {
        let mut buf = [1u8; 5];
        let w = BufferWriter::new(&mut buf[..])
            .write_byte(42).expect("failed writing buffer")
            .write_byte(43).expect("failed writing buffer")
            .write_bytes(&[2, 3]).expect("failed writing buffer");

        assert_eq!(4, w.written_bytes());
        assert_eq!(1, w.remaining_bytes());
        assert_eq!([42, 43, 2, 3, 1], buf);
    }
}