stun_bytes/
lib.rs

1#![no_std]
2
3use r_ex::prelude::*;
4
5pub struct ByteMsg<'a> {
6    buf: &'a [u8],
7}
8
9impl<'a> From<&'a [u8]> for ByteMsg<'a> {
10    fn from(buf: &'a [u8]) -> Self {
11        Self {
12            buf
13        }
14    }
15}
16
17impl<'a> ByteMsg<'a> {
18    pub fn typ(&self) -> Option<&'a [u8; 2]> {
19        self.buf.carved()
20    }
21
22    pub fn len(&self) -> Option<&'a [u8; 2]> {
23        self.buf.carve(2)
24    }
25
26    #[cfg(feature = "cookie")]
27    pub fn cookie(&self) -> Option<&'a [u8; 4]> {
28        self.buf.carve(4)
29    }
30
31    #[cfg(feature = "cookie")]
32    pub fn tid(&self) -> Option<&'a [u8; 12]> {
33        self.buf.carve(8)
34    }
35
36    #[cfg(not(feature = "cookie"))]
37    pub fn tid(&self) -> Option<&'a [u8; 16]> {
38        self.buf.carve(4)
39    }
40
41    pub fn attrs(&self) -> Option<&'a [u8]> {
42        let len = self.len().map(u16::from_be_ref)? as usize;
43        self.buf.get(20..20 + len)
44    }
45
46    pub fn attrs_iter(&self) -> ByteAttrIter {
47        ByteAttrIter::from(self.attrs().unwrap_or(&[]))
48    }
49
50    pub fn size(&self) -> usize {
51        self.len()
52            .map(u16::from_be_ref)
53            .map(|len| 20 + len as usize)
54            .unwrap_or(self.buf.len())
55    }
56}
57
58pub struct ByteAttr<'a> {
59    buf: &'a [u8],
60}
61
62impl<'a> From<&'a [u8]> for ByteAttr<'a> {
63    fn from(buf: &'a [u8]) -> Self {
64        Self {
65            buf
66        }
67    }
68}
69
70impl<'a> ByteAttr<'a> {
71    pub fn typ(&self) -> Option<&'a [u8; 2]> {
72        self.buf.carved()
73    }
74
75    pub fn len(&self) -> Option<&'a [u8; 2]> {
76        self.buf.carve(2)
77    }
78
79    pub fn val(&self) -> Option<&'a [u8]> {
80        self.len()
81            .map(u16::from_be_ref)
82            .map(|len| len + 3 & !3)
83            .map(|len| self.buf.get(4..4 + len as usize))
84            .flatten()
85            .or(self.buf.get(4..))
86    }
87}
88
89#[derive(Copy, Clone)]
90pub struct ByteAttrIter<'a> {
91    buf: &'a [u8],
92}
93
94impl<'a> From<&'a [u8]> for ByteAttrIter<'a> {
95    fn from(buf: &'a [u8]) -> Self {
96        Self {
97            buf
98        }
99    }
100}
101
102impl<'a> Iterator for ByteAttrIter<'a> {
103    type Item = ByteAttr<'a>;
104
105    fn next(&mut self) -> Option<Self::Item> {
106        if self.buf.is_empty() { return None; }
107        let attr = ByteAttr::from(self.buf);
108        self.buf = attr.val()
109            .map(<[u8]>::len)
110            .map(|len| 4 + len)
111            .map(|len| self.buf.get(len..))
112            .flatten()
113            .unwrap_or(&[]);
114        Some(attr)
115    }
116}
117
118pub struct ByteMsgMut<'a> {
119    buf: &'a mut [u8],
120}
121
122impl<'a> From<&'a mut [u8]> for ByteMsgMut<'a> {
123    fn from(buf: &'a mut [u8]) -> Self {
124        Self {
125            buf
126        }
127    }
128}
129
130impl<'a> ByteMsgMut<'a> {
131    pub fn typ(&mut self) -> Option<&mut [u8; 2]> {
132        self.buf.carved_mut()
133    }
134
135    pub fn len(&mut self) -> Option<&mut [u8; 2]> {
136        self.buf.carve_mut(2)
137    }
138
139    #[cfg(feature = "cookie")]
140    pub fn cookie(&mut self) -> Option<&mut [u8; 4]> {
141        self.buf.carve_mut(4)
142    }
143
144    #[cfg(feature = "cookie")]
145    pub fn tid(&mut self) -> Option<&mut[u8; 12]> {
146        self.buf.carve_mut(8)
147    }
148
149    #[cfg(not(feature = "cookie"))]
150    pub fn tid(&mut self) -> Option<&mut [u8; 16]> {
151        self.buf.carve_mut(4)
152    }
153
154    pub fn attrs(&mut self) -> Option<&mut [u8]> {
155        let len = self.len().map(|r| u16::from_be_ref(r))? as usize;
156        self.buf.get_mut(20..20 + len)
157    }
158
159    pub fn attrs_iter(&mut self) -> ByteAttrIterMut {
160        ByteAttrIterMut::from(self.attrs().unwrap_or(&mut []))
161    }
162
163    pub fn add_attr(&mut self, typ: &[u8; 2], len: &[u8; 2], val: &[u8]) -> Option<()> {
164        let curr_len = self.len().map(|r| u16::from_be_ref(r))? as usize;
165        let (typ_buf, buf) = self.buf.get_mut(20 + curr_len..)?.splice_mut()?;
166        let (len_buf, buf) = buf.splice_mut()?;
167        let val_buf = buf.get_mut(..val.len())?;
168
169        typ_buf.copy_from(typ);
170        len_buf.copy_from(len);
171        val_buf.copy_from_slice(val);
172
173        self.len()?.set_be((curr_len + 4 + (val.len() + 3) & !3) as u16);
174        Some(())
175    }
176
177    pub fn add_attr2<F: Fn(&mut [u8; 2], &mut [u8; 2], &mut [u8]) -> Option<usize>>(&mut self, callback: F) -> Option<()> {
178        let curr_len = self.len().map(|r| u16::from_be_ref(r))? as usize;
179
180        let (typ_buf, buf) = self.buf.get_mut(20 + curr_len..)?.splice_mut()?;
181        let (len_buf, val_buf) = buf.splice_mut()?;
182
183        let size = callback(typ_buf, len_buf, val_buf)?;
184        self.len()?.set_be((curr_len + 4 + (size + 3) & !3) as u16);
185        Some(())
186    }
187
188    pub fn as_bytes(&mut self) -> &mut [u8] {
189        let len = self.len()
190            .map(|r| u16::from_be_ref(r))
191            .unwrap_or(0) as usize;
192        if self.buf.len() < (len + 3) & !3 {
193            self.buf
194        } else {
195            self.buf.get_mut(0..20 + ((len as usize) + 3) & !3).unwrap_or(&mut [])
196        }
197    }
198}
199
200pub struct ByteAttrMut<'a> {
201    buf: &'a mut [u8],
202}
203
204impl<'a> From<&'a mut [u8]> for ByteAttrMut<'a> {
205    fn from(buf: &'a mut [u8]) -> Self {
206        Self {
207            buf
208        }
209    }
210}
211
212impl<'a> ByteAttrMut<'a> {
213    pub fn typ(&mut self) -> Option<&mut [u8; 2]> {
214        self.buf.carved_mut()
215    }
216
217    pub fn len(&mut self) -> Option<&mut [u8; 2]> {
218        self.buf.carve_mut(2)
219    }
220
221    pub fn val(&mut self) -> Option<&mut [u8]> {
222        let len = (self.len().map(|r| u16::from_be_ref(r))? + 3) & !3;
223        let val = if self.buf.len() > (4 + len) as usize {
224            self.buf.get_mut(4..4 + len as usize)
225        } else {
226            self.buf.get_mut(4..)
227        };
228        val
229    }
230}
231
232pub struct ByteAttrIterMut<'a> {
233    buf: &'a mut [u8],
234}
235
236impl<'a> From<&'a mut [u8]> for ByteAttrIterMut<'a> {
237    fn from(buf: &'a mut [u8]) -> Self {
238        Self {
239            buf
240        }
241    }
242}
243
244impl<'a> Iterator for ByteAttrIterMut<'a> {
245    type Item = ByteAttrMut<'a>;
246
247    fn next(&mut self) -> Option<Self::Item> {
248        if self.buf.is_empty() { return None; }
249        let mut tmp_attr = unsafe {
250            ByteAttrMut::from(core::mem::transmute::<&mut [u8], &mut [u8]>(self.buf))
251        };
252        let declared_val_size = tmp_attr.len().map(|r| u16::from_be_ref(r))?;
253        let total_attr_size = (4 + declared_val_size + 3) & !3;
254        if self.buf.len() > total_attr_size as usize {
255            let (head, tail) = self.buf.split_at_mut(total_attr_size as usize);
256            self.buf = unsafe {
257                core::mem::transmute(tail)
258            };
259            unsafe {
260                core::mem::transmute(Some(ByteAttrMut::from(head)))
261            }
262        } else {
263            let attr = unsafe {
264                Some(ByteAttrMut::from(core::mem::transmute::<&mut [u8], &mut [u8]>(self.buf)))
265            };
266            self.buf = &mut [];
267            attr
268        }
269    }
270}
271
272#[cfg(test)]
273mod tests {
274    use super::*;
275
276    const MSG: [u8; 28] = [
277        0x00, 0x01,                     // type: Binding Request
278        0x00, 0x08,                     // length: 8 (header does not count)
279        0x21, 0x12, 0xA4, 0x42,         // magic cookie
280        0x00, 0x00, 0x00, 0x00,
281        0x00, 0x00, 0x00, 0x00,
282        0x00, 0x00, 0x00, 0x01,         // transaction id
283        0x00, 0x03,                     // type: ChangeRequest
284        0x00, 0x04,                     // length: 4 (only value bytes count)
285        0x00, 0x00, 0x00, 0x40 | 0x20,  // change both ip and port
286    ];
287
288    #[test]
289    fn read() {
290        let msg = ByteMsg::from_arr(&MSG);
291
292        assert_eq!(&MSG[0..2], msg.typ().unwrap());
293        assert_eq!(&MSG[2..4], msg.len().unwrap());
294        assert_eq!(&MSG[4..8], msg.cookie().unwrap());
295        assert_eq!(&MSG[8..20], msg.tid().unwrap());
296        assert_eq!(&MSG[20..28], msg.attrs().unwrap());
297        assert_eq!(1, msg.attrs_iter().count());
298
299        let attr = msg.attrs_iter().next().unwrap();
300
301        assert_eq!(&MSG[20..22], attr.typ().unwrap());
302        assert_eq!(&MSG[22..24], attr.len().unwrap());
303        assert_eq!(&MSG[24..28], attr.val().unwrap());
304    }
305
306    #[test]
307    fn read_mut() {
308        let mut buf = MSG.clone();
309        let mut msg = ByteMsgMut::from_arr_mut(&mut buf);
310
311        assert_eq!(&MSG[0..2], msg.typ().unwrap());
312        assert_eq!(&MSG[2..4], msg.len().unwrap());
313        assert_eq!(&MSG[4..8], msg.cookie().unwrap());
314        assert_eq!(&MSG[8..20], msg.tid().unwrap());
315        assert_eq!(&MSG[20..28], msg.attrs().unwrap());
316        assert_eq!(1, msg.attrs_iter().count());
317
318        let mut attr = msg.attrs_iter().next().unwrap();
319
320        assert_eq!(&MSG[20..22], attr.typ().unwrap());
321        assert_eq!(&MSG[22..24], attr.len().unwrap());
322        assert_eq!(&MSG[24..28], attr.val().unwrap());
323    }
324
325    #[test]
326    fn write() {
327        let mut buf = [0u8; MSG.len()];
328        let mut msg = ByteMsgMut::from_arr_mut(&mut buf);
329
330        msg.typ().unwrap().copy_from(MSG.carved().unwrap());
331        // msg.len().unwrap().copy_from(MSG.carve(2..4).unwrap()); // length should be updated automatically
332        msg.cookie().unwrap().copy_from(MSG.carve(4).unwrap());
333        msg.tid().unwrap().copy_from(MSG.carve(8).unwrap());
334        msg.add_attr(MSG.carve(20).unwrap(), MSG.carve(22).unwrap(), MSG.get(24..28).unwrap());
335
336        assert_eq!(&MSG, &buf);
337
338        let mut msg = ByteMsgMut::from_arr_mut(&mut buf);
339        msg.len().unwrap().copy_from(&[0, 0]); // just to double-check len() works
340        assert_ne!(&MSG, &buf);
341
342        let mut msg = ByteMsgMut::from_arr_mut(&mut buf);
343        msg.len().unwrap().copy_from(MSG.carve(2).unwrap());
344        assert_eq!(&MSG, msg.as_bytes());
345    }
346
347    #[test]
348    fn write2() {
349        let mut buf = [0u8; MSG.len()];
350        let mut msg = ByteMsgMut::from_arr_mut(&mut buf);
351
352        msg.typ().unwrap().copy_from(MSG.carved().unwrap());
353        // msg.len().unwrap().copy_from(MSG.carve(2..4).unwrap()); // length should be updated automatically
354        msg.cookie().unwrap().copy_from(MSG.carve(4).unwrap());
355        msg.tid().unwrap().copy_from(MSG.carve(8).unwrap());
356        msg.add_attr2(|typ, len, val| {
357            typ.copy_from(MSG.carve(20)?);
358            len.copy_from(MSG.carve(22)?);
359            val.get_mut(0..4)?.copy_from_slice(MSG.get(24..28)?);
360            Some(4)
361        });
362
363        assert_eq!(&MSG, msg.as_bytes());
364    }
365}
stun_bytes/lib.rs

stun_bytes/
lib.rs
