exocore_core/framing/
padded.rs

1use std::io;
2
3use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
4use bytes::Bytes;
5
6use super::{check_from_size, check_into_size, Error, FrameBuilder, FrameReader};
7
8/// Frame that pads an underlying frame so that it has a minimum size.
9pub struct PaddedFrame<I: FrameReader> {
10    inner: I,
11    padding_size: usize,
12}
13
14impl<I: FrameReader> PaddedFrame<I> {
15    pub fn new(inner: I) -> Result<PaddedFrame<I>, Error> {
16        let exposed_data = inner.exposed_data();
17        check_from_size(4, exposed_data)?;
18
19        let padding_size =
20            (&exposed_data[exposed_data.len() - 4..]).read_u32::<LittleEndian>()? as usize;
21        check_from_size(4 + padding_size, exposed_data)?;
22
23        Ok(PaddedFrame {
24            inner,
25            padding_size,
26        })
27    }
28}
29
30impl<I: FrameReader> FrameReader for PaddedFrame<I> {
31    type OwnedType = PaddedFrame<I::OwnedType>;
32
33    fn exposed_data(&self) -> &[u8] {
34        let exposed_data = self.inner.exposed_data();
35        &exposed_data[..exposed_data.len() - 4 - self.padding_size]
36    }
37
38    fn whole_data(&self) -> &[u8] {
39        self.inner.whole_data()
40    }
41
42    fn to_owned_frame(&self) -> Self::OwnedType {
43        PaddedFrame {
44            inner: self.inner.to_owned_frame(),
45            padding_size: self.padding_size,
46        }
47    }
48}
49
50impl<I: FrameReader + Clone> Clone for PaddedFrame<I> {
51    fn clone(&self) -> Self {
52        PaddedFrame {
53            inner: self.inner.clone(),
54            padding_size: self.padding_size,
55        }
56    }
57}
58
59/// Padded frame builder
60pub struct PaddedFrameBuilder<I: FrameBuilder> {
61    inner: I,
62    minimum_size: usize,
63}
64
65impl<I: FrameBuilder> PaddedFrameBuilder<I> {
66    pub fn new(inner: I, minimum_size: usize) -> PaddedFrameBuilder<I> {
67        PaddedFrameBuilder {
68            inner,
69            minimum_size,
70        }
71    }
72
73    pub fn inner_mut(&mut self) -> &mut I {
74        &mut self.inner
75    }
76
77    pub fn set_minimum_size(&mut self, minimum_size: usize) {
78        self.minimum_size = minimum_size;
79    }
80}
81
82impl<I: FrameBuilder> FrameBuilder for PaddedFrameBuilder<I> {
83    type OwnedFrameType = PaddedFrame<Bytes>;
84
85    fn write_to<W: io::Write>(&self, writer: &mut W) -> Result<usize, Error> {
86        let inner_size = self.inner.write_to(writer)?;
87
88        let padding_size = if inner_size < self.minimum_size {
89            let required_padding = self.minimum_size - inner_size;
90            for _i in 0..required_padding {
91                writer.write_u8(0)?;
92            }
93            required_padding
94        } else {
95            0
96        };
97
98        writer.write_u32::<LittleEndian>(padding_size as u32)?;
99        Ok(inner_size + padding_size + 4)
100    }
101
102    fn write_into(&self, into: &mut [u8]) -> Result<usize, Error> {
103        let inner_size = self.inner.write_into(into)?;
104
105        let padding_size = if inner_size < self.minimum_size {
106            let required_padding = self.minimum_size - inner_size;
107            check_into_size(inner_size + required_padding, into)?;
108
109            for i in 0..required_padding {
110                into[inner_size + i] = 0;
111            }
112            required_padding
113        } else {
114            0
115        };
116
117        let total_size = inner_size + padding_size + 4;
118        check_into_size(padding_size, into)?;
119        (&mut into[inner_size + padding_size..]).write_u32::<LittleEndian>(padding_size as u32)?;
120
121        Ok(total_size)
122    }
123
124    fn expected_size(&self) -> Option<usize> {
125        self.inner.expected_size().map(|inner_size| {
126            if inner_size < self.minimum_size {
127                self.minimum_size + 4
128            } else {
129                inner_size + 4
130            }
131        })
132    }
133
134    fn as_owned_frame(&self) -> Self::OwnedFrameType {
135        PaddedFrame::new(self.as_bytes()).expect("Couldn't read just-created frame")
136    }
137}
138
139#[cfg(test)]
140mod tests {
141    use super::*;
142    use crate::framing::assert_builder_equals;
143
144    #[test]
145    fn can_build_and_read() -> anyhow::Result<()> {
146        let builder = PaddedFrameBuilder::new(Bytes::from(vec![1; 10]), 0);
147        assert_builder_equals(&builder)?;
148
149        let frame = PaddedFrame::new(builder.as_bytes())?;
150        assert_eq!(vec![1; 10], frame.exposed_data());
151
152        let builder = PaddedFrameBuilder::new(Bytes::from(vec![1; 10]), 10);
153        assert_builder_equals(&builder)?;
154        let frame = PaddedFrame::new(builder.as_bytes())?;
155        assert_eq!(0, frame.padding_size);
156        assert_eq!(vec![1; 10], frame.exposed_data());
157
158        let builder = PaddedFrameBuilder::new(Bytes::from(vec![1; 10]), 20);
159        assert_builder_equals(&builder)?;
160        let frame = PaddedFrame::new(builder.as_bytes())?;
161        assert_eq!(vec![1; 10], frame.exposed_data());
162        assert_eq!(10, frame.padding_size);
163        assert!(frame.whole_data().len() > 20);
164
165        Ok(())
166    }
167
168    #[test]
169    fn can_build_to_owned() {
170        let builder = PaddedFrameBuilder::new(Bytes::from(vec![1; 10]), 0);
171        let frame = builder.as_owned_frame();
172        assert_eq!(vec![1; 10], frame.exposed_data());
173    }
174}