librqbit_peer_protocol/
double_buf.rs

1use std::io::IoSlice;
2
3/// A helper for working with a buffer split into 2.
4/// You can advance it forward (like you would do with buf=&buf[idx..])
5#[derive(Clone, Copy)]
6pub struct DoubleBufHelper<'a> {
7    buf_0: &'a [u8],
8    buf_1: &'a [u8],
9}
10
11impl<'a> DoubleBufHelper<'a> {
12    pub fn new(buf: &'a [u8], buf2: &'a [u8]) -> Self {
13        Self {
14            buf_0: buf,
15            buf_1: buf2,
16        }
17    }
18
19    /// Consume len bytes and return them as 2 slices. Advances the buffer forward if successful.
20    /// On error returns how many bytes are missing.
21    pub fn consume_variable(&mut self, len: usize) -> Result<(&'a [u8], &'a [u8]), usize> {
22        let available = self.buf_0.len() + self.buf_1.len();
23        if available < len {
24            return Err(len - available);
25        }
26
27        let first_len = self.buf_0.len().min(len);
28        let (first_consumed, first_remaining) = self.buf_0.split_at(first_len);
29
30        let second_len = (len - first_len).min(self.buf_1.len()); // the .min() here is just for split_at() to be optimized without panic
31        let (second_consumed, second_remaining) = self.buf_1.split_at(second_len);
32
33        self.buf_0 = first_remaining;
34        self.buf_1 = second_remaining;
35
36        Ok((first_consumed, second_consumed))
37    }
38
39    /// Read N bytes and advance the buffer by N if successful.
40    /// Error returns how many missing bytes are there.
41    pub fn consume<const N: usize>(&mut self) -> Result<[u8; N], usize> {
42        match (self.buf_0.len(), self.buf_1.len()) {
43            (l, _) if l >= N => {
44                let (chunk, rem) = self.buf_0.split_at(N);
45                self.buf_0 = rem;
46                return Ok(chunk.try_into().unwrap());
47            }
48            (0, l) if l >= N => {
49                let (chunk, rem) = self.buf_1.split_at(N);
50                self.buf_1 = rem;
51                return Ok(chunk.try_into().unwrap());
52            }
53            _ => {}
54        }
55
56        let mut res = [0u8; N];
57
58        let first = self.buf_0.len().min(N);
59        let second = self.buf_1.len().min(N.saturating_sub(first));
60
61        let missing = N - first - second;
62        if missing > 0 {
63            return Err(missing);
64        }
65
66        res[..first].copy_from_slice(&self.buf_0[..first]);
67        res[first..].copy_from_slice(&self.buf_1[..second]);
68        self.buf_0 = &self.buf_0[first..];
69        self.buf_1 = &self.buf_1[second..];
70        Ok(res)
71    }
72
73    pub fn get(&self) -> [&'a [u8]; 2] {
74        [self.buf_0, self.buf_1]
75    }
76
77    /// Read 4 big endian bytes and advance the buffer by 4 if successful.
78    /// Error returns how many missing bytes are there.
79    pub fn read_u32_be(&mut self) -> Result<u32, usize> {
80        let data = self.consume::<4>()?;
81        Ok(u32::from_be_bytes(data))
82    }
83
84    /// Read 1 byte and advance. Returns 1
85    pub fn read_u8(&mut self) -> Option<u8> {
86        let b = if !self.buf_0.is_empty() {
87            &mut self.buf_0
88        } else if !self.buf_1.is_empty() {
89            &mut self.buf_1
90        } else {
91            return None;
92        };
93        let value = b[0];
94        *b = &b[1..];
95        Some(value)
96    }
97
98    /// Get a contiguous slice at the start if it exists.
99    pub fn get_contiguous(&self, len: usize) -> Option<&'a [u8]> {
100        match (self.buf_0.len(), self.buf_1.len()) {
101            (l, _) if l >= len => Some(&self.buf_0[..len]),
102            (0, l) if l >= len => Some(&self.buf_1[..len]),
103            _ => None,
104        }
105    }
106
107    pub fn len(&self) -> usize {
108        self.buf_0.len() + self.buf_1.len()
109    }
110
111    pub fn is_empty(&self) -> bool {
112        self.buf_0.len() == 0 && self.buf_1.len() == 0
113    }
114
115    /// Advance it forward. If it was a single buffer this would be equivalent to buf=&buf[idx..]).
116    /// If offset is too large, will set itself empty.
117    pub fn advance(&mut self, offset: usize) {
118        let buf_0_adv = self.buf_0.len().min(offset);
119        self.buf_0 = &self.buf_0[buf_0_adv..];
120        let buf_1_adv = (offset - buf_0_adv).min(self.buf_1.len());
121        self.buf_1 = &self.buf_1[buf_1_adv..];
122    }
123
124    pub fn with_max_len(&self, max_len: usize) -> DoubleBufHelper<'a> {
125        let buf_0_len = self.buf_0.len().min(max_len);
126        let buf_1_len = (max_len - buf_0_len).min(self.buf_1.len());
127        DoubleBufHelper {
128            buf_0: &self.buf_0[..buf_0_len],
129            buf_1: &self.buf_1[..buf_1_len],
130        }
131    }
132
133    pub fn as_ioslices(&self, len_limit: usize) -> [IoSlice<'a>; 2] {
134        let buf_0_len = self.buf_0.len().min(len_limit);
135        let buf_1_len = (len_limit - buf_0_len).min(self.buf_1.len());
136        [
137            IoSlice::new(&self.buf_0[..buf_0_len]),
138            IoSlice::new(&self.buf_1[..buf_1_len]),
139        ]
140    }
141}
142
143#[cfg(test)]
144mod tests {
145    use crate::double_buf::DoubleBufHelper;
146
147    #[test]
148    fn test_get_contiguous() {
149        let d = DoubleBufHelper::new(&[], &[]);
150        assert_eq!(d.get_contiguous(0).unwrap(), &[]);
151        assert_eq!(d.get_contiguous(1), None);
152
153        let d = DoubleBufHelper::new(&[42u8; 43], &[]);
154        assert_eq!(d.get_contiguous(42).unwrap(), &[42u8; 42]);
155        assert_eq!(d.get_contiguous(43).unwrap(), &[42u8; 43]);
156        assert_eq!(d.get_contiguous(44), None);
157
158        let d = DoubleBufHelper::new(&[], &[42u8; 43]);
159        assert_eq!(d.get_contiguous(42).unwrap(), &[42u8; 42]);
160        assert_eq!(d.get_contiguous(43).unwrap(), &[42u8; 43]);
161        assert_eq!(d.get_contiguous(44), None);
162
163        let d = DoubleBufHelper::new(&[42u8; 43], &[43u8; 52]);
164        assert_eq!(d.get_contiguous(42).unwrap(), &[42u8; 42]);
165        assert_eq!(d.get_contiguous(43).unwrap(), &[42u8; 43]);
166        assert_eq!(d.get_contiguous(44), None);
167
168        let d = DoubleBufHelper::new(&[], &[43u8; 52]);
169        assert_eq!(d.get_contiguous(42).unwrap(), &[43u8; 42]);
170        assert_eq!(d.get_contiguous(43).unwrap(), &[43u8; 43]);
171        assert_eq!(d.get_contiguous(52).unwrap(), &[43u8; 52]);
172        assert_eq!(d.get_contiguous(53), None);
173
174        let d = DoubleBufHelper::new(&[42u8; 43], &[43u8; 52]);
175        assert_eq!(d.get_contiguous(42).unwrap(), &[42u8; 42]);
176    }
177
178    #[test]
179    fn test_consume() {
180        for (first, second) in [(&[0, 1][..], &[][..]), (&[0], &[1]), (&[], &[0, 1])] {
181            let mut d = DoubleBufHelper::new(first, second);
182            assert_eq!(d.consume::<0>(), Ok([]));
183            assert_eq!(d.len(), 2);
184
185            let mut d = DoubleBufHelper::new(first, second);
186            assert_eq!(d.consume::<1>(), Ok([0]));
187            assert_eq!(d.len(), 1);
188
189            let mut d = DoubleBufHelper::new(first, second);
190            assert_eq!(d.consume::<2>(), Ok([0, 1]));
191            assert_eq!(d.len(), 0);
192
193            let mut d = DoubleBufHelper::new(first, second);
194            assert_eq!(d.consume::<3>(), Err(1));
195            assert_eq!(d.len(), 2);
196        }
197    }
198
199    #[test]
200    fn test_consume_variable() {
201        let mut d = DoubleBufHelper::new(&[], &[]);
202        assert_eq!(d.consume_variable(0), Ok((&[][..], &[][..])));
203        assert_eq!(d.len(), 0);
204        assert_eq!(d.consume_variable(1), Err(1));
205
206        let mut d = DoubleBufHelper::new(&[0, 1], &[]);
207        assert_eq!(d.consume_variable(0), Ok((&[][..], &[][..])));
208        assert_eq!(d.len(), 2);
209
210        let mut d = DoubleBufHelper::new(&[0, 1], &[]);
211        assert_eq!(d.consume_variable(1), Ok((&[0][..], &[][..])));
212        assert_eq!(d.len(), 1);
213        assert_eq!(d.buf_0, &[1]);
214        assert_eq!(d.buf_1, &[]);
215
216        let mut d = DoubleBufHelper::new(&[0, 1], &[]);
217        assert_eq!(d.consume_variable(2), Ok((&[0, 1][..], &[][..])));
218        assert_eq!(d.len(), 0);
219        assert_eq!(d.buf_0, &[]);
220        assert_eq!(d.buf_1, &[]);
221
222        let mut d = DoubleBufHelper::new(&[0, 1], &[]);
223        assert_eq!(d.consume_variable(3), Err(1));
224        assert_eq!(d.len(), 2);
225        assert_eq!(d.buf_0, &[0, 1]);
226        assert_eq!(d.buf_1, &[]);
227
228        let mut d = DoubleBufHelper::new(&[0], &[1]);
229        assert_eq!(d.consume_variable(0), Ok((&[][..], &[][..])));
230        assert_eq!(d.len(), 2);
231
232        let mut d = DoubleBufHelper::new(&[0], &[1]);
233        assert_eq!(d.consume_variable(1), Ok((&[0][..], &[][..])));
234        assert_eq!(d.len(), 1);
235        assert_eq!(d.buf_0, &[]);
236        assert_eq!(d.buf_1, &[1]);
237
238        let mut d = DoubleBufHelper::new(&[0], &[1]);
239        assert_eq!(d.consume_variable(2), Ok((&[0][..], &[1][..])));
240        assert_eq!(d.len(), 0);
241        assert_eq!(d.buf_0, &[]);
242        assert_eq!(d.buf_1, &[]);
243
244        let mut d = DoubleBufHelper::new(&[0], &[1]);
245        assert_eq!(d.consume_variable(3), Err(1));
246        assert_eq!(d.len(), 2);
247        assert_eq!(d.buf_0, &[0]);
248        assert_eq!(d.buf_1, &[1]);
249    }
250
251    #[test]
252    fn test_advance_out_of_bounds_0() {
253        let mut d = DoubleBufHelper::new(&[], &[]);
254        d.advance(1);
255        assert!(d.is_empty());
256        assert_eq!(d.buf_0, &[]);
257        assert_eq!(d.buf_1, &[]);
258    }
259
260    #[test]
261    fn test_advance_out_of_bounds_1() {
262        let mut d = DoubleBufHelper::new(&[42], &[]);
263        d.advance(2);
264        assert!(d.is_empty());
265        assert_eq!(d.buf_0, &[]);
266        assert_eq!(d.buf_1, &[]);
267    }
268
269    #[test]
270    fn test_advance_out_of_bounds_2() {
271        let mut d = DoubleBufHelper::new(&[42], &[43]);
272        d.advance(3);
273        assert!(d.is_empty());
274        assert_eq!(d.buf_0, &[]);
275        assert_eq!(d.buf_1, &[]);
276    }
277
278    #[test]
279    fn test_advance_out_of_bounds_3() {
280        let mut d = DoubleBufHelper::new(&[], &[42, 43]);
281        d.advance(3);
282        assert!(d.is_empty());
283        assert_eq!(d.buf_0, &[]);
284        assert_eq!(d.buf_1, &[]);
285    }
286
287    #[test]
288    fn test_advance() {
289        let mut d = DoubleBufHelper::new(&[], &[]);
290        d.advance(0);
291        assert_eq!(d.len(), 0);
292        assert_eq!(d.consume_variable(1), Err(1));
293
294        let mut d = DoubleBufHelper::new(&[0, 1], &[]);
295        d.advance(0);
296        assert_eq!(d.len(), 2);
297
298        let mut d = DoubleBufHelper::new(&[0, 1], &[]);
299        d.advance(1);
300        assert_eq!(d.len(), 1);
301        assert_eq!(d.buf_0, &[1]);
302        assert_eq!(d.buf_1, &[]);
303
304        let mut d = DoubleBufHelper::new(&[0, 1], &[]);
305        d.advance(2);
306        assert_eq!(d.len(), 0);
307        assert_eq!(d.buf_0, &[]);
308        assert_eq!(d.buf_1, &[]);
309
310        let mut d = DoubleBufHelper::new(&[0], &[1]);
311        d.advance(0);
312        assert_eq!(d.len(), 2);
313
314        let mut d = DoubleBufHelper::new(&[0], &[1]);
315        d.advance(1);
316        assert_eq!(d.len(), 1);
317        assert_eq!(d.buf_0, &[]);
318        assert_eq!(d.buf_1, &[1]);
319
320        let mut d = DoubleBufHelper::new(&[0], &[1]);
321        d.advance(2);
322        assert_eq!(d.len(), 0);
323        assert_eq!(d.buf_0, &[]);
324        assert_eq!(d.buf_1, &[]);
325    }
326}
librqbit_peer_protocol/double_buf.rs

librqbit_peer_protocol/
double_buf.rs
