Merge pull request #33 from openwallet-foundation-labs/unify-generic-…

…msgs remove ROUTE_MSG as a seperate encoding
openwallet-foundation-labs · Jul 22, 2024 · a1ddc4f · a1ddc4f
2 parents e350852 + 9cc2a4c
commit a1ddc4f
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 29 deletions.
diff --git a/fuzz/fuzz_targets/payload_encode_decode.rs b/fuzz/fuzz_targets/payload_encode_decode.rs
@@ -4,23 +4,16 @@ use libfuzzer_sys::fuzz_target;
 
 fuzz_target!(|data: tsp::cesr::fuzzing::Wrapper| {
     let mut buf = Vec::new();
-    tsp::cesr::encode_payload(&data.0, None, &mut buf).unwrap();
+    match tsp::cesr::encode_payload(&data.0, None, &mut buf) {
+        Ok(()) => {
+            let result = tsp::cesr::decode_payload(&buf).unwrap();
 
-    match tsp::cesr::decode_payload(&buf) {
-        Ok(result) => assert_eq!(data, result.payload),
-        Err(e) => {
-            // most of these errors should be unreachable if encoding succeeded
-            match e {
-                tsp::cesr::error::DecodeError::UnexpectedData => todo!(),
-                tsp::cesr::error::DecodeError::UnexpectedMsgType => todo!(),
-                tsp::cesr::error::DecodeError::TrailingGarbage => todo!(),
-                tsp::cesr::error::DecodeError::SignatureError => todo!(),
-                tsp::cesr::error::DecodeError::VidError => todo!(),
-                tsp::cesr::error::DecodeError::VersionMismatch => todo!(),
-                tsp::cesr::error::DecodeError::InvalidCryptoType => todo!(),
-                tsp::cesr::error::DecodeError::InvalidSignatureType => todo!(),
-                tsp::cesr::error::DecodeError::MissingHops => (),
-            }
+            assert_eq!(data, result.payload);
         }
+        Err(tsp::cesr::error::EncodeError::MissingHops) => match &data.0 {
+            tsp::cesr::Payload::RoutedMessage(route, _) => assert!(route.is_empty()),
+            _ => todo!(),
+        },
+        _ => todo!(),
     }
 });
diff --git a/tsp/src/cesr/error.rs b/tsp/src/cesr/error.rs
@@ -2,6 +2,7 @@
 #[derive(Clone, Copy, Debug)]
 pub enum EncodeError {
     PayloadTooLarge,
+    MissingHops,
 }
 
 /// An error type to indicate something went wrong with decoding
@@ -13,7 +14,6 @@ pub enum DecodeError {
     SignatureError,
     VidError,
     VersionMismatch,
-    MissingHops,
     InvalidCryptoType,
     InvalidSignatureType,
 }

diff --git a/tsp/src/cesr/packet.rs b/tsp/src/cesr/packet.rs
@@ -20,7 +20,6 @@ const TSP_PAYLOAD: u16 = (b'Z' - b'A') as u16;
 mod msgtype {
     pub(super) const GEN_MSG: [u8; 2] = [0, 0];
     pub(super) const NEST_MSG: [u8; 2] = [0, 1];
-    pub(super) const ROUTE_MSG: [u8; 2] = [0, 2];
     pub(super) const NEW_REL: [u8; 2] = [1, 0];
     pub(super) const NEW_REL_REPLY: [u8; 2] = [1, 1];
     pub(super) const NEW_NEST_REL: [u8; 2] = [1, 2];
@@ -217,7 +216,10 @@ pub fn encode_payload(
             checked_encode_variable_data(TSP_PLAINTEXT, data.as_ref(), output)?;
         }
         Payload::RoutedMessage(hops, data) => {
-            encode_fixed_data(TSP_TYPECODE, &msgtype::ROUTE_MSG, output);
+            encode_fixed_data(TSP_TYPECODE, &msgtype::GEN_MSG, output);
+            if hops.is_empty() {
+                return Err(EncodeError::MissingHops);
+            }
             encode_hops(hops, output)?;
             checked_encode_variable_data(TSP_PLAINTEXT, data.as_ref(), output)?;
         }
@@ -317,7 +319,13 @@ pub fn decode_payload(mut stream: &[u8]) -> Result<DecodedPayload, DecodeError>
         .ok_or(DecodeError::UnexpectedData)?
     {
         msgtype::GEN_MSG => {
-            decode_variable_data(TSP_PLAINTEXT, &mut stream).map(Payload::GenericMessage)
+            let hop_list = decode_hops(&mut stream)?;
+            if hop_list.is_empty() {
+                decode_variable_data(TSP_PLAINTEXT, &mut stream).map(Payload::GenericMessage)
+            } else {
+                decode_variable_data(TSP_PLAINTEXT, &mut stream)
+                    .map(|msg| Payload::RoutedMessage(hop_list, msg))
+            }
         }
         msgtype::NEW_REL => {
             let hop_list = decode_hops(&mut stream)?;
@@ -330,15 +338,6 @@ pub fn decode_payload(mut stream: &[u8]) -> Result<DecodedPayload, DecodeError>
         msgtype::NEST_MSG => {
             decode_variable_data(TSP_PLAINTEXT, &mut stream).map(Payload::NestedMessage)
         }
-        msgtype::ROUTE_MSG => {
-            let hop_list = decode_hops(&mut stream)?;
-            if hop_list.is_empty() {
-                return Err(DecodeError::MissingHops);
-            }
-
-            decode_variable_data(TSP_PLAINTEXT, &mut stream)
-                .map(|msg| Payload::RoutedMessage(hop_list, msg))
-        }
         msgtype::NEW_REL_REPLY => decode_fixed_data(TSP_SHA256, &mut stream)
             .map(|reply| Payload::DirectRelationAffirm { reply }),
         msgtype::NEW_NEST_REL => {