Add header support to request & response (#200)

CHANGELOG.md

@@ -7,6 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- `Request` and `Response` types from `distant-net` now support an optional
+  `Header` to send miscellaneous information
+
 ### Changed
 
 - `Change` structure now provides a single `path` instead of `paths` with the

Cargo.lock

@@ -935,6 +935,7 @@ dependencies = [
  "p256",
  "paste",
  "rand",
+ "rmp",
  "rmp-serde",
  "serde",
  "serde_bytes",

distant-net/Cargo.toml

@@ -25,10 +25,12 @@ log = "0.4.18"
 paste = "1.0.12"
 p256 = { version = "0.13.2", features = ["ecdh", "pem"] }
 rand = { version = "0.8.5", features = ["getrandom"] }
+rmp = "0.8.11"
 rmp-serde = "1.1.1"
 sha2 = "0.10.6"
 serde = { version = "1.0.163", features = ["derive"] }
 serde_bytes = "0.11.9"
+serde_json = "1.0.96"
 strum = { version = "0.24.1", features = ["derive"] }
 tokio = { version = "1.28.2", features = ["full"] }
 

distant-net/src/common.rs

@@ -20,4 +20,5 @@ pub use listener::*;
 pub use map::*;
 pub use packet::*;
 pub use port::*;
+pub use serde_json::Value;
 pub use transport::*;

distant-net/src/common/packet/header.rs

@@ -0,0 +1,80 @@
+use crate::common::{utils, Value};
+use derive_more::IntoIterator;
+use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
+use std::io;
+use std::ops::{Deref, DerefMut};
+
+/// Generates a new [`Header`] of key/value pairs based on literals.
+///
+/// ```
+/// use distant_net::header;
+///
+/// let _header = header!("key" -> "value", "key2" -> 123);
+/// ```
+#[macro_export]
+macro_rules! header {
+    ($($key:literal -> $value:expr),* $(,)?) => {{
+        let mut _header = $crate::common::Header::default();
+
+        $(
+            _header.insert($key, $value);
+        )*
+
+        _header
+    }};
+}
+
+/// Represents a packet header comprised of arbitrary data tied to string keys.
+#[derive(Clone, Debug, Default, PartialEq, Eq, IntoIterator, Serialize, Deserialize)]
+#[serde(transparent)]
+pub struct Header(HashMap<String, Value>);
+
+impl Header {
+    /// Creates an empty [`Header`] newtype wrapper.
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    /// Exists purely to support serde serialization checks.
+    #[inline]
+    pub(crate) fn is_empty(&self) -> bool {
+        self.0.is_empty()
+    }
+
+    /// Inserts a key-value pair into the map.
+    ///
+    /// If the map did not have this key present, [`None`] is returned.
+    ///
+    /// If the map did have this key present, the value is updated, and the old value is returned.
+    /// The key is not updated, though; this matters for types that can be `==` without being
+    /// identical. See the [module-level documentation](std::collections#insert-and-complex-keys)
+    /// for more.
+    pub fn insert(&mut self, key: impl Into<String>, value: impl Into<Value>) -> Option<Value> {
+        self.0.insert(key.into(), value.into())
+    }
+
+    /// Serializes the header into bytes.
+    pub fn to_vec(&self) -> io::Result<Vec<u8>> {
+        utils::serialize_to_vec(self)
+    }
+
+    /// Deserializes the header from bytes.
+    pub fn from_slice(slice: &[u8]) -> io::Result<Self> {
+        utils::deserialize_from_slice(slice)
+    }
+}
+
+impl Deref for Header {
+    type Target = HashMap<String, Value>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl DerefMut for Header {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}

distant-net/src/common/packet/request.rs

@@ -5,12 +5,17 @@ use derive_more::{Display, Error};
 use serde::de::DeserializeOwned;
 use serde::{Deserialize, Serialize};
 
-use super::{parse_msg_pack_str, write_str_msg_pack, Id};
+use super::{read_header_bytes, read_key_eq, read_str_bytes, Header, Id};
 use crate::common::utils;
+use crate::header;
 
 /// Represents a request to send
-#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
 pub struct Request<T> {
+    /// Optional header data to include with request
+    #[serde(default, skip_serializing_if = "Header::is_empty")]
+    pub header: Header,
+
     /// Unique id associated with the request
     pub id: Id,
 
@@ -19,9 +24,10 @@ pub struct Request<T> {
 }
 
 impl<T> Request<T> {
-    /// Creates a new request with a random, unique id
+    /// Creates a new request with a random, unique id and no header data
     pub fn new(payload: T) -> Self {
         Self {
+            header: header!(),
             id: rand::random::<u64>().to_string(),
             payload,
         }
@@ -45,6 +51,11 @@ where
     /// Attempts to convert a typed request to an untyped request
     pub fn to_untyped_request(&self) -> io::Result<UntypedRequest> {
         Ok(UntypedRequest {
+            header: Cow::Owned(if !self.header.is_empty() {
+                utils::serialize_to_vec(&self.header)?
+            } else {
+                Vec::new()
+            }),
             id: Cow::Borrowed(&self.id),
             payload: Cow::Owned(self.to_payload_vec()?),
         })
@@ -73,13 +84,34 @@ pub enum UntypedRequestParseError {
     /// When the bytes do not represent a request
     WrongType,
 
+    /// When a header should be present, but the key is wrong
+    InvalidHeaderKey,
+
+    /// When a header should be present, but the header bytes are wrong
+    InvalidHeader,
+
+    /// When the key for the id is wrong
+    InvalidIdKey,
+
     /// When the id is not a valid UTF-8 string
     InvalidId,
+
+    /// When the key for the payload is wrong
+    InvalidPayloadKey,
+}
+
+#[inline]
+fn header_is_empty(header: &[u8]) -> bool {
+    header.is_empty()
 }
 
 /// Represents a request to send whose payload is bytes instead of a specific type
-#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
 pub struct UntypedRequest<'a> {
+    /// Header data associated with the request as bytes
+    #[serde(default, skip_serializing_if = "header_is_empty")]
+    pub header: Cow<'a, [u8]>,
+
     /// Unique id associated with the request
     pub id: Cow<'a, str>,
 
@@ -91,6 +123,11 @@ impl<'a> UntypedRequest<'a> {
     /// Attempts to convert an untyped request to a typed request
     pub fn to_typed_request<T: DeserializeOwned>(&self) -> io::Result<Request<T>> {
         Ok(Request {
+            header: if header_is_empty(&self.header) {
+                header!()
+            } else {
+                utils::deserialize_from_slice(&self.header)?
+            },
             id: self.id.to_string(),
             payload: utils::deserialize_from_slice(&self.payload)?,
         })
@@ -99,6 +136,10 @@ impl<'a> UntypedRequest<'a> {
     /// Convert into a borrowed version
     pub fn as_borrowed(&self) -> UntypedRequest<'_> {
         UntypedRequest {
+            header: match &self.header {
+                Cow::Borrowed(x) => Cow::Borrowed(x),
+                Cow::Owned(x) => Cow::Borrowed(x.as_slice()),
+            },
             id: match &self.id {
                 Cow::Borrowed(x) => Cow::Borrowed(x),
                 Cow::Owned(x) => Cow::Borrowed(x.as_str()),
@@ -113,6 +154,10 @@ impl<'a> UntypedRequest<'a> {
     /// Convert into an owned version
     pub fn into_owned(self) -> UntypedRequest<'static> {
         UntypedRequest {
+            header: match self.header {
+                Cow::Borrowed(x) => Cow::Owned(x.to_vec()),
+                Cow::Owned(x) => Cow::Owned(x),
+            },
             id: match self.id {
                 Cow::Borrowed(x) => Cow::Owned(x.to_string()),
                 Cow::Owned(x) => Cow::Owned(x),
@@ -124,6 +169,11 @@ impl<'a> UntypedRequest<'a> {
         }
     }
 
+    /// Updates the header of the request to the given `header`.
+    pub fn set_header(&mut self, header: impl IntoIterator<Item = u8>) {
+        self.header = Cow::Owned(header.into_iter().collect());
+    }
+
     /// Updates the id of the request to the given `id`.
     pub fn set_id(&mut self, id: impl Into<String>) {
         self.id = Cow::Owned(id.into());
@@ -131,61 +181,80 @@ impl<'a> UntypedRequest<'a> {
 
     /// Allocates a new collection of bytes representing the request.
     pub fn to_bytes(&self) -> Vec<u8> {
-        let mut bytes = vec![0x82];
+        let mut bytes = vec![];
+
+        let has_header = !header_is_empty(&self.header);
+        if has_header {
+            rmp::encode::write_map_len(&mut bytes, 3).unwrap();
+        } else {
+            rmp::encode::write_map_len(&mut bytes, 2).unwrap();
+        }
+
+        if has_header {
+            rmp::encode::write_str(&mut bytes, "header").unwrap();
+            bytes.extend_from_slice(&self.header);
+        }
 
-        write_str_msg_pack("id", &mut bytes);
-        write_str_msg_pack(&self.id, &mut bytes);
+        rmp::encode::write_str(&mut bytes, "id").unwrap();
+        rmp::encode::write_str(&mut bytes, &self.id).unwrap();
 
-        write_str_msg_pack("payload", &mut bytes);
+        rmp::encode::write_str(&mut bytes, "payload").unwrap();
         bytes.extend_from_slice(&self.payload);
 
         bytes
     }
 
     /// Parses a collection of bytes, returning a partial request if it can be potentially
-    /// represented as a [`Request`] depending on the payload, or the original bytes if it does not
-    /// represent a [`Request`]
+    /// represented as a [`Request`] depending on the payload.
     ///
     /// NOTE: This supports parsing an invalid request where the payload would not properly
     /// deserialize, but the bytes themselves represent a complete request of some kind.
     pub fn from_slice(input: &'a [u8]) -> Result<Self, UntypedRequestParseError> {
-        if input.len() < 2 {
+        if input.is_empty() {
             return Err(UntypedRequestParseError::WrongType);
         }
 
-        // MsgPack marks a fixmap using 0x80 - 0x8f to indicate the size (up to 15 elements).
-        //
-        // In the case of the request, there are only two elements: id and payload. So the first
-        // byte should ALWAYS be 0x82 (130).
-        if input[0] != 0x82 {
-            return Err(UntypedRequestParseError::WrongType);
-        }
+        let has_header = match rmp::Marker::from_u8(input[0]) {
+            rmp::Marker::FixMap(2) => false,
+            rmp::Marker::FixMap(3) => true,
+            _ => return Err(UntypedRequestParseError::WrongType),
+        };
 
-        // Skip the first byte representing the fixmap
+        // Advance position by marker
         let input = &input[1..];
 
-        // Validate that first field is id
-        let (input, id_key) =
-            parse_msg_pack_str(input).map_err(|_| UntypedRequestParseError::WrongType)?;
-        if id_key != "id" {
-            return Err(UntypedRequestParseError::WrongType);
-        }
+        // Parse the header if we have one
+        let (header, input) = if has_header {
+            let (_, input) = read_key_eq(input, "header")
+                .map_err(|_| UntypedRequestParseError::InvalidHeaderKey)?;
+
+            let (header, input) =
+                read_header_bytes(input).map_err(|_| UntypedRequestParseError::InvalidHeader)?;
+            (header, input)
+        } else {
+            ([0u8; 0].as_slice(), input)
+        };
+
+        // Validate that next field is id
+        let (_, input) =
+            read_key_eq(input, "id").map_err(|_| UntypedRequestParseError::InvalidIdKey)?;
 
         // Get the id itself
-        let (input, id) =
-            parse_msg_pack_str(input).map_err(|_| UntypedRequestParseError::InvalidId)?;
+        let (id, input) = read_str_bytes(input).map_err(|_| UntypedRequestParseError::InvalidId)?;
 
-        // Validate that second field is payload
-        let (input, payload_key) =
-            parse_msg_pack_str(input).map_err(|_| UntypedRequestParseError::WrongType)?;
-        if payload_key != "payload" {
-            return Err(UntypedRequestParseError::WrongType);
-        }
+        // Validate that final field is payload
+        let (_, input) = read_key_eq(input, "payload")
+            .map_err(|_| UntypedRequestParseError::InvalidPayloadKey)?;
 
+        let header = Cow::Borrowed(header);
         let id = Cow::Borrowed(id);
         let payload = Cow::Borrowed(input);
 
-        Ok(Self { id, payload })
+        Ok(Self {
+            header,
+            id,
+            payload,
+        })
     }
 }
 
@@ -198,18 +267,33 @@ mod tests {
     const TRUE_BYTE: u8 = 0xc3;
     const NEVER_USED_BYTE: u8 = 0xc1;
 
+    // fixstr of 6 bytes with str "header"
+    const HEADER_FIELD_BYTES: &[u8] = &[0xa6, b'h', b'e', b'a', b'd', b'e', b'r'];
+
+    // fixmap of 2 objects with
+    // 1. key fixstr "key" and value fixstr "value"
+    // 1. key fixstr "num" and value fixint 123
+    const HEADER_BYTES: &[u8] = &[
+        0x82, // valid map with 2 pair
+        0xa3, b'k', b'e', b'y', // key: "key"
+        0xa5, b'v', b'a', b'l', b'u', b'e', // value: "value"
+        0xa3, b'n', b'u', b'm', // key: "num"
+        0x7b, // value: 123
+    ];
+
     // fixstr of 2 bytes with str "id"
-    const ID_FIELD_BYTES: &[u8] = &[0xa2, 0x69, 0x64];
+    const ID_FIELD_BYTES: &[u8] = &[0xa2, b'i', b'd'];
 
     // fixstr of 7 bytes with str "payload"
-    const PAYLOAD_FIELD_BYTES: &[u8] = &[0xa7, 0x70, 0x61, 0x79, 0x6c, 0x6f, 0x61, 0x64];
+    const PAYLOAD_FIELD_BYTES: &[u8] = &[0xa7, b'p', b'a', b'y', b'l', b'o', b'a', b'd'];
 
-    /// fixstr of 4 bytes with str "test"
-    const TEST_STR_BYTES: &[u8] = &[0xa4, 0x74, 0x65, 0x73, 0x74];
+    // fixstr of 4 bytes with str "test"
+    const TEST_STR_BYTES: &[u8] = &[0xa4, b't', b'e', b's', b't'];
 
     #[test]
     fn untyped_request_should_support_converting_to_bytes() {
         let bytes = Request {
+            header: header!(),
             id: "some id".to_string(),
             payload: true,
         }
@@ -220,9 +304,44 @@ mod tests {
         assert_eq!(untyped_request.to_bytes(), bytes);
     }
 
+    #[test]
+    fn untyped_request_should_support_converting_to_bytes_with_header() {
+        let bytes = Request {
+            header: header!("key" -> 123),
+            id: "some id".to_string(),
+            payload: true,
+        }
+        .to_vec()
+        .unwrap();
+
+        let untyped_request = UntypedRequest::from_slice(&bytes).unwrap();
+        assert_eq!(untyped_request.to_bytes(), bytes);
+    }
+
+    #[test]
+    fn untyped_request_should_support_parsing_from_request_bytes_with_header() {
+        let bytes = Request {
+            header: header!("key" -> 123),
+            id: "some id".to_string(),
+            payload: true,
+        }
+        .to_vec()
+        .unwrap();
+
+        assert_eq!(
+            UntypedRequest::from_slice(&bytes),
+            Ok(UntypedRequest {
+                header: Cow::Owned(utils::serialize_to_vec(&header!("key" -> 123)).unwrap()),
+                id: Cow::Borrowed("some id"),
+                payload: Cow::Owned(vec![TRUE_BYTE]),
+            })
+        );
+    }
+
     #[test]
     fn untyped_request_should_support_parsing_from_request_bytes_with_valid_payload() {
         let bytes = Request {
+            header: header!(),
             id: "some id".to_string(),
             payload: true,
         }
@@ -232,6 +351,7 @@ mod tests {
         assert_eq!(
             UntypedRequest::from_slice(&bytes),
             Ok(UntypedRequest {
+                header: Cow::Owned(vec![]),
                 id: Cow::Borrowed("some id"),
                 payload: Cow::Owned(vec![TRUE_BYTE]),
             })
@@ -242,6 +362,7 @@ mod tests {
     fn untyped_request_should_support_parsing_from_request_bytes_with_invalid_payload() {
         // Request with id < 32 bytes
         let mut bytes = Request {
+            header: header!(),
             id: "".to_string(),
             payload: true,
         }
@@ -255,12 +376,35 @@ mod tests {
         assert_eq!(
             UntypedRequest::from_slice(&bytes),
             Ok(UntypedRequest {
+                header: Cow::Owned(vec![]),
                 id: Cow::Owned("".to_string()),
                 payload: Cow::Owned(vec![TRUE_BYTE, NEVER_USED_BYTE]),
             })
         );
     }
 
+    #[test]
+    fn untyped_request_should_support_parsing_full_request() {
+        let input = [
+            &[0x83],
+            HEADER_FIELD_BYTES,
+            HEADER_BYTES,
+            ID_FIELD_BYTES,
+            TEST_STR_BYTES,
+            PAYLOAD_FIELD_BYTES,
+            &[TRUE_BYTE],
+        ]
+        .concat();
+
+        // Convert into typed so we can test
+        let untyped_request = UntypedRequest::from_slice(&input).unwrap();
+        let request: Request<bool> = untyped_request.to_typed_request().unwrap();
+
+        assert_eq!(request.header, header!("key" -> "value", "num" -> 123));
+        assert_eq!(request.id, "test");
+        assert!(request.payload);
+    }
+
     #[test]
     fn untyped_request_should_fail_to_parse_if_given_bytes_not_representing_a_request() {
         // Empty byte slice
@@ -281,10 +425,46 @@ mod tests {
             Err(UntypedRequestParseError::WrongType)
         );
 
+        // Invalid header key
+        assert_eq!(
+            UntypedRequest::from_slice(
+                [
+                    &[0x83],
+                    &[0xa0], // header key would be defined here, set to empty str
+                    HEADER_BYTES,
+                    ID_FIELD_BYTES,
+                    TEST_STR_BYTES,
+                    PAYLOAD_FIELD_BYTES,
+                    &[TRUE_BYTE],
+                ]
+                .concat()
+                .as_slice()
+            ),
+            Err(UntypedRequestParseError::InvalidHeaderKey)
+        );
+
+        // Invalid header bytes
+        assert_eq!(
+            UntypedRequest::from_slice(
+                [
+                    &[0x83],
+                    HEADER_FIELD_BYTES,
+                    &[0xa0], // header would be defined here, set to empty str
+                    ID_FIELD_BYTES,
+                    TEST_STR_BYTES,
+                    PAYLOAD_FIELD_BYTES,
+                    &[TRUE_BYTE],
+                ]
+                .concat()
+                .as_slice()
+            ),
+            Err(UntypedRequestParseError::InvalidHeader)
+        );
+
         // Missing fields (corrupt data)
         assert_eq!(
             UntypedRequest::from_slice(&[0x82]),
-            Err(UntypedRequestParseError::WrongType)
+            Err(UntypedRequestParseError::InvalidIdKey)
         );
 
         // Missing id field (has valid data itself)
@@ -300,7 +480,7 @@ mod tests {
                 .concat()
                 .as_slice()
             ),
-            Err(UntypedRequestParseError::WrongType)
+            Err(UntypedRequestParseError::InvalidIdKey)
         );
 
         // Non-str id field value
@@ -348,7 +528,7 @@ mod tests {
                 .concat()
                 .as_slice()
             ),
-            Err(UntypedRequestParseError::WrongType)
+            Err(UntypedRequestParseError::InvalidPayloadKey)
         );
     }
 }

distant-net/src/common/packet/response.rs

@@ -5,12 +5,17 @@ use derive_more::{Display, Error};
 use serde::de::DeserializeOwned;
 use serde::{Deserialize, Serialize};
 
-use super::{parse_msg_pack_str, write_str_msg_pack, Id};
+use super::{read_header_bytes, read_key_eq, read_str_bytes, Header, Id};
 use crate::common::utils;
+use crate::header;
 
 /// Represents a response received related to some response
-#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
 pub struct Response<T> {
+    /// Optional header data to include with response
+    #[serde(default, skip_serializing_if = "Header::is_empty")]
+    pub header: Header,
+
     /// Unique id associated with the response
     pub id: Id,
 
@@ -22,9 +27,10 @@ pub struct Response<T> {
 }
 
 impl<T> Response<T> {
-    /// Creates a new response with a random, unique id
+    /// Creates a new response with a random, unique id and no header data
     pub fn new(origin_id: Id, payload: T) -> Self {
         Self {
+            header: header!(),
             id: rand::random::<u64>().to_string(),
             origin_id,
             payload,
@@ -49,6 +55,11 @@ where
     /// Attempts to convert a typed response to an untyped response
     pub fn to_untyped_response(&self) -> io::Result<UntypedResponse> {
         Ok(UntypedResponse {
+            header: Cow::Owned(if !self.header.is_empty() {
+                utils::serialize_to_vec(&self.header)?
+            } else {
+                Vec::new()
+            }),
             id: Cow::Borrowed(&self.id),
             origin_id: Cow::Borrowed(&self.origin_id),
             payload: Cow::Owned(self.to_payload_vec()?),
@@ -72,16 +83,40 @@ pub enum UntypedResponseParseError {
     /// When the bytes do not represent a response
     WrongType,
 
+    /// When a header should be present, but the key is wrong
+    InvalidHeaderKey,
+
+    /// When a header should be present, but the header bytes are wrong
+    InvalidHeader,
+
+    /// When the key for the id is wrong
+    InvalidIdKey,
+
     /// When the id is not a valid UTF-8 string
     InvalidId,
 
+    /// When the key for the origin id is wrong
+    InvalidOriginIdKey,
+
     /// When the origin id is not a valid UTF-8 string
     InvalidOriginId,
+
+    /// When the key for the payload is wrong
+    InvalidPayloadKey,
+}
+
+#[inline]
+fn header_is_empty(header: &[u8]) -> bool {
+    header.is_empty()
 }
 
 /// Represents a response to send whose payload is bytes instead of a specific type
 #[derive(Clone, Debug, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
 pub struct UntypedResponse<'a> {
+    /// Header data associated with the response as bytes
+    #[serde(default, skip_serializing_if = "header_is_empty")]
+    pub header: Cow<'a, [u8]>,
+
     /// Unique id associated with the response
     pub id: Cow<'a, str>,
 
@@ -93,9 +128,14 @@ pub struct UntypedResponse<'a> {
 }
 
 impl<'a> UntypedResponse<'a> {
-    /// Attempts to convert an untyped request to a typed request
+    /// Attempts to convert an untyped response to a typed response
     pub fn to_typed_response<T: DeserializeOwned>(&self) -> io::Result<Response<T>> {
         Ok(Response {
+            header: if header_is_empty(&self.header) {
+                header!()
+            } else {
+                utils::deserialize_from_slice(&self.header)?
+            },
             id: self.id.to_string(),
             origin_id: self.origin_id.to_string(),
             payload: utils::deserialize_from_slice(&self.payload)?,
@@ -105,6 +145,10 @@ impl<'a> UntypedResponse<'a> {
     /// Convert into a borrowed version
     pub fn as_borrowed(&self) -> UntypedResponse<'_> {
         UntypedResponse {
+            header: match &self.header {
+                Cow::Borrowed(x) => Cow::Borrowed(x),
+                Cow::Owned(x) => Cow::Borrowed(x.as_slice()),
+            },
             id: match &self.id {
                 Cow::Borrowed(x) => Cow::Borrowed(x),
                 Cow::Owned(x) => Cow::Borrowed(x.as_str()),
@@ -123,6 +167,10 @@ impl<'a> UntypedResponse<'a> {
     /// Convert into an owned version
     pub fn into_owned(self) -> UntypedResponse<'static> {
         UntypedResponse {
+            header: match self.header {
+                Cow::Borrowed(x) => Cow::Owned(x.to_vec()),
+                Cow::Owned(x) => Cow::Owned(x),
+            },
             id: match self.id {
                 Cow::Borrowed(x) => Cow::Owned(x.to_string()),
                 Cow::Owned(x) => Cow::Owned(x),
@@ -138,6 +186,11 @@ impl<'a> UntypedResponse<'a> {
         }
     }
 
+    /// Updates the header of the response to the given `header`.
+    pub fn set_header(&mut self, header: impl IntoIterator<Item = u8>) {
+        self.header = Cow::Owned(header.into_iter().collect());
+    }
+
     /// Updates the id of the response to the given `id`.
     pub fn set_id(&mut self, id: impl Into<String>) {
         self.id = Cow::Owned(id.into());
@@ -150,76 +203,90 @@ impl<'a> UntypedResponse<'a> {
 
     /// Allocates a new collection of bytes representing the response.
     pub fn to_bytes(&self) -> Vec<u8> {
-        let mut bytes = vec![0x83];
+        let mut bytes = vec![];
 
-        write_str_msg_pack("id", &mut bytes);
-        write_str_msg_pack(&self.id, &mut bytes);
+        let has_header = !header_is_empty(&self.header);
+        if has_header {
+            rmp::encode::write_map_len(&mut bytes, 4).unwrap();
+        } else {
+            rmp::encode::write_map_len(&mut bytes, 3).unwrap();
+        }
 
-        write_str_msg_pack("origin_id", &mut bytes);
-        write_str_msg_pack(&self.origin_id, &mut bytes);
+        if has_header {
+            rmp::encode::write_str(&mut bytes, "header").unwrap();
+            bytes.extend_from_slice(&self.header);
+        }
 
-        write_str_msg_pack("payload", &mut bytes);
+        rmp::encode::write_str(&mut bytes, "id").unwrap();
+        rmp::encode::write_str(&mut bytes, &self.id).unwrap();
+
+        rmp::encode::write_str(&mut bytes, "origin_id").unwrap();
+        rmp::encode::write_str(&mut bytes, &self.origin_id).unwrap();
+
+        rmp::encode::write_str(&mut bytes, "payload").unwrap();
         bytes.extend_from_slice(&self.payload);
 
         bytes
     }
 
     /// Parses a collection of bytes, returning an untyped response if it can be potentially
-    /// represented as a [`Response`] depending on the payload, or the original bytes if it does not
-    /// represent a [`Response`].
+    /// represented as a [`Response`] depending on the payload.
     ///
     /// NOTE: This supports parsing an invalid response where the payload would not properly
     /// deserialize, but the bytes themselves represent a complete response of some kind.
     pub fn from_slice(input: &'a [u8]) -> Result<Self, UntypedResponseParseError> {
-        if input.len() < 2 {
+        if input.is_empty() {
             return Err(UntypedResponseParseError::WrongType);
         }
 
-        // MsgPack marks a fixmap using 0x80 - 0x8f to indicate the size (up to 15 elements).
-        //
-        // In the case of the request, there are only three elements: id, origin_id, and payload.
-        // So the first byte should ALWAYS be 0x83 (131).
-        if input[0] != 0x83 {
-            return Err(UntypedResponseParseError::WrongType);
-        }
+        let has_header = match rmp::Marker::from_u8(input[0]) {
+            rmp::Marker::FixMap(3) => false,
+            rmp::Marker::FixMap(4) => true,
+            _ => return Err(UntypedResponseParseError::WrongType),
+        };
 
-        // Skip the first byte representing the fixmap
+        // Advance position by marker
         let input = &input[1..];
 
-        // Validate that first field is id
-        let (input, id_key) =
-            parse_msg_pack_str(input).map_err(|_| UntypedResponseParseError::WrongType)?;
-        if id_key != "id" {
-            return Err(UntypedResponseParseError::WrongType);
-        }
+        // Parse the header if we have one
+        let (header, input) = if has_header {
+            let (_, input) = read_key_eq(input, "header")
+                .map_err(|_| UntypedResponseParseError::InvalidHeaderKey)?;
+
+            let (header, input) =
+                read_header_bytes(input).map_err(|_| UntypedResponseParseError::InvalidHeader)?;
+            (header, input)
+        } else {
+            ([0u8; 0].as_slice(), input)
+        };
+
+        // Validate that next field is id
+        let (_, input) =
+            read_key_eq(input, "id").map_err(|_| UntypedResponseParseError::InvalidIdKey)?;
 
         // Get the id itself
-        let (input, id) =
-            parse_msg_pack_str(input).map_err(|_| UntypedResponseParseError::InvalidId)?;
+        let (id, input) =
+            read_str_bytes(input).map_err(|_| UntypedResponseParseError::InvalidId)?;
 
-        // Validate that second field is origin_id
-        let (input, origin_id_key) =
-            parse_msg_pack_str(input).map_err(|_| UntypedResponseParseError::WrongType)?;
-        if origin_id_key != "origin_id" {
-            return Err(UntypedResponseParseError::WrongType);
-        }
+        // Validate that next field is origin_id
+        let (_, input) = read_key_eq(input, "origin_id")
+            .map_err(|_| UntypedResponseParseError::InvalidOriginIdKey)?;
 
         // Get the origin_id itself
-        let (input, origin_id) =
-            parse_msg_pack_str(input).map_err(|_| UntypedResponseParseError::InvalidOriginId)?;
+        let (origin_id, input) =
+            read_str_bytes(input).map_err(|_| UntypedResponseParseError::InvalidOriginId)?;
 
-        // Validate that second field is payload
-        let (input, payload_key) =
-            parse_msg_pack_str(input).map_err(|_| UntypedResponseParseError::WrongType)?;
-        if payload_key != "payload" {
-            return Err(UntypedResponseParseError::WrongType);
-        }
+        // Validate that final field is payload
+        let (_, input) = read_key_eq(input, "payload")
+            .map_err(|_| UntypedResponseParseError::InvalidPayloadKey)?;
 
+        let header = Cow::Borrowed(header);
         let id = Cow::Borrowed(id);
         let origin_id = Cow::Borrowed(origin_id);
         let payload = Cow::Borrowed(input);
 
         Ok(Self {
+            header,
             id,
             origin_id,
             payload,
@@ -236,22 +303,52 @@ mod tests {
     const TRUE_BYTE: u8 = 0xc3;
     const NEVER_USED_BYTE: u8 = 0xc1;
 
+    // fixstr of 6 bytes with str "header"
+    const HEADER_FIELD_BYTES: &[u8] = &[0xa6, b'h', b'e', b'a', b'd', b'e', b'r'];
+
+    // fixmap of 2 objects with
+    // 1. key fixstr "key" and value fixstr "value"
+    // 1. key fixstr "num" and value fixint 123
+    const HEADER_BYTES: &[u8] = &[
+        0x82, // valid map with 2 pair
+        0xa3, b'k', b'e', b'y', // key: "key"
+        0xa5, b'v', b'a', b'l', b'u', b'e', // value: "value"
+        0xa3, b'n', b'u', b'm', // key: "num"
+        0x7b, // value: 123
+    ];
+
     // fixstr of 2 bytes with str "id"
-    const ID_FIELD_BYTES: &[u8] = &[0xa2, 0x69, 0x64];
+    const ID_FIELD_BYTES: &[u8] = &[0xa2, b'i', b'd'];
 
     // fixstr of 9 bytes with str "origin_id"
     const ORIGIN_ID_FIELD_BYTES: &[u8] =
         &[0xa9, 0x6f, 0x72, 0x69, 0x67, 0x69, 0x6e, 0x5f, 0x69, 0x64];
 
     // fixstr of 7 bytes with str "payload"
-    const PAYLOAD_FIELD_BYTES: &[u8] = &[0xa7, 0x70, 0x61, 0x79, 0x6c, 0x6f, 0x61, 0x64];
+    const PAYLOAD_FIELD_BYTES: &[u8] = &[0xa7, b'p', b'a', b'y', b'l', b'o', b'a', b'd'];
 
     /// fixstr of 4 bytes with str "test"
-    const TEST_STR_BYTES: &[u8] = &[0xa4, 0x74, 0x65, 0x73, 0x74];
+    const TEST_STR_BYTES: &[u8] = &[0xa4, b't', b'e', b's', b't'];
 
     #[test]
     fn untyped_response_should_support_converting_to_bytes() {
         let bytes = Response {
+            header: header!(),
+            id: "some id".to_string(),
+            origin_id: "some origin id".to_string(),
+            payload: true,
+        }
+        .to_vec()
+        .unwrap();
+
+        let untyped_response = UntypedResponse::from_slice(&bytes).unwrap();
+        assert_eq!(untyped_response.to_bytes(), bytes);
+    }
+
+    #[test]
+    fn untyped_response_should_support_converting_to_bytes_with_header() {
+        let bytes = Response {
+            header: header!("key" -> 123),
             id: "some id".to_string(),
             origin_id: "some origin id".to_string(),
             payload: true,
@@ -263,9 +360,32 @@ mod tests {
         assert_eq!(untyped_response.to_bytes(), bytes);
     }
 
+    #[test]
+    fn untyped_response_should_support_parsing_from_response_bytes_with_header() {
+        let bytes = Response {
+            header: header!("key" -> 123),
+            id: "some id".to_string(),
+            origin_id: "some origin id".to_string(),
+            payload: true,
+        }
+        .to_vec()
+        .unwrap();
+
+        assert_eq!(
+            UntypedResponse::from_slice(&bytes),
+            Ok(UntypedResponse {
+                header: Cow::Owned(utils::serialize_to_vec(&header!("key" -> 123)).unwrap()),
+                id: Cow::Borrowed("some id"),
+                origin_id: Cow::Borrowed("some origin id"),
+                payload: Cow::Owned(vec![TRUE_BYTE]),
+            })
+        );
+    }
+
     #[test]
     fn untyped_response_should_support_parsing_from_response_bytes_with_valid_payload() {
         let bytes = Response {
+            header: header!(),
             id: "some id".to_string(),
             origin_id: "some origin id".to_string(),
             payload: true,
@@ -276,6 +396,7 @@ mod tests {
         assert_eq!(
             UntypedResponse::from_slice(&bytes),
             Ok(UntypedResponse {
+                header: Cow::Owned(vec![]),
                 id: Cow::Borrowed("some id"),
                 origin_id: Cow::Borrowed("some origin id"),
                 payload: Cow::Owned(vec![TRUE_BYTE]),
@@ -287,6 +408,7 @@ mod tests {
     fn untyped_response_should_support_parsing_from_response_bytes_with_invalid_payload() {
         // Response with id < 32 bytes
         let mut bytes = Response {
+            header: header!(),
             id: "".to_string(),
             origin_id: "".to_string(),
             payload: true,
@@ -301,6 +423,7 @@ mod tests {
         assert_eq!(
             UntypedResponse::from_slice(&bytes),
             Ok(UntypedResponse {
+                header: Cow::Owned(vec![]),
                 id: Cow::Owned("".to_string()),
                 origin_id: Cow::Owned("".to_string()),
                 payload: Cow::Owned(vec![TRUE_BYTE, NEVER_USED_BYTE]),
@@ -308,6 +431,31 @@ mod tests {
         );
     }
 
+    #[test]
+    fn untyped_response_should_support_parsing_full_request() {
+        let input = [
+            &[0x84],
+            HEADER_FIELD_BYTES,
+            HEADER_BYTES,
+            ID_FIELD_BYTES,
+            TEST_STR_BYTES,
+            ORIGIN_ID_FIELD_BYTES,
+            &[0xa2, b'o', b'g'],
+            PAYLOAD_FIELD_BYTES,
+            &[TRUE_BYTE],
+        ]
+        .concat();
+
+        // Convert into typed so we can test
+        let untyped_response = UntypedResponse::from_slice(&input).unwrap();
+        let response: Response<bool> = untyped_response.to_typed_response().unwrap();
+
+        assert_eq!(response.header, header!("key" -> "value", "num" -> 123));
+        assert_eq!(response.id, "test");
+        assert_eq!(response.origin_id, "og");
+        assert!(response.payload);
+    }
+
     #[test]
     fn untyped_response_should_fail_to_parse_if_given_bytes_not_representing_a_response() {
         // Empty byte slice
@@ -328,10 +476,50 @@ mod tests {
             Err(UntypedResponseParseError::WrongType)
         );
 
+        // Invalid header key
+        assert_eq!(
+            UntypedResponse::from_slice(
+                [
+                    &[0x84],
+                    &[0xa0], // header key would be defined here, set to empty str
+                    HEADER_BYTES,
+                    ID_FIELD_BYTES,
+                    TEST_STR_BYTES,
+                    ORIGIN_ID_FIELD_BYTES,
+                    TEST_STR_BYTES,
+                    PAYLOAD_FIELD_BYTES,
+                    &[TRUE_BYTE],
+                ]
+                .concat()
+                .as_slice()
+            ),
+            Err(UntypedResponseParseError::InvalidHeaderKey)
+        );
+
+        // Invalid header bytes
+        assert_eq!(
+            UntypedResponse::from_slice(
+                [
+                    &[0x84],
+                    HEADER_FIELD_BYTES,
+                    &[0xa0], // header would be defined here, set to empty str
+                    ID_FIELD_BYTES,
+                    TEST_STR_BYTES,
+                    ORIGIN_ID_FIELD_BYTES,
+                    TEST_STR_BYTES,
+                    PAYLOAD_FIELD_BYTES,
+                    &[TRUE_BYTE],
+                ]
+                .concat()
+                .as_slice()
+            ),
+            Err(UntypedResponseParseError::InvalidHeader)
+        );
+
         // Missing fields (corrupt data)
         assert_eq!(
             UntypedResponse::from_slice(&[0x83]),
-            Err(UntypedResponseParseError::WrongType)
+            Err(UntypedResponseParseError::InvalidIdKey)
         );
 
         // Missing id field (has valid data itself)
@@ -349,7 +537,7 @@ mod tests {
                 .concat()
                 .as_slice()
             ),
-            Err(UntypedResponseParseError::WrongType)
+            Err(UntypedResponseParseError::InvalidIdKey)
         );
 
         // Non-str id field value
@@ -403,7 +591,7 @@ mod tests {
                 .concat()
                 .as_slice()
             ),
-            Err(UntypedResponseParseError::WrongType)
+            Err(UntypedResponseParseError::InvalidOriginIdKey)
         );
 
         // Non-str origin_id field value
@@ -457,7 +645,7 @@ mod tests {
                 .concat()
                 .as_slice()
             ),
-            Err(UntypedResponseParseError::WrongType)
+            Err(UntypedResponseParseError::InvalidPayloadKey)
         );
     }
 }

distant-protocol/src/response.rs

@@ -631,7 +631,7 @@ mod tests {
                 value,
                 serde_json::json!({
                     "type": "changed",
-                    "ts": u64::MAX,
+                    "timestamp": u64::MAX,
                     "kind": "access",
                     "path": "path",
                 })
@@ -657,13 +657,13 @@ mod tests {
                 value,
                 serde_json::json!({
                     "type": "changed",
-                    "ts": u64::MAX,
+                    "timestamp": u64::MAX,
                     "kind": "access",
                     "path": "path",
                     "details": {
                         "attribute": "permissions",
                         "renamed": "renamed",
-                        "ts": u64::MAX,
+                        "timestamp": u64::MAX,
                         "extra": "info",
                     },
                 })
@@ -674,7 +674,7 @@ mod tests {
         fn should_be_able_to_deserialize_minimal_payload_from_json() {
             let value = serde_json::json!({
                 "type": "changed",
-                "ts": u64::MAX,
+                "timestamp": u64::MAX,
                 "kind": "access",
                 "path": "path",
             });
@@ -695,13 +695,13 @@ mod tests {
         fn should_be_able_to_deserialize_full_payload_from_json() {
             let value = serde_json::json!({
                 "type": "changed",
-                "ts": u64::MAX,
+                "timestamp": u64::MAX,
                 "kind": "access",
                 "path": "path",
                 "details": {
                     "attribute": "permissions",
                     "renamed": "renamed",
-                    "ts": u64::MAX,
+                    "timestamp": u64::MAX,
                     "extra": "info",
                 },
             });