path: root/src/transport/wire.rs

//! The wire format of a message: one byte saying what it is (`MessageKind`), then the
//! message itself in MessagePack, as derived from its type in `transport` (structs as
//! arrays, fields in declaration order). The codecs below make hashes, IDs and bytes
//! compact (MessagePack binaries, not arrays of numbers) and check their lengths.
//!
//! The wire format may change freely until the first release: there is no compatibility
//! with older nodes to keep.

use thiserror::Error;

use super::{MessageKind, ParsedMessage};

/// What the first byte of a payload says.
fn tag(kind: MessageKind) -> u8 {
    match kind {
        MessageKind::Hello => 0,
        MessageKind::HelloReply => 1,
        MessageKind::Sync => 2,
        MessageKind::Ping => 3,
        MessageKind::Pong => 4,
        MessageKind::Blob => 5,
    }
}

#[derive(Error, Debug)]
pub enum WireError {
    #[error("empty payload")]
    Empty,

    #[error("unknown message kind {0}")]
    UnknownKind(u8),

    #[error("MessagePack serialization failed: {0}")]
    Encode(#[from] rmp_serde::encode::Error),

    #[error("MessagePack deserialization failed: {0}")]
    Decode(#[from] rmp_serde::decode::Error),
}

/// The payload of `message`.
pub(crate) fn encode(message: &ParsedMessage) -> Result<Vec<u8>, WireError> {
    let mut payload = vec![tag(message.kind())];
    match message {
        ParsedMessage::Hello(hello) => rmp_serde::encode::write(&mut payload, hello)?,
        ParsedMessage::HelloReply(reply) => rmp_serde::encode::write(&mut payload, reply)?,
        ParsedMessage::Sync(envelope) => rmp_serde::encode::write(&mut payload, envelope)?,
        ParsedMessage::Ping(keepalive) | ParsedMessage::Pong(keepalive) => {
            rmp_serde::encode::write(&mut payload, keepalive)?
        }
        ParsedMessage::Blob(blob) => rmp_serde::encode::write(&mut payload, blob)?,
    }
    Ok(payload)
}

/// The message in `payload` (as made by [`encode`]).
pub(crate) fn decode(payload: &[u8]) -> Result<ParsedMessage, WireError> {
    let (&kind, body) = payload.split_first().ok_or(WireError::Empty)?;
    let is = |wanted: MessageKind| kind == tag(wanted);
    Ok(if is(MessageKind::Hello) {
        ParsedMessage::Hello(rmp_serde::from_slice(body)?)
    } else if is(MessageKind::HelloReply) {
        ParsedMessage::HelloReply(rmp_serde::from_slice(body)?)
    } else if is(MessageKind::Sync) {
        ParsedMessage::Sync(rmp_serde::from_slice(body)?)
    } else if is(MessageKind::Ping) {
        ParsedMessage::Ping(rmp_serde::from_slice(body)?)
    } else if is(MessageKind::Pong) {
        ParsedMessage::Pong(rmp_serde::from_slice(body)?)
    } else if is(MessageKind::Blob) {
        ParsedMessage::Blob(rmp_serde::from_slice(body)?)
    } else {
        return Err(WireError::UnknownKind(kind));
    })
}

/// Field codecs (`#[serde(with = "...")]`) for the message types.
pub(crate) mod codec {
    use serde::de::Error as _;
    use serde::{Deserialize, Deserializer, Serialize, Serializer};
    use serde_bytes::{ByteBuf, Bytes};

    /// Exactly `N` bytes, or a decoding error naming `what`.
    fn fixed<const N: usize, E: serde::de::Error>(bytes: &[u8], what: &str) -> Result<[u8; N], E> {
        bytes
            .try_into()
            .map_err(|_| E::custom(format!("{what} must be {N} bytes")))
    }

    /// A run ID (`RunId`).
    pub(crate) mod run_id {
        use super::*;

        pub(crate) fn serialize<S: Serializer>(value: &[u8; 16], s: S) -> Result<S::Ok, S::Error> {
            Bytes::new(value).serialize(s)
        }

        pub(crate) fn deserialize<'de, D: Deserializer<'de>>(d: D) -> Result<[u8; 16], D::Error> {
            fixed(&ByteBuf::deserialize(d)?, "a run id")
        }
    }

    /// A 32-byte hash.
    pub(crate) mod bytes32 {
        use super::*;

        pub(crate) fn serialize<S: Serializer>(value: &[u8; 32], s: S) -> Result<S::Ok, S::Error> {
            Bytes::new(value).serialize(s)
        }

        pub(crate) fn deserialize<'de, D: Deserializer<'de>>(d: D) -> Result<[u8; 32], D::Error> {
            fixed(&ByteBuf::deserialize(d)?, "a hash")
        }
    }

    /// 32-byte hashes.
    pub(crate) mod hashes32 {
        use super::*;

        pub(crate) fn serialize<S: Serializer>(
            value: &[[u8; 32]],
            s: S,
        ) -> Result<S::Ok, S::Error> {
            s.collect_seq(value.iter().map(|hash| Bytes::new(hash)))
        }

        pub(crate) fn deserialize<'de, D: Deserializer<'de>>(
            d: D,
        ) -> Result<Vec<[u8; 32]>, D::Error> {
            Vec::<ByteBuf>::deserialize(d)?
                .iter()
                .map(|hash| fixed(hash, "a hash"))
                .collect()
        }
    }

    /// An Automerge change hash.
    pub(crate) mod change_hash {
        use super::*;
        use automerge::ChangeHash;

        pub(crate) fn serialize<S: Serializer>(
            value: &ChangeHash,
            s: S,
        ) -> Result<S::Ok, S::Error> {
            Bytes::new(&value.0).serialize(s)
        }

        pub(crate) fn deserialize<'de, D: Deserializer<'de>>(d: D) -> Result<ChangeHash, D::Error> {
            Ok(ChangeHash(fixed(
                &ByteBuf::deserialize(d)?,
                "a change hash",
            )?))
        }
    }

    /// An epoch: generation and ID.
    pub(crate) mod epoch {
        use super::*;
        use crate::document::replica::Epoch;

        pub(crate) fn serialize<S: Serializer>(value: &Epoch, s: S) -> Result<S::Ok, S::Error> {
            (value.generation, Bytes::new(&value.id)).serialize(s)
        }

        pub(crate) fn deserialize<'de, D: Deserializer<'de>>(d: D) -> Result<Epoch, D::Error> {
            let (generation, id): (u64, ByteBuf) = Deserialize::deserialize(d)?;
            Ok(Epoch {
                generation,
                id: fixed(&id, "an epoch id")?,
            })
        }
    }

    /// Heads as a keepalive carries them: listed (at most `MAX_LISTED_HEADS`, digest
    /// empty), or as a digest (none listed).
    pub(crate) mod heads {
        use super::*;
        use crate::transport::{HeadsSummary, MAX_LISTED_HEADS};
        use automerge::ChangeHash;

        pub(crate) fn serialize<S: Serializer>(
            value: &HeadsSummary,
            s: S,
        ) -> Result<S::Ok, S::Error> {
            let (listed, digest): (Vec<&Bytes>, &Bytes) = match value {
                HeadsSummary::Heads(heads) => (
                    heads.iter().map(|head| Bytes::new(&head.0)).collect(),
                    Bytes::new(&[]),
                ),
                HeadsSummary::Digest(digest) => (Vec::new(), Bytes::new(digest)),
            };
            (listed, digest).serialize(s)
        }

        pub(crate) fn deserialize<'de, D: Deserializer<'de>>(
            d: D,
        ) -> Result<HeadsSummary, D::Error> {
            let (listed, digest): (Vec<ByteBuf>, ByteBuf) = Deserialize::deserialize(d)?;
            if !digest.is_empty() {
                return Ok(HeadsSummary::Digest(fixed(&digest, "a heads digest")?));
            }
            if listed.len() > MAX_LISTED_HEADS {
                return Err(D::Error::custom(format!(
                    "more than {MAX_LISTED_HEADS} heads listed"
                )));
            }
            listed
                .iter()
                .map(|head| -> Result<ChangeHash, D::Error> {
                    Ok(ChangeHash(fixed(head, "a head")?))
                })
                .collect::<Result<Vec<_>, _>>()
                .map(HeadsSummary::Heads)
        }
    }

    /// An Automerge sync message, in Automerge's own encoding.
    pub(crate) mod sync_message {
        use super::*;
        use automerge::sync::Message;

        pub(crate) fn serialize<S: Serializer>(value: &Message, s: S) -> Result<S::Ok, S::Error> {
            Bytes::new(&value.clone().encode()).serialize(s)
        }

        pub(crate) fn deserialize<'de, D: Deserializer<'de>>(d: D) -> Result<Message, D::Error> {
            Message::decode(&ByteBuf::deserialize(d)?)
                .map_err(|error| D::Error::custom(format!("Automerge decode failed: {error:?}")))
        }
    }

    /// Shared bytes (`bytes::Bytes`).
    pub(crate) mod shared_bytes {
        use super::*;

        pub(crate) fn serialize<S: Serializer>(
            value: &bytes::Bytes,
            s: S,
        ) -> Result<S::Ok, S::Error> {
            Bytes::new(value).serialize(s)
        }

        pub(crate) fn deserialize<'de, D: Deserializer<'de>>(
            d: D,
        ) -> Result<bytes::Bytes, D::Error> {
            Ok(ByteBuf::deserialize(d)?.into_vec().into())
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::document::replica::Epoch;
    use crate::transport::{
        BlobBody, BlobId, BlobMessage, HeadsSummary, HelloReply, HelloRequest, Keepalive,
        MAX_LISTED_HEADS,
    };
    use automerge::ChangeHash;

    #[test]
    fn handshakes_round_trip() {
        let epoch = Epoch {
            generation: 2,
            id: [3; 32],
        };
        let hello = ParsedMessage::Hello(HelloRequest {
            swarm_id: Some("swarm".into()),
            epoch,
            epoch_seq: 1,
            auth: vec![1, 2, 3],
            ready: true,
            run_id: [4; 16],
        });
        let ParsedMessage::Hello(decoded) = decode(&encode(&hello).unwrap()).unwrap() else {
            panic!("not a HELLO");
        };
        assert_eq!(decoded.swarm_id.as_deref(), Some("swarm"));
        assert_eq!((decoded.epoch, decoded.epoch_seq), (epoch, 1));
        assert_eq!(decoded.auth, vec![1, 2, 3]);
        assert_eq!(decoded.run_id, [4; 16]);

        let reply = ParsedMessage::HelloReply(HelloReply {
            swarm_id: "swarm".into(),
            peers: vec!["p".into()],
            epoch,
            epoch_seq: 1,
            auth: Vec::new(),
            ready: false,
            run_id: [4; 16],
            away_ms: Some(5),
            sender_run_id: [6; 16],
        });
        let ParsedMessage::HelloReply(decoded) = decode(&encode(&reply).unwrap()).unwrap() else {
            panic!("not a HELLO_REPLY");
        };
        assert_eq!(decoded.peers, vec!["p".to_owned()]);
        assert_eq!(decoded.away_ms, Some(5));
        assert_eq!(decoded.sender_run_id, [6; 16]);
    }

    #[test]
    fn keepalives_keep_their_kind_and_heads() {
        let keepalive = Keepalive {
            heads: HeadsSummary::Digest([8; 32]),
            epoch: Epoch::GENESIS,
            epoch_seq: 0,
            ready: true,
            syncs_with_you: false,
        };
        let decoded = decode(&encode(&ParsedMessage::Pong(keepalive.clone())).unwrap()).unwrap();
        let ParsedMessage::Pong(decoded) = decoded else {
            panic!("not a PONG");
        };
        assert_eq!(decoded.heads, keepalive.heads);
    }

    #[test]
    fn too_many_listed_heads_are_refused() {
        let keepalive = Keepalive {
            heads: HeadsSummary::Heads(
                (0..=MAX_LISTED_HEADS as u8)
                    .map(|n| ChangeHash([n; 32]))
                    .collect(),
            ),
            epoch: Epoch::GENESIS,
            epoch_seq: 0,
            ready: true,
            syncs_with_you: true,
        };
        assert!(decode(&encode(&ParsedMessage::Ping(keepalive)).unwrap()).is_err());
    }

    #[test]
    fn unknown_kinds_and_bad_payloads_are_refused() {
        assert!(matches!(decode(&[]), Err(WireError::Empty)));
        assert!(matches!(decode(&[99]), Err(WireError::UnknownKind(99))));

        let blob = ParsedMessage::Blob(BlobMessage {
            epoch: Epoch::GENESIS,
            epoch_seq: 0,
            id: BlobId::Content([1; 32]),
            body: BlobBody::Missing,
        });
        let payload = encode(&blob).unwrap();
        assert!(matches!(decode(&payload).unwrap(), ParsedMessage::Blob(_)));
        // A payload cut short doesn't decode.
        assert!(decode(&payload[..payload.len() - 1]).is_err());
        // Nor one of another kind.
        let mut relabelled = payload.clone();
        relabelled[0] = tag(MessageKind::Hello);
        assert!(decode(&relabelled).is_err());
    }
}