tor-bytes: Add a new writer implementation for fixed-size objects

Because the API assumes that many writes are infallible, this writer
takes ownership of the backing object, and will only return it to
you if you didn't run over the end.

I'm going to use this to save some allocations in relay cell bodies
This commit is contained in:
Nick Mathewson 2023-02-10 12:09:24 -05:00
parent 424b2ae4a5
commit 9e2b6f3aed
3 changed files with 125 additions and 0 deletions

View File

@ -1 +1,2 @@
MODIFIED: Error::BadMessage is deprecated, Error::InvalidMessage is new.
MODIFIED: New SliceWriter type.

View File

@ -42,11 +42,13 @@ mod err;
mod impls;
mod reader;
mod secretbuf;
mod slicewriter;
mod writer;
pub use err::{EncodeError, Error};
pub use reader::Reader;
pub use secretbuf::SecretBuf;
pub use slicewriter::{SliceWriter, SliceWriterError};
pub use writer::Writer;
use arrayref::array_ref;

View File

@ -0,0 +1,122 @@
//! A Writer that can put its results into an buffer of known byte size without
//! changing that size.
use thiserror::Error;
use crate::Writer;
/// An error that occurred while trying to unwrap a SliceWriter.
#[non_exhaustive]
#[derive(Clone, Debug, Error)]
pub enum SliceWriterError {
/// We've
#[error("Tried to write more than would fit into a fixed-size slice.")]
Truncated,
}
/// An object that supports writing into a byte-slice of fixed size.
///
/// Since the writer API does not allow all `write_*` functions to report errors,
/// this type defers any truncated-data errors until you try to retrieve the
/// inner data.
///
//
// TODO: in theory we could have a version of this that used MaybeUninit, but I
// don't think that would be worth it.
pub struct SliceWriter<T> {
/// The object we're writing into. Must have fewer than usize::LEN bytes.
data: T,
/// Our current write position within that object.
offset: usize,
}
impl<T> Writer for SliceWriter<T>
where
T: AsMut<[u8]>,
{
fn write_all(&mut self, b: &[u8]) {
let new_len = self.offset.saturating_add(b.len());
if new_len <= self.data.as_mut().len() {
// Note that if we reach this case, the addition was not saturating.
self.data.as_mut()[self.offset..new_len].copy_from_slice(b);
self.offset = new_len;
} else {
self.offset = usize::MAX;
}
}
}
impl<T> SliceWriter<T> {
/// Construct a new SliceWriter
///
/// Typically, you would want to use this on a type that implements
/// `AsMut<[u8]>`, or else it won't be very useful.
///
/// Preexisting bytes in the `data` object will be unchanged, unless you use
/// the [`Writer`] API to write to them.
pub fn new(data: T) -> Self {
Self { data, offset: 0 }
}
/// Try to extract the data from this `SliceWriter`.
///
/// On success (if we did not write "off the end" of the underlying object),
/// return the object and the number of bytes we wrote into it. (Bytes
/// after that position are unchanged.)
///
/// On failure (if we tried to write too much), return an error.
pub fn try_unwrap(self) -> Result<(T, usize), SliceWriterError> {
if self.offset != usize::MAX {
Ok((self.data, self.offset))
} else {
Err(SliceWriterError::Truncated)
}
}
}
#[cfg(test)]
mod test {
// @@ begin test lint list maintained by maint/add_warning @@
#![allow(clippy::bool_assert_comparison)]
#![allow(clippy::clone_on_copy)]
#![allow(clippy::dbg_macro)]
#![allow(clippy::print_stderr)]
#![allow(clippy::print_stdout)]
#![allow(clippy::single_char_pattern)]
#![allow(clippy::unwrap_used)]
#![allow(clippy::unchecked_duration_subtraction)]
//! <!-- @@ end test lint list maintained by maint/add_warning @@ -->
use super::*;
#[test]
fn basics() {
let mut w = SliceWriter::new([0_u8; 16]);
w.write_u8(b'h');
w.write_u16(0x656c);
w.write_u32(0x6c6f2077);
w.write_all(b"orld!");
let (a, len) = w.try_unwrap().unwrap();
assert_eq!(a.as_ref(), b"hello world!\0\0\0\0");
assert_eq!(len, 12);
}
#[test]
fn full_is_ok() {
let mut w = SliceWriter::new([0_u8; 4]);
w.write_u8(1);
w.write_u16(0x0203);
w.write_u8(4);
let (a, len) = w.try_unwrap().unwrap();
assert_eq!(a.as_ref(), [1, 2, 3, 4]);
assert_eq!(len, 4);
}
#[test]
fn too_full_is_not_ok() {
let mut w = SliceWriter::new([0_u8; 5]);
w.write_u32(12);
w.write_u32(12);
assert!(matches!(w.try_unwrap(), Err(SliceWriterError::Truncated)));
}
}