Skip to main content

leodos_protocols/datalink/spp/segmentation/
reassembler.rs

1use crate::network::spp::SequenceCount;
2use crate::network::spp::SequenceFlag;
3use crate::network::spp::SpacePacket;
4
5/// An error that can occur during the reassembly of segmented packets.
6#[derive(Debug, Copy, Clone, Eq, PartialEq, thiserror::Error)]
7pub enum ReassemblyError {
8    /// A `Continuation` or `Last` packet was received, but no `First` packet
9    /// was processed to start the sequence.
10    #[error("Continuation or Last packet received before First packet")]
11    ContinuationBeforeFirst,
12    /// A `First` packet was received, but a reassembly for this sequence
13    /// is already in progress.
14    #[error("Duplicate First packet received")]
15    DuplicateFirstPacket,
16    /// A packet was received with a `SequenceCount` that was not the expected
17    /// next value in the sequence, indicating a lost packet.
18    #[error("Packet out of order: expected {expected}, got {got}")]
19    PacketOutOfOrder {
20        /// The expected sequence count value.
21        expected: u16,
22        /// The actual sequence count value received.
23        got: u16,
24    },
25    /// A `Unsegmented` packet was passed to the reassembler, which only
26    /// handles segmented sequences.
27    #[error("Unexpected Unsegmented packet received")]
28    UnexpectedUnsegmentedPacket,
29    /// The provided user buffer is too small to hold the incoming data.
30    #[error("User buffer too small for reassembled data")]
31    BufferTooSmall,
32}
33
34/// The state of a reassembly process.
35#[derive(Debug, Clone, Eq, PartialEq)]
36pub enum ReassemblyState<'a> {
37    /// More packets are needed to complete the data block.
38    InProgress,
39    /// The `Last` packet has been received and the data block is complete.
40    /// Contains a slice of the user-provided buffer holding the reassembled data.
41    Complete(&'a [u8]),
42}
43
44/// A stateful helper that reassembles data from segmented packets into a user-provided buffer.
45///
46/// This version is allocator-free. It does not own its buffer; it only borrows it mutably.
47/// An application would typically manage a pool of these `Reassembler` instances.
48pub struct Reassembler<'a> {
49    buffer: &'a mut [u8],
50    write_position: usize,
51    expected_sequence_count: SequenceCount,
52    is_started: bool,
53}
54
55impl<'a> Reassembler<'a> {
56    /// Creates a new `Reassembler` that will write into the provided buffer.
57    pub fn new(buffer: &'a mut [u8]) -> Self {
58        Self {
59            buffer,
60            write_position: 0,
61            expected_sequence_count: SequenceCount::new(),
62            is_started: false,
63        }
64    }
65
66    /// Resets the reassembler to its initial state, allowing its buffer to be reused.
67    pub fn reset(&mut self) {
68        self.write_position = 0;
69        self.is_started = false;
70        // The buffer itself is not cleared for performance; it will be overwritten.
71    }
72
73    /// Processes an incoming `SpacePacket` and writes its data into the buffer.
74    pub fn process_packet(
75        &'a mut self,
76        packet: &SpacePacket,
77    ) -> Result<ReassemblyState<'a>, ReassemblyError> {
78        let payload = packet.data_field();
79
80        // Check if the payload will fit
81        if self.write_position + payload.len() > self.buffer.len() {
82            return Err(ReassemblyError::BufferTooSmall);
83        }
84
85        match packet.sequence_flag() {
86            SequenceFlag::First => {
87                if self.is_started {
88                    return Err(ReassemblyError::DuplicateFirstPacket);
89                }
90                self.reset(); // Ensure we're in a clean state
91                self.is_started = true;
92                self.expected_sequence_count = packet.sequence_count();
93                self.expected_sequence_count.increment();
94            }
95            SequenceFlag::Continuation | SequenceFlag::Last => {
96                if !self.is_started {
97                    return Err(ReassemblyError::ContinuationBeforeFirst);
98                }
99                if packet.sequence_count() != self.expected_sequence_count {
100                    return Err(ReassemblyError::PacketOutOfOrder {
101                        expected: self.expected_sequence_count.value(),
102                        got: packet.sequence_count().value(),
103                    });
104                }
105                self.expected_sequence_count.increment();
106            }
107            SequenceFlag::Unsegmented => return Err(ReassemblyError::UnexpectedUnsegmentedPacket),
108        }
109
110        // If all checks pass, write the data
111        let write_end = self.write_position + payload.len();
112        self.buffer[self.write_position..write_end].copy_from_slice(payload);
113        self.write_position = write_end;
114
115        if packet.sequence_flag() == SequenceFlag::Last {
116            let data = &self.buffer[..self.write_position];
117            Ok(ReassemblyState::Complete(data))
118        } else {
119            Ok(ReassemblyState::InProgress)
120        }
121    }
122}