wireparse/io/
checked.rs

1use core::mem::MaybeUninit;
2
3use crate::error::{NetpackError, Result};
4
5pub trait CheckedReader: Sized {
6    fn read_slice_checked(&mut self, size: usize) -> Result<&[u8]>;
7    #[inline(always)]
8    fn read_checked<T: ReadDataChecked>(&mut self) -> Result<T> {
9        T::read_checked(self)
10    }
11}
12
13pub trait CheckedWriter: Sized {
14    fn write_slice(&mut self, data: &[u8]) -> Result<()>;
15    #[inline(always)]
16    fn write<T: WriteDataChecked>(&mut self, value: T) -> Result<()> {
17        T::write_to_checked(value, self)
18    }
19}
20
21impl<'a> CheckedReader for &'a [u8] {
22    #[inline(always)]
23    fn read_slice_checked(&mut self, size: usize) -> Result<&'a [u8]> {
24        if self.len() < size {
25            return Err(NetpackError::BufferTooSmall {
26                provided_size: self.len(),
27                requried_size: size,
28            });
29        }
30        let (data, remaining) = self.split_at(size);
31        *self = remaining;
32        Ok(data)
33    }
34}
35
36impl<'a> CheckedWriter for &'a mut [u8] {
37    #[inline(always)]
38    fn write_slice(&mut self, data: &[u8]) -> Result<()> {
39        if self.len() < data.len() {
40            return Err(NetpackError::BufferTooSmall {
41                provided_size: self.len(),
42                requried_size: data.len(),
43            });
44        }
45        let this = core::mem::take(self);
46        let (write_buffer, remaining) = this.split_at_mut(data.len());
47        *self = remaining;
48        write_buffer.copy_from_slice(data);
49        Ok(())
50    }
51}
52
53pub trait ReadDataChecked: Sized {
54    fn read_checked(reader: &mut impl CheckedReader) -> Result<Self>;
55}
56
57pub trait WriteDataChecked {
58    fn write_to_checked(self, writer: &mut impl CheckedWriter) -> Result<()>;
59}
60
61impl<const N: usize> ReadDataChecked for [u8; N] {
62    #[cfg_attr(feature = "fast-rw", inline(always))]
63    fn read_checked(reader: &mut impl CheckedReader) -> Result<Self> {
64        let slice = reader.read_slice_checked(N)?;
65        unsafe { Ok(*slice.as_ptr().cast()) }
66    }
67}
68
69impl<const N: usize> WriteDataChecked for [u8; N] {
70    fn write_to_checked(self, writer: &mut impl CheckedWriter) -> Result<()> {
71        writer.write_slice(&self)
72    }
73}
74
75impl<const N: usize> ReadDataChecked for [u16; N] {
76    fn read_checked(reader: &mut impl CheckedReader) -> Result<Self> {
77        let mut read_buffer = reader.read_slice_checked(core::mem::size_of::<u16>() * N)?;
78
79        let mut write_buffer: [MaybeUninit<u16>; N] =
80            unsafe { MaybeUninit::uninit().assume_init() };
81        let read_ptr = &mut read_buffer;
82
83        for elem in write_buffer.iter_mut() {
84            elem.write(u16::read_checked(read_ptr)?);
85        }
86
87        Ok(unsafe { *write_buffer.as_ptr().cast() })
88    }
89}
90impl<const N: usize> WriteDataChecked for [u16; N] {
91    fn write_to_checked(self, writer: &mut impl CheckedWriter) -> Result<()> {
92        for v in self.iter() {
93            writer.write(*v)?;
94        }
95        Ok(())
96    }
97}
98
99impl ReadDataChecked for u8 {
100    fn read_checked(reader: &mut impl CheckedReader) -> Result<Self> {
101        reader.read_slice_checked(1).map(|m| m[0])
102    }
103}
104
105impl WriteDataChecked for u8 {
106    fn write_to_checked(self, writer: &mut impl CheckedWriter) -> Result<()> {
107        writer.write_slice(&[self])
108    }
109}
110
111macro_rules! impl_read_write_data {
112    ($($t:ty),+) => {
113        $(
114            impl ReadDataChecked for $t {
115                #[inline(always)]
116                fn read_checked(reader: &mut impl CheckedReader) -> Result<$t> {
117                    let data = reader.read_checked();
118                    data.map(<$t>::from_be_bytes)
119                }
120            }
121
122            impl WriteDataChecked for $t {
123                #[inline(always)]
124                fn write_to_checked(self, writer: &mut impl CheckedWriter) -> Result<()> {
125                    let data = self.to_be_bytes();
126                    writer.write(data)
127                }
128            }
129         )*
130    };
131}
132
133impl_read_write_data! {
134    u16,
135    u32,
136    u64
137}
wireparse/io/checked.rs

wireparse/io/
checked.rs
