donet_core/datagram/
datagram.rs

1/*
2    This file is part of Donet.
3
4    Copyright © 2024 Max Rodriguez <[email protected]>
5
6    Donet is free software; you can redistribute it and/or modify
7    it under the terms of the GNU Affero General Public License,
8    as published by the Free Software Foundation, either version 3
9    of the License, or (at your option) any later version.
10
11    Donet is distributed in the hope that it will be useful,
12    but WITHOUT ANY WARRANTY; without even the implied warranty of
13    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
14    GNU Affero General Public License for more details.
15
16    You should have received a copy of the GNU Affero General Public
17    License along with Donet. If not, see <https://www.gnu.org/licenses/>.
18*/
19
20//! Provides structure to write network packets (datagrams).
21
22use crate::datagram::byte_order as endianness;
23use crate::globals::*;
24use anyhow::Result;
25use thiserror::Error;
26
27/// Custom error type for [`Datagram`].
28#[derive(Debug, Error, PartialEq)]
29pub enum DatagramError {
30    #[error("datagram overflow; {0}")]
31    DatagramOverflow(&'static str),
32    #[error("impossible cast; {0}")]
33    ImpossibleCast(&'static str),
34}
35
36impl From<DatagramError> for std::io::Error {
37    fn from(value: DatagramError) -> Self {
38        std::io::Error::new(std::io::ErrorKind::Other, value.to_string())
39    }
40}
41
42/// Representation of a new network message (datagram) to be sent.
43#[derive(Debug, Clone)]
44pub struct Datagram {
45    buffer: Vec<u8>,
46    index: usize,
47    /// See [`Datagram::override_cap`].
48    cap: usize,
49}
50
51impl Default for Datagram {
52    fn default() -> Self {
53        Self {
54            buffer: vec![],
55            index: 0,
56            cap: usize::from(DgSizeTag::MAX),
57        }
58    }
59}
60
61/// Appends another datagram's raw bytes to this datagram.
62///
63/// Consumes the right-hand-side [`Datagram`].
64impl std::ops::Add for Datagram {
65    type Output = Result<Datagram, DatagramError>;
66
67    fn add(mut self, rhs: Self) -> Self::Output {
68        let dg_buffer: Vec<u8> = rhs.get_data();
69
70        self.add_data(dg_buffer)?;
71        Ok(self)
72    }
73}
74
75impl Datagram {
76    /// Checks if we can add `length` number of bytes to the datagram.
77    fn check_add_length(&mut self, length: usize) -> Result<(), DatagramError> {
78        let new_index: usize = self.index + length;
79
80        if new_index > self.cap {
81            return Err(DatagramError::DatagramOverflow(
82                "Tried to add data to the datagram past its maximum size!",
83            ));
84        }
85        Ok(())
86    }
87
88    /// Overrides the byte limit for this [`Datagram`].
89    ///
90    /// It should **always** be set to the MAX of the size
91    /// tag type (u16), but it can be overridden for special
92    /// uses, such as processing large buffers that combine
93    /// multiple TCP segments.
94    pub fn override_cap(&mut self, cap: usize) {
95        self.cap = cap
96    }
97
98    /// Adds an unsigned 8-bit integer to the datagram that is
99    /// guaranteed to be one of the values 0x00 (false) or 0x01 (true).
100    pub fn add_bool(&mut self, v: bool) -> Result<(), DatagramError> {
101        self.check_add_length(1)?;
102        if v {
103            self.add_u8(1)?;
104        } else {
105            self.add_u8(0)?;
106        }
107        Ok(())
108    }
109
110    /// Adds an unsigned 8-bit integer value to the datagram.
111    pub fn add_u8(&mut self, v: u8) -> Result<(), DatagramError> {
112        self.check_add_length(1)?;
113        self.buffer.push(v);
114        self.index += 1;
115        Ok(())
116    }
117
118    /// Adds an unsigned 16-bit integer value to the datagram.
119    pub fn add_u16(&mut self, mut v: u16) -> Result<(), DatagramError> {
120        self.check_add_length(2)?;
121
122        v = endianness::swap_le_16(v);
123
124        self.buffer.push((v & 0x00ff) as u8);
125        self.buffer.push(((v & 0xff00) >> 8) as u8);
126
127        self.index += 2;
128        Ok(())
129    }
130
131    /// Adds an unsigned 32-bit integer value to the datagram.
132    pub fn add_u32(&mut self, mut v: u32) -> Result<(), DatagramError> {
133        self.check_add_length(4)?;
134
135        v = endianness::swap_le_32(v);
136
137        self.buffer.push((v & 0x000000ff) as u8);
138        self.buffer.push(((v & 0x0000ff00) >> 8) as u8);
139        self.buffer.push(((v & 0x00ff0000) >> 16) as u8);
140        self.buffer.push(((v & 0xff000000) >> 24) as u8);
141
142        self.index += 4;
143        Ok(())
144    }
145
146    /// Adds an unsigned 64-bit integer value to the datagram.
147    pub fn add_u64(&mut self, mut v: u64) -> Result<(), DatagramError> {
148        self.check_add_length(8)?;
149
150        v = endianness::swap_le_64(v);
151
152        self.buffer.push((v & 0x00000000000000ff) as u8);
153        self.buffer.push(((v & 0x000000000000ff00) >> 8) as u8);
154        self.buffer.push(((v & 0x0000000000ff0000) >> 16) as u8);
155        self.buffer.push(((v & 0x00000000ff000000) >> 24) as u8);
156        self.buffer.push(((v & 0x000000ff00000000) >> 32) as u8);
157        self.buffer.push(((v & 0x0000ff0000000000) >> 40) as u8);
158        self.buffer.push(((v & 0x00ff000000000000) >> 48) as u8);
159        self.buffer.push(((v & 0xff00000000000000) >> 56) as u8);
160
161        self.index += 8;
162        Ok(())
163    }
164
165    // signed integer aliases. same bitwise operations.
166    #[inline(always)]
167    pub fn add_i8(&mut self, v: i8) -> Result<(), DatagramError> {
168        self.add_u8(v as u8)
169    }
170
171    #[inline(always)]
172    pub fn add_i16(&mut self, v: i16) -> Result<(), DatagramError> {
173        self.add_u16(v as u16)
174    }
175
176    #[inline(always)]
177    pub fn add_i32(&mut self, v: i32) -> Result<(), DatagramError> {
178        self.add_u32(v as u32)
179    }
180
181    #[inline(always)]
182    pub fn add_i64(&mut self, v: i64) -> Result<(), DatagramError> {
183        self.add_u64(v as u64)
184    }
185
186    /// 32-bit IEEE 754 floating point. same bitwise operations.
187    #[inline(always)]
188    pub fn add_f32(&mut self, v: f32) -> Result<(), DatagramError> {
189        self.add_u32(v as u32)
190    }
191
192    /// 64-bit IEEE 754 floating point. same bitwise operations.
193    #[inline(always)]
194    pub fn add_f64(&mut self, v: f64) -> Result<(), DatagramError> {
195        self.add_u64(v as u64)
196    }
197
198    /// Adds a Datagram / Field length tag to the end of the datagram.
199    #[inline(always)]
200    pub fn add_size(&mut self, v: DgSizeTag) -> Result<(), DatagramError> {
201        self.add_u16(v)
202    }
203
204    /// Adds a 64-bit channel ID to the end of the datagram.
205    #[inline(always)]
206    pub fn add_channel(&mut self, v: Channel) -> Result<(), DatagramError> {
207        self.add_u64(v)
208    }
209
210    /// Adds a 32-bit Distributed Object ID to the end of the datagram.
211    #[inline(always)]
212    pub fn add_doid(&mut self, v: DoId) -> Result<(), DatagramError> {
213        self.add_u32(v)
214    }
215
216    /// Adds a 32-bit zone ID to the end of the datagram.
217    #[inline(always)]
218    pub fn add_zone(&mut self, v: Zone) -> Result<(), DatagramError> {
219        self.add_u32(v)
220    }
221
222    /// Added for convenience, rather than adding the parent and the zone separately.
223    #[inline(always)]
224    pub fn add_location(&mut self, parent: DoId, zone: Zone) -> Result<(), DatagramError> {
225        self.add_u32(parent)?;
226        self.add_u32(zone)
227    }
228
229    /// Adds raw bytes to the datagram via an unsigned 8-bit integer vector.
230    ///
231    /// **NOTE**: not to be confused with [`Datagram::add_blob`], which
232    /// adds a dclass blob to the datagram.
233    ///
234    pub fn add_data(&mut self, mut bytes: Vec<u8>) -> Result<(), DatagramError> {
235        let size: usize = bytes.len();
236
237        self.check_add_length(size)?;
238        self.buffer.append(&mut bytes);
239
240        self.index += size;
241        Ok(())
242    }
243
244    /// Adds a dclass string value to the end of the datagram.
245    /// A 16-bit length tag prefix with the string's size in bytes is added.
246    pub fn add_string(&mut self, str: &str) -> Result<(), DatagramError> {
247        let size: usize = str.len();
248
249        if size > DG_SIZE_MAX.into() {
250            // The string is too big to be described with a 16-bit length tag.
251            return Err(DatagramError::DatagramOverflow(
252                "Given string will overflow datagram.",
253            ));
254        }
255        // Add string length to the datagram
256        self.add_u16(size.try_into().expect("String size should fit u16."))?;
257
258        // convert the string into a byte array, as a vector
259        let mut str_bytes: Vec<u8> = str.as_bytes().to_vec();
260
261        // make sure the byte array won't overflow the datagram
262        self.check_add_length(str_bytes.len())?;
263        self.buffer.append(&mut str_bytes);
264
265        self.index += size;
266        Ok(())
267    }
268
269    /// Adds a dclass blob value (binary data) to the end of the datagram.
270    /// A 16-bit length tag prefix with the blob's size in bytes is added.
271    pub fn add_blob(&mut self, mut bytes: Vec<u8>) -> Result<(), DatagramError> {
272        let size: usize = bytes.len();
273
274        // add blob size in bytes
275        self.add_size(match size.try_into() {
276            Ok(n) => n,
277            Err(_) => {
278                return Err(DatagramError::ImpossibleCast(
279                    "Given blob has a size that does not fit in dg size tag.",
280                ))
281            }
282        })?;
283
284        // manually check add length before appending byte array
285        self.check_add_length(size)?;
286        self.buffer.append(&mut bytes);
287
288        self.index += size;
289        Ok(())
290    }
291
292    /// Reserves an amount of bytes in the datagram buffer.
293    pub fn add_buffer(&mut self, size: usize) -> Result<usize, DatagramError> {
294        self.check_add_length(size)?;
295
296        // get start length (before push)
297        let start: usize = self.index;
298
299        for _ in 1..size {
300            self.buffer.push(0);
301        }
302        self.index += size;
303        Ok(start)
304    }
305
306    /// Appends a generic header for messages that are to be routed to
307    /// one or more role instances within the server cluster.
308    ///
309    /// Use this method to avoid repetitive code for every internal message.
310    ///
311    /// The header is formatted as shown below:
312    ///
313    /// (recipients: [`u8`], recipients: [`Vec<Channel>`],
314    /// sender: [`globals::Channel`], message_type: [`u16`])
315    ///
316    /// # Errors
317    ///
318    /// It is an error for the given `recipients` vector to have a size
319    /// larger than [`std::u8::MAX`]. Else, [`DatagramError::ImpossibleCast`]
320    /// will be returned.
321    ///
322    pub fn add_internal_header(
323        &mut self,
324        recipients: Vec<Channel>,
325        sender: Channel,
326        msg_type: MsgType,
327    ) -> Result<(), DatagramError> {
328        let n_recipients: usize = recipients.len();
329
330        if n_recipients > usize::from(u8::MAX) {
331            return Err(DatagramError::ImpossibleCast(
332                "Cannot convert recipient vec size to u8.",
333            ));
334        }
335        // Add recipient(s) count
336        self.add_u8(n_recipients.try_into().expect("Recipients exceeds u8 limit."))?;
337
338        for recipient in recipients {
339            // append each recipient in vector given
340            self.add_channel(recipient)?;
341        }
342        self.add_channel(sender)?;
343        self.add_u16(msg_type)
344    }
345
346    /// Appends a control header, which is very similar to a server header,
347    /// but it always has only one recipient, which is the control channel,
348    /// and does not require a sender (or 'from') channel to be provided.
349    ///
350    /// The sender field is not required as control messages are not
351    /// routed, meaning that the message director will ONLY receive this
352    /// message DIRECTLY from a cluster subscriber, so it can be speculated
353    /// that the sender is the participant on the other end of the connection.
354    pub fn add_control_header(&mut self, msg_type: MsgType) -> Result<(), DatagramError> {
355        self.add_u8(1)?;
356        self.add_channel(CONTROL_CHANNEL)?;
357        self.add_u16(msg_type)
358    }
359
360    /// Returns the size of this [`Datagram`].
361    pub fn size(&self) -> usize {
362        self.buffer.len()
363    }
364
365    /// Returns a reference to this [`Datagram`]'s byte buffer.
366    pub fn get_buffer(&self) -> &[u8] {
367        &self.buffer
368    }
369
370    /// Similar to [`Self::get_buffer`], but returns a copy of the buffer.
371    pub fn get_data(&self) -> Vec<u8> {
372        // we can't give out ownership of our vector,
373        // so a copy of the vector is made instead
374        let mut vec_copy: Vec<u8> = vec![];
375        for byte in &self.buffer {
376            // dereference the borrowed 'byte'
377            vec_copy.push(*byte);
378        }
379        vec_copy
380    }
381}
382
383#[cfg(test)]
384mod tests {
385    use super::*;
386    use crate::Protocol;
387
388    #[test]
389    fn add_boolean() {
390        let mut dg: Datagram = Datagram::default();
391        let mut res: Result<(), DatagramError> = dg.add_bool(true);
392
393        match res {
394            Ok(_) => {}
395            Err(err) => panic!("add_bool(true) error: {:?}", err),
396        }
397
398        res = dg.add_bool(false);
399
400        match res {
401            Ok(_) => {}
402            Err(err) => panic!("add_bool(false) error: {:?}", err),
403        }
404    }
405
406    #[test]
407    fn add_integers_and_types() {
408        // A bit repetitive, but we need coverage on all of these methods.
409        let mut dg: Datagram = Datagram::default();
410        let mut results: Vec<Result<(), DatagramError>> = vec![];
411
412        // Signed integers
413        results.push(dg.add_i8(i8::MAX));
414        results.push(dg.add_i16(i16::MAX));
415        results.push(dg.add_i32(i32::MAX));
416        results.push(dg.add_i64(i64::MAX));
417
418        // Unsigned integers
419        results.push(dg.add_u8(u8::MAX));
420        results.push(dg.add_u16(u16::MAX));
421        results.push(dg.add_u32(u32::MAX));
422        results.push(dg.add_u64(u64::MAX));
423
424        // 32-bit/64-bit floats
425        results.push(dg.add_f32(f32::MAX));
426        results.push(dg.add_f64(f64::MAX));
427
428        // Types (aliases)
429        results.push(dg.add_size(DG_SIZE_MAX));
430        results.push(dg.add_channel(CHANNEL_MAX));
431        results.push(dg.add_doid(DOID_MAX));
432        results.push(dg.add_zone(ZONE_MAX));
433        results.push(dg.add_location(DOID_MAX, ZONE_MAX));
434        results.push(dg.add_string("TEST")); // 16-bit length prefix + # of chars
435        results.push(dg.add_blob(vec![u8::MAX, u8::MAX])); // same prefix as above
436        results.push(dg.add_data(vec![u8::MAX, u8::MAX, u8::MAX, u8::MAX]));
437
438        for dg_res in &results {
439            assert!(dg_res.is_ok());
440        }
441
442        let dg_size: usize = dg.size();
443        let dg_buffer: Vec<u8> = dg.get_data();
444
445        assert_eq!(dg_buffer.len(), dg_size); // verify buffer length
446        assert_eq!(dg_size, 82); // total in bytes
447    }
448
449    #[test]
450    #[rustfmt::skip]
451    fn add_datagram() {
452        let mut dg: Datagram = Datagram::default();
453        let mut dg_2: Datagram = Datagram::default();
454
455        assert!(dg.add_channel(CHANNEL_MAX).is_ok());
456        assert!(dg_2.add_blob(vec![0, 125, u8::MAX]).is_ok());
457
458        let addition = dg.clone() + dg_2;
459
460        assert!(addition.is_ok());
461        dg = addition.unwrap();
462
463        let dg_size: usize = dg.size();
464        let dg_buffer: Vec<u8> = dg.get_data();
465
466        assert_eq!(dg_buffer.len(), dg_size);
467        assert_eq!(dg_buffer, vec![
468            u8::MAX, u8::MAX, u8::MAX, u8::MAX,
469            u8::MAX, u8::MAX, u8::MAX, u8::MAX,
470            3, 0, 0, 125, u8::MAX,
471        ]);
472    }
473
474    #[test]
475    #[rustfmt::skip]
476    fn message_headers() {
477        let mut dg: Datagram = Datagram::default();
478        let mut results: Vec<Result<(), DatagramError>> = vec![];
479
480        results.push(dg.add_internal_header(
481            vec![CHANNEL_MAX], // recipients
482            0, // sender
483            Protocol::MDAddChannel.into(), // msg type
484        ));
485
486        results.push(dg.add_control_header(Protocol::MDAddChannel.into()));
487
488        for dg_res in &results {
489            assert!(dg_res.is_ok());
490        }
491        let dg_size: usize = dg.size();
492        let dg_buffer: Vec<u8> = dg.get_data();
493
494        assert_eq!(dg_buffer.len(), dg_size);
495        assert_eq!(dg_buffer, vec![
496            1, u8::MAX, u8::MAX, u8::MAX, u8::MAX, // recipients
497            u8::MAX, u8::MAX, u8::MAX, u8::MAX,
498            0, 0, 0, 0, 0, 0, 0, 0, // sender
499            40, 35, // message type (9000; 0x2823, or 40, 35)
500            1, 1, 0, 0, 0, 0, 0, 0, 0, // recipients (control)
501            40, 35, // message type
502        ]);
503    }
504
505    #[test]
506    fn overflow_test() {
507        let mut dg: Datagram = Datagram::default();
508        let res_1: Result<usize, DatagramError> = dg.add_buffer(DG_SIZE_MAX.into());
509
510        assert!(!res_1.is_err(), "Could not append 2^16 bytes to datagram buffer.");
511        assert_eq!(res_1.unwrap(), 0, "add_buffer() didn't return start of reserve.");
512        assert_eq!(
513            dg.size() + 1,
514            usize::from(DG_SIZE_MAX),
515            "Datagram didn't add 2^16 bytes to the buffer."
516        );
517
518        let res_2: Result<(), DatagramError> = dg.add_u16(0);
519        assert!(
520            res_2.is_err(),
521            "Datagram overflow occurred, but did not throw an error."
522        );
523
524        assert_eq!(
525            res_2.unwrap_err(),
526            DatagramError::DatagramOverflow("Tried to add data to the datagram past its maximum size!"),
527            "Datagram overflow occurred, but failed to respond with DatagramOverflow err."
528        );
529    }
530}
donet_core/datagram/datagram.rs

donet_core/datagram/
datagram.rs
