rsrs_core/fragmentation/
fragment.rs

1use recode::bytes::BytesMut;
2
3use crate::frame::Payload;
4use crate::Frame;
5use super::payload::PayloadFrame;
6
7/// A helper that is used to fragment a frame that would not fit in a single
8/// transmission unit.
9#[derive(Debug)]
10pub struct Fragmenter {
11    mtu: usize,
12}
13
14impl Fragmenter {
15    /// Creates a new instance of [`Fragmenter`].
16    ///
17    /// # Parameters
18    /// * `mtu` - The maximum transmission unit (MTU) to fragment frames into.
19    #[inline(always)]
20    pub const fn new(mtu: usize) -> Self {
21        Self { mtu }
22    }
23
24    /// Fragments the provided frame into a stream of frames with the configured
25    /// MTU.
26    #[inline]
27    pub fn fragment(&self, mut frame: Frame) -> impl Iterator<Item = Frame> {
28        let payload = self.trim_frame(&mut frame);
29
30        std::iter::once(frame).chain(PayloadIterator {
31            payload,
32            mtu: self.mtu,
33        })
34    }
35
36    #[inline]
37    fn trim_frame(&self, frame: &mut Frame) -> Option<(BytesMut, BytesMut)> {
38        match frame {
39            | Frame::RequestResponse(f) => f.trim_to(self.mtu),
40            | Frame::RequestFNF(f) => f.trim_to(self.mtu),
41            | Frame::RequestStream(f) => f.trim_to(self.mtu),
42            | Frame::RequestChannel(f) => f.trim_to(self.mtu),
43            | Frame::Payload(f) => f.trim_to(self.mtu),
44            | _ => Default::default(),
45        }
46    }
47}
48
49#[derive(Debug)]
50struct PayloadIterator {
51    payload: Option<(BytesMut, BytesMut)>,
52    mtu: usize,
53}
54
55impl Iterator for PayloadIterator {
56    type Item = Frame;
57
58    fn next(&mut self) -> Option<Self::Item> {
59        let Some((metadata, data)) = self.payload.take() else {
60            return None;
61        };
62
63        let mut payload = Payload::builder()
64            .metadata(metadata)
65            .data(data)
66            .build()
67            .unwrap();
68
69        if let Some((metadata, data)) = payload.trim_to(self.mtu) {
70            self.payload = Some((metadata, data));
71            payload.set_follow(true);
72        } else {
73            payload.set_follow(false);
74        }
75
76        Some(payload.into())
77    }
78}