librqbit_utp/raw/
selective_ack.rs

1use bitvec::{BitArr, order::Lsb0, slice::BitSlice};
2
3use crate::constants::SACK_DEPTH;
4
5type SelectiveAckData = BitArr!(for SACK_DEPTH, in u8, Lsb0);
6
7#[derive(Debug, Clone, Copy, PartialEq, Eq)]
8pub struct SelectiveAck {
9    data: SelectiveAckData,
10    len: usize,
11}
12
13impl SelectiveAck {
14    pub fn new(unacked: impl Iterator<Item = usize>) -> Self {
15        let mut data = SelectiveAckData::default();
16        for idx in unacked.take_while(|i| *i < SACK_DEPTH) {
17            data.set(idx, true);
18        }
19        Self {
20            data,
21            len: SACK_DEPTH,
22        }
23    }
24
25    #[cfg(test)]
26    pub fn new_test(acked: impl IntoIterator<Item = usize>) -> Option<Self> {
27        let mut data = SelectiveAckData::default();
28        let mut count = 0;
29        for idx in acked {
30            data.get_mut(idx).unwrap().set(true);
31            count += 1;
32        }
33        if count > 0 {
34            Some(Self {
35                data,
36                len: SACK_DEPTH,
37            })
38        } else {
39            None
40        }
41    }
42
43    #[allow(clippy::len_without_is_empty)]
44    pub fn len(&self) -> usize {
45        self.len
46    }
47
48    pub fn as_bytes(&self) -> &[u8] {
49        self.data.as_raw_slice()
50    }
51
52    pub fn deserialize(bytes: &[u8]) -> Self {
53        // The spec says the len must be a multiple of 4, but there's a ton of messages
54        // in the wild that are 1 bytes long (probably coming from libutp). So we can
55        // thus deserialize any payload.
56        //
57        // If it's longer than 8 bytes (unlikely), it will truncate the end, which is fine, as
58        // we'll just resend that data if anything.
59        let len = bytes.len().min(std::mem::size_of::<SelectiveAckData>());
60        let mut data = SelectiveAckData::default();
61        data.as_raw_mut_slice()[..len].copy_from_slice(&bytes[..len]);
62        Self {
63            data,
64            len: bytes.len() * 8,
65        }
66    }
67
68    pub fn as_bitslice(&self) -> &BitSlice<u8> {
69        self.data.as_bitslice()
70    }
71
72    pub fn iter(&self) -> impl Iterator<Item = bool> + '_ {
73        self.data.iter().map(|b| *b)
74    }
75}
76
77#[cfg(test)]
78mod tests {
79    use std::num::NonZeroUsize;
80
81    use crate::{
82        message::UtpMessage,
83        raw::{Type::ST_DATA, selective_ack::SelectiveAck},
84        stream_rx::OutOfOrderQueue,
85    };
86
87    fn asm() -> OutOfOrderQueue {
88        OutOfOrderQueue::new(NonZeroUsize::new(65).unwrap())
89    }
90
91    fn msg() -> UtpMessage {
92        UtpMessage::new_test(
93            crate::raw::UtpHeader {
94                htype: ST_DATA,
95                ..Default::default()
96            },
97            b"a",
98        )
99    }
100
101    #[test]
102    fn test_empty_is_none() {
103        let asm = asm();
104        assert!(asm.selective_ack().is_none());
105    }
106
107    #[test]
108    fn test_holes() {
109        let mut asm = asm();
110        asm.add_remove(msg(), 8).unwrap();
111        asm.add_remove(msg(), 1).unwrap();
112        asm.add_remove(msg(), 2).unwrap();
113        asm.add_remove(msg(), 64).unwrap();
114
115        let sack = asm.selective_ack().unwrap();
116        assert_eq!(
117            sack.data.as_raw_slice(),
118            [0b1000_0011, 0, 0, 0, 0, 0, 0, 0b1000_0000]
119        );
120
121        assert_eq!(sack, SelectiveAck::deserialize(sack.as_bytes()));
122    }
123}
librqbit_utp/raw/selective_ack.rs

librqbit_utp/raw/
selective_ack.rs
