ddk_messages/segmentation/
segment_reader.rs

1//! Module helping with processing message segmentation related messages.
2
3use super::{SegmentChunk, SegmentStart, MAX_CHUNK_SIZE, MAX_SEGMENTS, MAX_START_DATA_SIZE};
4
5/// Struct helping with processing message segmentation related messages.
6pub struct SegmentReader {
7    cur_data: Vec<u8>,
8    remaining_segments: u16,
9}
10
11#[derive(Debug)]
12/// An error that occured while processing message segmentation related messages.
13pub enum Error {
14    /// The reader is in a state that is invalid.
15    InvalidState(String),
16    /// A parameter received by the reader was not in accordance with its state.
17    InvalidParameter(String),
18}
19
20impl std::fmt::Display for Error {
21    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
22        match *self {
23            Error::InvalidState(ref s) => write!(f, "Invalid state {s}"),
24            Error::InvalidParameter(ref s) => write!(f, "Invalid parameters were provided: {s}"),
25        }
26    }
27}
28
29#[cfg(not(feature = "no-std"))]
30impl std::error::Error for Error {
31    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
32        match self {
33            Error::InvalidState(_) => None,
34            Error::InvalidParameter(_) => None,
35        }
36    }
37}
38
39impl Default for SegmentReader {
40    fn default() -> Self {
41        Self::new()
42    }
43}
44
45impl SegmentReader {
46    /// Returns a new instance of [`Self`].
47    pub fn new() -> Self {
48        SegmentReader {
49            cur_data: Vec::new(),
50            remaining_segments: 0,
51        }
52    }
53
54    /// Reset the state of the reader
55    pub fn reset(&mut self) {
56        self.cur_data = Vec::new();
57        self.remaining_segments = 0;
58    }
59
60    /// Whether the reader is waiting for an incoming chunk.
61    pub fn expecting_chunk(&self) -> bool {
62        self.remaining_segments != 0
63    }
64
65    /// Process a [`super::SegmentStart`] message.
66    pub fn process_segment_start(&mut self, segment_start: SegmentStart) -> Result<(), Error> {
67        if !self.cur_data.is_empty() {
68            return Err(Error::InvalidState(
69                "Received segment start while cur data buffer is not empty.".to_string(),
70            ));
71        }
72
73        if segment_start.nb_segments < 2 || segment_start.nb_segments > (MAX_SEGMENTS as u16) {
74            return Err(Error::InvalidParameter(
75                "Segment start must specify at least two chunks and maximum a thousand."
76                    .to_string(),
77            ));
78        }
79
80        if segment_start.data.len() < MAX_START_DATA_SIZE {
81            return Err(Error::InvalidParameter(
82                "Segment start data should be filled to its maximum capacity.".to_string(),
83            ));
84        }
85
86        let SegmentStart { nb_segments, data } = segment_start;
87
88        self.remaining_segments = nb_segments - 1;
89
90        self.cur_data = data;
91
92        Ok(())
93    }
94
95    /// Process a [`super::SegmentChunk`] message.
96    pub fn process_segment_chunk(
97        &mut self,
98        mut segment_chunk: SegmentChunk,
99    ) -> Result<Option<Vec<u8>>, Error> {
100        if self.cur_data.is_empty() {
101            return Err(Error::InvalidState(
102                "Received segment chunk while cur data buffer is empty.".to_string(),
103            ));
104        }
105
106        if self.remaining_segments > 1 && segment_chunk.data.len() != MAX_CHUNK_SIZE {
107            return Err(Error::InvalidParameter(
108                "Receive non final segment chunk that was not not filled.".to_string(),
109            ));
110        }
111
112        self.cur_data.append(&mut segment_chunk.data);
113        self.remaining_segments -= 1;
114
115        if self.remaining_segments == 0 {
116            let mut res = Vec::new();
117            std::mem::swap(&mut self.cur_data, &mut res);
118            Ok(Some(res))
119        } else {
120            Ok(None)
121        }
122    }
123}
124
125#[cfg(test)]
126mod tests {
127    use crate::segmentation::MAX_DATA_SIZE;
128
129    use super::*;
130
131    fn segments() -> (SegmentStart, Vec<SegmentChunk>) {
132        let mut buf = Vec::new();
133        buf.resize(MAX_DATA_SIZE * 4, 1);
134        super::super::get_segments(buf, 2)
135    }
136
137    #[test]
138    fn read_segments_test() {
139        let mut segment_reader = SegmentReader::new();
140        let (segment_start, segment_chunks) = segments();
141
142        assert!(!segment_reader.expecting_chunk());
143
144        segment_reader
145            .process_segment_start(segment_start)
146            .expect("to be able to process the segment start");
147
148        assert!(segment_reader.expecting_chunk());
149
150        for chunk in segment_chunks {
151            assert!(segment_reader.expecting_chunk());
152            segment_reader
153                .process_segment_chunk(chunk)
154                .expect("to be able to process the segment chunk");
155        }
156
157        assert!(!segment_reader.expecting_chunk());
158    }
159
160    #[test]
161    fn chunk_no_start_fails_test() {
162        let mut segment_reader = SegmentReader::new();
163        let (_, mut segment_chunks) = segments();
164        segment_reader
165            .process_segment_chunk(segment_chunks.pop().unwrap())
166            .expect_err("Should not process a chunk without having had a start first.");
167    }
168
169    #[test]
170    fn start_not_finished_previous_fails_test() {
171        let mut segment_reader = SegmentReader::new();
172        let (segment_start, segment_chunks) = segments();
173        let (segment_start2, _) = segments();
174        segment_reader
175            .process_segment_start(segment_start)
176            .expect("to be able to process the first segment start");
177        segment_reader
178            .process_segment_chunk(segment_chunks[0].clone())
179            .expect("to be able to process the segment chunk");
180        segment_reader
181            .process_segment_start(segment_start2)
182            .expect_err("should not process new start before finishing previous segment");
183    }
184
185    #[test]
186    fn start_reset_start_accepted_test() {
187        let mut segment_reader = SegmentReader::new();
188        let (segment_start, _) = segments();
189        let (segment_start2, _) = segments();
190        segment_reader
191            .process_segment_start(segment_start)
192            .expect("to be able to process the first segment start");
193        segment_reader.reset();
194        segment_reader
195            .process_segment_start(segment_start2)
196            .expect("to be able to process the same segment start after reset");
197    }
198
199    #[test]
200    fn too_few_chunk_in_start_fails_test() {
201        let mut segment_reader = SegmentReader::new();
202        let (mut segment_start, _) = segments();
203        segment_start.nb_segments = 1;
204        segment_reader
205            .process_segment_start(segment_start)
206            .expect_err("should not accept segment with less than 2 elements");
207    }
208
209    #[test]
210    fn too_many_chunks_in_start_fails_test() {
211        let mut segment_reader = SegmentReader::new();
212        let (mut segment_start, _) = segments();
213        segment_start.nb_segments = 1001;
214        segment_reader
215            .process_segment_start(segment_start)
216            .expect_err("should not accept segments with more than 1000 elements");
217    }
218
219    #[test]
220    fn segment_start_not_full_fails_test() {
221        let mut segment_reader = SegmentReader::new();
222        let (mut segment_start, _) = segments();
223        segment_start.data.pop();
224
225        segment_reader
226            .process_segment_start(segment_start)
227            .expect_err("Should error on non full segment start message.");
228    }
229
230    #[test]
231    fn non_final_chunk_not_full_fails_test() {
232        let mut segment_reader = SegmentReader::new();
233        let (segment_start, mut segment_chunks) = segments();
234        segment_reader
235            .process_segment_start(segment_start)
236            .expect("to be able to process the segment start");
237
238        segment_chunks[0].data.pop();
239
240        segment_reader
241            .process_segment_chunk(segment_chunks[0].clone())
242            .expect_err("should not accept not full segment that is not the last one");
243    }
244}
ddk_messages/segmentation/segment_reader.rs

ddk_messages/segmentation/
segment_reader.rs
