veloren_network_protocol/
frame.rs

1use crate::types::{Bandwidth, Mid, Pid, Prio, Promises, Sid};
2use bytes::{Buf, BufMut, Bytes, BytesMut};
3
4// const FRAME_RESERVED_1: u8 = 0;
5const FRAME_HANDSHAKE: u8 = 1;
6const FRAME_INIT: u8 = 2;
7const FRAME_SHUTDOWN: u8 = 3;
8const FRAME_OPEN_STREAM: u8 = 4;
9const FRAME_CLOSE_STREAM: u8 = 5;
10const FRAME_DATA_HEADER: u8 = 6;
11const FRAME_DATA: u8 = 7;
12const FRAME_RAW: u8 = 8;
13//const FRAME_RESERVED_2: u8 = 10;
14//const FRAME_RESERVED_3: u8 = 13;
15
16/// Used for Communication between Channel <----(TCP/UDP)----> Channel
17#[derive(Debug, PartialEq, Eq, Clone)]
18pub enum InitFrame {
19    Handshake {
20        magic_number: [u8; 7],
21        version: [u32; 3],
22    },
23    Init {
24        pid: Pid,
25        secret: u128,
26    },
27    /// WARNING: sending RAW is only for debug purposes and will drop the
28    /// connection
29    Raw(Vec<u8>),
30}
31
32/// Used for OUT TCP Communication between Channel --(TCP)--> Channel
33#[derive(Debug, PartialEq, Eq, Clone)]
34pub enum OTFrame {
35    Shutdown, /* Shutdown this channel gracefully, if all channels are shutdown (gracefully),
36               * Participant is deleted */
37    OpenStream {
38        sid: Sid,
39        prio: Prio,
40        promises: Promises,
41        guaranteed_bandwidth: Bandwidth,
42    },
43    CloseStream {
44        sid: Sid,
45    },
46    DataHeader {
47        mid: Mid,
48        sid: Sid,
49        length: u64,
50    },
51    Data {
52        mid: Mid,
53        data: Bytes,
54    },
55}
56
57/// Used for IN TCP Communication between Channel <--(TCP)-- Channel
58#[derive(Debug, PartialEq, Eq, Clone)]
59pub enum ITFrame {
60    Shutdown, /* Shutdown this channel gracefully, if all channels are shutdown (gracefully),
61               * Participant is deleted */
62    OpenStream {
63        sid: Sid,
64        prio: Prio,
65        promises: Promises,
66        guaranteed_bandwidth: Bandwidth,
67    },
68    CloseStream {
69        sid: Sid,
70    },
71    DataHeader {
72        mid: Mid,
73        sid: Sid,
74        length: u64,
75    },
76    Data {
77        mid: Mid,
78        data: BytesMut,
79    },
80}
81
82impl InitFrame {
83    // Size WITHOUT the 1rst indicating byte
84    pub(crate) const HANDSHAKE_CNS: usize = 19;
85    pub(crate) const INIT_CNS: usize = 32;
86    /// const part of the RAW frame, actual size is variable
87    pub(crate) const RAW_CNS: usize = 2;
88
89    //provide an appropriate buffer size. > 1500
90    pub(crate) fn write_bytes(self, bytes: &mut BytesMut) {
91        match self {
92            InitFrame::Handshake {
93                magic_number,
94                version,
95            } => {
96                bytes.put_u8(FRAME_HANDSHAKE);
97                bytes.put_slice(&magic_number);
98                bytes.put_u32_le(version[0]);
99                bytes.put_u32_le(version[1]);
100                bytes.put_u32_le(version[2]);
101            },
102            InitFrame::Init { pid, secret } => {
103                bytes.put_u8(FRAME_INIT);
104                pid.to_bytes(bytes);
105                bytes.put_u128_le(secret);
106            },
107            InitFrame::Raw(data) => {
108                bytes.put_u8(FRAME_RAW);
109                bytes.put_u16_le(data.len() as u16);
110                bytes.put_slice(&data);
111            },
112        }
113    }
114
115    pub(crate) fn read_frame(bytes: &mut BytesMut) -> Option<Self> {
116        let frame_no = match bytes.first() {
117            Some(&f) => f,
118            None => return None,
119        };
120        let frame = match frame_no {
121            FRAME_HANDSHAKE => {
122                if bytes.len() < Self::HANDSHAKE_CNS + 1 {
123                    return None;
124                }
125                bytes.advance(1);
126                let mut magic_number_bytes = bytes.copy_to_bytes(7);
127                let mut magic_number = [0u8; 7];
128                magic_number_bytes.copy_to_slice(&mut magic_number);
129                InitFrame::Handshake {
130                    magic_number,
131                    version: [bytes.get_u32_le(), bytes.get_u32_le(), bytes.get_u32_le()],
132                }
133            },
134            FRAME_INIT => {
135                if bytes.len() < Self::INIT_CNS + 1 {
136                    return None;
137                }
138                bytes.advance(1);
139                InitFrame::Init {
140                    pid: Pid::from_bytes(bytes),
141                    secret: bytes.get_u128_le(),
142                }
143            },
144            FRAME_RAW => {
145                if bytes.len() < Self::RAW_CNS + 1 {
146                    return None;
147                }
148                bytes.advance(1);
149                let length = bytes.get_u16_le() as usize;
150                // lower length is allowed
151                let max_length = length.min(bytes.len());
152                let mut data = vec![0; max_length];
153                data.copy_from_slice(&bytes[..max_length]);
154                InitFrame::Raw(data)
155            },
156            _ => InitFrame::Raw(bytes.to_vec()),
157        };
158        Some(frame)
159    }
160}
161
162pub(crate) const TCP_CLOSE_STREAM_CNS: usize = 8;
163/// const part of the DATA frame, actual size is variable
164pub(crate) const TCP_DATA_CNS: usize = 10;
165pub(crate) const TCP_DATA_HEADER_CNS: usize = 24;
166pub(crate) const TCP_OPEN_STREAM_CNS: usize = 18;
167// Size WITHOUT the 1rst indicating byte
168pub(crate) const TCP_SHUTDOWN_CNS: usize = 0;
169
170impl OTFrame {
171    pub fn write_bytes(self, bytes: &mut BytesMut) {
172        match self {
173            Self::Shutdown => {
174                bytes.put_u8(FRAME_SHUTDOWN);
175            },
176            Self::OpenStream {
177                sid,
178                prio,
179                promises,
180                guaranteed_bandwidth,
181            } => {
182                bytes.put_u8(FRAME_OPEN_STREAM);
183                sid.to_bytes(bytes);
184                bytes.put_u8(prio);
185                bytes.put_u8(promises.to_le_bytes()[0]);
186                bytes.put_u64_le(guaranteed_bandwidth);
187            },
188            Self::CloseStream { sid } => {
189                bytes.put_u8(FRAME_CLOSE_STREAM);
190                sid.to_bytes(bytes);
191            },
192            Self::DataHeader { mid, sid, length } => {
193                bytes.put_u8(FRAME_DATA_HEADER);
194                bytes.put_u64_le(mid);
195                sid.to_bytes(bytes);
196                bytes.put_u64_le(length);
197            },
198            Self::Data { mid, data } => {
199                bytes.put_u8(FRAME_DATA);
200                bytes.put_u64_le(mid);
201                bytes.put_u16_le(data.len() as u16);
202                bytes.put_slice(&data);
203            },
204        }
205    }
206}
207
208impl ITFrame {
209    /// Err => cannot recover
210    /// Ok(None) => waiting for more data
211    pub(crate) fn read_frame(bytes: &mut BytesMut) -> Result<Option<Self>, ()> {
212        let frame_no = match bytes.first() {
213            Some(&f) => f,
214            None => return Ok(None),
215        };
216        let size = match frame_no {
217            FRAME_SHUTDOWN => TCP_SHUTDOWN_CNS,
218            FRAME_OPEN_STREAM => TCP_OPEN_STREAM_CNS,
219            FRAME_CLOSE_STREAM => TCP_CLOSE_STREAM_CNS,
220            FRAME_DATA_HEADER => TCP_DATA_HEADER_CNS,
221            FRAME_DATA => {
222                if bytes.len() < 9 + 1 + 1 {
223                    return Ok(None);
224                }
225                u16::from_le_bytes([bytes[8 + 1], bytes[9 + 1]]) as usize + TCP_DATA_CNS
226            },
227            _ => return Err(()),
228        };
229
230        if bytes.len() < size + 1 {
231            return Ok(None);
232        }
233
234        let frame = match frame_no {
235            FRAME_SHUTDOWN => {
236                let _ = bytes.split_to(size + 1);
237                Self::Shutdown
238            },
239            FRAME_OPEN_STREAM => {
240                let mut bytes = bytes.split_to(size + 1);
241                bytes.advance(1);
242                Self::OpenStream {
243                    sid: Sid::from_bytes(&mut bytes),
244                    prio: bytes.get_u8(),
245                    promises: Promises::from_bits_truncate(bytes.get_u8()),
246                    guaranteed_bandwidth: bytes.get_u64_le(),
247                }
248            },
249            FRAME_CLOSE_STREAM => {
250                let mut bytes = bytes.split_to(size + 1);
251                bytes.advance(1);
252                Self::CloseStream {
253                    sid: Sid::from_bytes(&mut bytes),
254                }
255            },
256            FRAME_DATA_HEADER => {
257                let mut bytes = bytes.split_to(size + 1);
258                bytes.advance(1);
259                Self::DataHeader {
260                    mid: bytes.get_u64_le(),
261                    sid: Sid::from_bytes(&mut bytes),
262                    length: bytes.get_u64_le(),
263                }
264            },
265            FRAME_DATA => {
266                bytes.advance(1);
267                let mid = bytes.get_u64_le();
268                let length = bytes.get_u16_le();
269                debug_assert_eq!(length as usize, size - TCP_DATA_CNS);
270                let data = bytes.split_to(length as usize);
271                Self::Data { mid, data }
272            },
273            _ => unreachable!("Frame::to_frame should be handled before!"),
274        };
275        Ok(Some(frame))
276    }
277}
278
279#[expect(unused_variables)]
280impl PartialEq<ITFrame> for OTFrame {
281    fn eq(&self, other: &ITFrame) -> bool {
282        match self {
283            Self::Shutdown => matches!(other, ITFrame::Shutdown),
284            Self::OpenStream {
285                sid,
286                prio,
287                promises,
288                guaranteed_bandwidth,
289            } => matches!(other, ITFrame::OpenStream {
290                sid,
291                prio,
292                promises,
293                guaranteed_bandwidth,
294            }),
295            Self::CloseStream { sid } => matches!(other, ITFrame::CloseStream { sid }),
296            Self::DataHeader { mid, sid, length } => {
297                matches!(other, ITFrame::DataHeader { mid, sid, length })
298            },
299            Self::Data { mid, data } => matches!(other, ITFrame::Data { mid, data }),
300        }
301    }
302}
303
304#[cfg(test)]
305mod tests {
306    use super::*;
307    use crate::types::{VELOREN_MAGIC_NUMBER, VELOREN_NETWORK_VERSION};
308
309    fn get_initframes() -> Vec<InitFrame> {
310        vec![
311            InitFrame::Handshake {
312                magic_number: VELOREN_MAGIC_NUMBER,
313                version: VELOREN_NETWORK_VERSION,
314            },
315            InitFrame::Init {
316                pid: Pid::fake(0),
317                secret: 0u128,
318            },
319            InitFrame::Raw(vec![1, 2, 3]),
320        ]
321    }
322
323    fn get_otframes() -> Vec<OTFrame> {
324        vec![
325            OTFrame::OpenStream {
326                sid: Sid::new(1337),
327                prio: 14,
328                promises: Promises::GUARANTEED_DELIVERY,
329                guaranteed_bandwidth: 1_000_000,
330            },
331            OTFrame::DataHeader {
332                sid: Sid::new(1337),
333                mid: 0,
334                length: 36,
335            },
336            OTFrame::Data {
337                mid: 0,
338                data: Bytes::from(&[77u8; 20][..]),
339            },
340            OTFrame::Data {
341                mid: 0,
342                data: Bytes::from(&[42u8; 16][..]),
343            },
344            OTFrame::CloseStream {
345                sid: Sid::new(1337),
346            },
347            OTFrame::Shutdown,
348        ]
349    }
350
351    #[test]
352    fn initframe_individual() {
353        let dupl = |frame: InitFrame| {
354            let mut buffer = BytesMut::with_capacity(1500);
355            InitFrame::write_bytes(frame, &mut buffer);
356            InitFrame::read_frame(&mut buffer)
357        };
358
359        for frame in get_initframes() {
360            println!("initframe: {:?}", &frame);
361            assert_eq!(Some(frame.clone()), dupl(frame));
362        }
363    }
364
365    #[test]
366    fn initframe_multiple() {
367        let mut buffer = BytesMut::with_capacity(3000);
368
369        let mut frames = get_initframes();
370        // to string
371        for f in &frames {
372            InitFrame::write_bytes(f.clone(), &mut buffer);
373        }
374
375        // from string
376        let mut framesd = frames
377            .iter()
378            .map(|&_| InitFrame::read_frame(&mut buffer))
379            .collect::<Vec<_>>();
380
381        // compare
382        for (f, fd) in frames.drain(..).zip(framesd.drain(..)) {
383            println!("initframe: {:?}", &f);
384            assert_eq!(Some(f), fd);
385        }
386    }
387
388    #[test]
389    fn frame_individual() {
390        let dupl = |frame: OTFrame| {
391            let mut buffer = BytesMut::with_capacity(1500);
392            OTFrame::write_bytes(frame, &mut buffer);
393            ITFrame::read_frame(&mut buffer)
394        };
395
396        for frame in get_otframes() {
397            println!("frame: {:?}", &frame);
398            assert_eq!(frame.clone(), dupl(frame).expect("ERR").expect("NONE"));
399        }
400    }
401
402    #[test]
403    fn frame_multiple() {
404        let mut buffer = BytesMut::with_capacity(3000);
405
406        let mut frames = get_otframes();
407        // to string
408        for f in &frames {
409            OTFrame::write_bytes(f.clone(), &mut buffer);
410        }
411
412        // from string
413        let mut framesd = frames
414            .iter()
415            .map(|&_| ITFrame::read_frame(&mut buffer))
416            .collect::<Vec<_>>();
417
418        // compare
419        for (f, fd) in frames.drain(..).zip(framesd.drain(..)) {
420            println!("frame: {:?}", &f);
421            assert_eq!(f, fd.expect("ERR").expect("NONE"));
422        }
423    }
424
425    #[test]
426    fn frame_exact_size() {
427        const SIZE: usize = TCP_CLOSE_STREAM_CNS+1/*first byte*/;
428        let mut buffer = BytesMut::with_capacity(SIZE);
429
430        let frame1 = OTFrame::CloseStream { sid: Sid::new(2) };
431        OTFrame::write_bytes(frame1.clone(), &mut buffer);
432        assert_eq!(buffer.len(), SIZE);
433        let mut deque = buffer.iter().copied().collect();
434        let frame2 = ITFrame::read_frame(&mut deque);
435        assert_eq!(frame1, frame2.expect("ERR").expect("NONE"));
436    }
437
438    #[test]
439    fn initframe_too_short_buffer() {
440        let mut buffer = BytesMut::with_capacity(10);
441
442        let frame1 = InitFrame::Handshake {
443            magic_number: VELOREN_MAGIC_NUMBER,
444            version: VELOREN_NETWORK_VERSION,
445        };
446        InitFrame::write_bytes(frame1, &mut buffer);
447    }
448
449    #[test]
450    fn initframe_too_less_data() {
451        let mut buffer = BytesMut::with_capacity(20);
452
453        let frame1 = InitFrame::Handshake {
454            magic_number: VELOREN_MAGIC_NUMBER,
455            version: VELOREN_NETWORK_VERSION,
456        };
457        InitFrame::write_bytes(frame1, &mut buffer);
458        buffer.truncate(6); // simulate partial retrieve
459        let frame1d = InitFrame::read_frame(&mut buffer);
460        assert_eq!(frame1d, None);
461    }
462
463    #[test]
464    fn initframe_rubbish() {
465        let mut buffer = BytesMut::from(&b"dtrgwcser"[..]);
466        assert_eq!(
467            InitFrame::read_frame(&mut buffer),
468            Some(InitFrame::Raw(b"dtrgwcser".to_vec()))
469        );
470    }
471
472    #[test]
473    fn initframe_attack_too_much_length() {
474        let mut buffer = BytesMut::with_capacity(50);
475
476        let frame1 = InitFrame::Raw(b"foobar".to_vec());
477        InitFrame::write_bytes(frame1.clone(), &mut buffer);
478        buffer[1] = 255;
479        let framed = InitFrame::read_frame(&mut buffer);
480        assert_eq!(framed, Some(frame1));
481    }
482
483    #[test]
484    fn initframe_attack_too_low_length() {
485        let mut buffer = BytesMut::with_capacity(50);
486
487        let frame1 = InitFrame::Raw(b"foobar".to_vec());
488        InitFrame::write_bytes(frame1, &mut buffer);
489        buffer[1] = 3;
490        let framed = InitFrame::read_frame(&mut buffer);
491        // we accept a different frame here, as it's RAW and debug only!
492        assert_eq!(framed, Some(InitFrame::Raw(b"foo".to_vec())));
493    }
494
495    #[test]
496    fn frame_too_short_buffer() {
497        let mut buffer = BytesMut::with_capacity(10);
498
499        let frame1 = OTFrame::OpenStream {
500            sid: Sid::new(88),
501            promises: Promises::ENCRYPTED,
502            prio: 88,
503            guaranteed_bandwidth: 1_000_000,
504        };
505        OTFrame::write_bytes(frame1, &mut buffer);
506    }
507
508    #[test]
509    fn frame_too_less_data() {
510        let mut buffer = BytesMut::with_capacity(20);
511
512        let frame1 = OTFrame::OpenStream {
513            sid: Sid::new(88),
514            promises: Promises::ENCRYPTED,
515            prio: 88,
516            guaranteed_bandwidth: 1_000_000,
517        };
518        OTFrame::write_bytes(frame1, &mut buffer);
519        buffer.truncate(6); // simulate partial retrieve
520        let frame1d = ITFrame::read_frame(&mut buffer);
521        assert_eq!(frame1d, Ok(None));
522    }
523
524    #[test]
525    fn frame_rubish() {
526        let mut buffer = BytesMut::from(&b"dtrgwcser"[..]);
527        assert_eq!(ITFrame::read_frame(&mut buffer), Err(()));
528    }
529
530    #[test]
531    fn frame_attack_too_much_length() {
532        let mut buffer = BytesMut::with_capacity(50);
533
534        let frame1 = OTFrame::Data {
535            mid: 7u64,
536            data: Bytes::from(&b"foobar"[..]),
537        };
538
539        OTFrame::write_bytes(frame1, &mut buffer);
540        buffer[9] = 255;
541        let framed = ITFrame::read_frame(&mut buffer);
542        assert_eq!(framed, Ok(None));
543    }
544
545    #[test]
546    fn frame_attack_too_low_length() {
547        let mut buffer = BytesMut::with_capacity(50);
548
549        let frame1 = OTFrame::Data {
550            mid: 7u64,
551            data: Bytes::from(&b"foobar"[..]),
552        };
553
554        OTFrame::write_bytes(frame1, &mut buffer);
555        buffer[9] = 3;
556        let framed = ITFrame::read_frame(&mut buffer);
557        assert_eq!(
558            framed,
559            Ok(Some(ITFrame::Data {
560                mid: 7u64,
561                data: BytesMut::from(&b"foo"[..]),
562            }))
563        );
564        //next = Invalid => Empty
565        let framed = ITFrame::read_frame(&mut buffer);
566        assert_eq!(framed, Err(()));
567    }
568}
veloren_network_protocol/frame.rs

veloren_network_protocol/
frame.rs
