feat(rust): improve ockam udp transport implementation robustness

implementations/rust/ockam/ockam_transport_udp/src/error.rs

@@ -0,0 +1,112 @@
+#![allow(missing_docs)]
+
+use crate::messages::RoutingNumber;
+use ockam_core::errcode::{Kind, Origin};
+use ockam_core::Error;
+
+/// UDP Transport error type
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub enum UdpTransportError {
+    InvalidTotalNumber(RoutingNumber),
+    TotalNumberMismatch {
+        routing_number: RoutingNumber,
+        original_total_number: u16,
+        actual_total_number: u16,
+    },
+    OutOfBounds {
+        routing_number: RoutingNumber,
+        total_number: u16,
+        offset: u16,
+    },
+    DuplicateMessage {
+        routing_number: RoutingNumber,
+        offset: u16,
+    },
+    DuplicateLastMessage(RoutingNumber),
+    UnexpectedLength {
+        routing_number: RoutingNumber,
+        expected_length: usize,
+        actual_length: usize,
+    },
+    EmptyPayload(RoutingNumber),
+    MessageLengthExceeded {
+        routing_number: RoutingNumber,
+        data_offset_end: usize,
+    },
+}
+
+impl ockam_core::compat::error::Error for UdpTransportError {}
+impl core::fmt::Display for UdpTransportError {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        match self {
+            Self::InvalidTotalNumber(r) => write!(
+                f,
+                "Received invalid message with total number 0 for Routing number: {r}",
+            ),
+            Self::TotalNumberMismatch {
+                routing_number,
+                original_total_number,
+                actual_total_number,
+            } => {
+                write!(f, "Received invalid message for Routing number: {routing_number}. Original total number: {original_total_number}, received total number: {actual_total_number}")
+            }
+            Self::OutOfBounds {
+                routing_number,
+                total_number,
+                offset,
+            } => {
+                write!(
+                    f,
+                    "Received invalid message for Routing number: {routing_number}. Total: {total_number}, Offset: {offset}",
+                )
+            }
+            Self::DuplicateMessage {
+                routing_number,
+                offset,
+            } => {
+                write!(
+                    f,
+                    "Received duplicate message for Routing number: {routing_number}. Offset: {offset}",
+                )
+            }
+            Self::DuplicateLastMessage(routing_number) => {
+                write!(
+                    f,
+                    "Received duplicate last message for Routing number: {routing_number}",
+                )
+            }
+            Self::UnexpectedLength {
+                routing_number,
+                expected_length,
+                actual_length,
+            } => {
+                write!(
+                    f,
+                    "Received invalid message for Routing number: {routing_number}. Expected len: {expected_length} Actual len: {actual_length}",
+                )
+            }
+            Self::EmptyPayload(routing_number) => {
+                write!(
+                    f,
+                    "Received invalid message with empty payload for Routing number: {routing_number}",
+                )
+            }
+            Self::MessageLengthExceeded {
+                routing_number,
+                data_offset_end,
+            } => {
+                write!(
+                    f,
+                    "Message exceeded maximum limit. Routing number: {routing_number}, Data offset end: {data_offset_end}",
+                )
+            }
+        }
+    }
+}
+
+impl From<UdpTransportError> for Error {
+    #[track_caller]
+    fn from(err: UdpTransportError) -> Error {
+        Error::new(Origin::Transport, Kind::Io, err)
+    }
+}

implementations/rust/ockam/ockam_transport_udp/src/lib.rs

@@ -17,12 +17,14 @@ extern crate core;
 #[cfg(feature = "alloc")]
 extern crate alloc;
 
+mod error;
 mod messages;
 mod options;
 mod puncture;
 mod transport;
 mod workers;
 
+pub use error::*;
 pub use options::UdpBindOptions;
 pub use puncture::*;
 pub use transport::{UdpBind, UdpBindArguments, UdpTransport, UdpTransportExtension};

implementations/rust/ockam/ockam_transport_udp/src/messages/routing_number.rs

@@ -2,6 +2,7 @@ use core::fmt::{Debug, Formatter};
 use minicbor::{CborLen, Decode, Encode};
 use rand::random;
 use std::cmp::Ordering;
+use std::ops::{Add, AddAssign, Sub};
 
 /// Number of the [`UdpRoutingMessage`]. Each [`UdpRoutingMessage`] is assigned a value, which
 /// helps the receiver to assemble the message. Start with a random value and uses overflowing
@@ -16,13 +17,19 @@ impl core::fmt::Display for RoutingNumber {
     }
 }
 
-impl RoutingNumber {
-    pub fn new() -> Self {
+impl Default for RoutingNumber {
+    fn default() -> Self {
         Self(random())
     }
+}
+
+impl RoutingNumber {
+    pub fn new(value: u16) -> Self {
+        Self(value)
+    }
 
     pub fn increment(&mut self) {
-        self.0 = self.0.overflowing_add(1).0;
+        *self += 1;
     }
 }
 
@@ -57,18 +64,42 @@ impl Ord for RoutingNumber {
     }
 }
 
+impl AddAssign<u16> for RoutingNumber {
+    fn add_assign(&mut self, rhs: u16) {
+        self.0 = self.0.wrapping_add(rhs);
+    }
+}
+
+impl Add<u16> for RoutingNumber {
+    type Output = u16;
+
+    fn add(self, rhs: u16) -> Self::Output {
+        let mut s = self;
+        s += rhs;
+        s.0
+    }
+}
+
+impl Sub<Self> for RoutingNumber {
+    type Output = u16;
+
+    fn sub(self, rhs: Self) -> Self::Output {
+        self.0.wrapping_sub(rhs.0)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate::messages::RoutingNumber;
 
     #[test]
-    fn new_is_random() {
-        assert_ne!(RoutingNumber::new(), RoutingNumber::new());
+    fn default_is_random() {
+        assert_ne!(RoutingNumber::default(), RoutingNumber::default());
     }
 
     #[test]
     fn increment() {
-        let mut number = RoutingNumber::new();
+        let mut number = RoutingNumber::default();
         let n1 = number.0;
         number.increment();
         assert_eq!(number.0, n1 + 1);

implementations/rust/ockam/ockam_transport_udp/src/messages/transport_message.rs

@@ -29,8 +29,8 @@ pub struct Version(#[n(0)] pub u8);
 pub struct UdpTransportMessage<'a> {
     #[n(0)] pub version: Version,
     #[n(1)] pub routing_number: RoutingNumber,
-    #[n(2)] pub offset: u32,
-    #[n(3)] pub is_last: bool,
+    #[n(2)] pub offset: u16,
+    #[n(3)] pub total: u16,
     #[b(4)] pub payload: CowBytes<'a>,
 }
 
@@ -39,15 +39,15 @@ impl<'a> UdpTransportMessage<'a> {
     pub fn new(
         version: Version,
         routing_number: RoutingNumber,
-        offset: u32,
-        is_last: bool,
+        offset: u16,
+        total: u16,
         payload: impl Into<CowBytes<'a>>,
     ) -> Self {
         Self {
             version,
             routing_number,
             offset,
-            is_last,
+            total,
             payload: payload.into(),
         }
     }
@@ -64,12 +64,12 @@ mod tests {
         let msg = UdpTransportMessage::new(
             Version(u8::MAX),
             RoutingNumber(u16::MAX),
-            u32::MAX,
-            true,
+            u16::MAX,
+            u16::MAX,
             vec![0u8; MAX_PAYLOAD_SIZE],
         );
 
-        let len = minicbor::to_vec(msg).unwrap().len();
+        let len = ockam_core::cbor_encode_preallocate(msg).unwrap().len();
 
         assert!(len <= MAX_ON_THE_WIRE_SIZE);
     }
@@ -79,12 +79,12 @@ mod tests {
         let msg = UdpTransportMessage::new(
             Version(u8::MAX),
             RoutingNumber(u16::MAX),
-            u32::MAX,
-            true,
+            u16::MAX,
+            u16::MAX,
             vec![0u8; MAX_PAYLOAD_SIZE],
         );
 
-        let len = minicbor::to_vec(msg).unwrap().len();
+        let len = ockam_core::cbor_encode_preallocate(msg).unwrap().len();
 
         assert_eq!(len, MAX_ON_THE_WIRE_SIZE);
     }

implementations/rust/ockam/ockam_transport_udp/src/workers/mod.rs

@@ -7,3 +7,5 @@ pub(crate) use addresses::*;
 pub(crate) use receiver::*;
 pub(crate) use sender::*;
 pub(crate) use socket_split::*;
+
+mod pending_messages;