plumtree

1 files changed, 21 insertions, 2 deletions

diff --git a/src/transport/wire.rs b/src/transport/wire.rs
index bf46fc9..729acd3 100644
--- a/src/transport/wire.rs
+++ b/src/transport/wire.rs
@@ -14,7 +14,7 @@ use super::{EncodedMessage, MessageKind, ParsedMessage};
 
 /// The wire format's version, the first byte of every payload. Bump it with every change
 /// to a message type or a codec (`golden_bytes` fails then).
-pub(crate) const VERSION: u8 = 36;
+pub(crate) const VERSION: u8 = 37;
 
 /// What the second byte of a payload says.
 fn tag(kind: MessageKind) -> u8 {
@@ -33,6 +33,8 @@ fn tag(kind: MessageKind) -> u8 {
         MessageKind::LinksReply => 11,
         MessageKind::Probe => 12,
         MessageKind::ProbeReply => 13,
+        MessageKind::Prune => 14,
+        MessageKind::Graft => 15,
     }
 }
 
@@ -89,6 +91,9 @@ pub(crate) fn encode(message: &ParsedMessage) -> Result<Vec<u8>, WireError> {
         ParsedMessage::LinksReply(reply) => rmp_serde::encode::write(&mut payload, reply)?,
         ParsedMessage::Probe(request) => rmp_serde::encode::write(&mut payload, request)?,
         ParsedMessage::ProbeReply(reply) => rmp_serde::encode::write(&mut payload, reply)?,
+        ParsedMessage::Prune(link) | ParsedMessage::Graft(link) => {
+            rmp_serde::encode::write(&mut payload, link)?
+        }
     }
     Ok(payload)
 }
@@ -162,6 +167,10 @@ pub(crate) fn decode(payload: &[u8]) -> Result<ParsedMessage, WireError> {
         ParsedMessage::Probe(rmp_serde::from_slice(body)?)
     } else if is(MessageKind::ProbeReply) {
         ParsedMessage::ProbeReply(rmp_serde::from_slice(body)?)
+    } else if is(MessageKind::Prune) {
+        ParsedMessage::Prune(rmp_serde::from_slice(body)?)
+    } else if is(MessageKind::Graft) {
+        ParsedMessage::Graft(rmp_serde::from_slice(body)?)
     } else {
         return Err(WireError::UnknownKind(kind));
     })
@@ -532,6 +541,16 @@ pub(crate) fn samples() -> Vec<ParsedMessage> {
             member: "q".into(),
             reached: None,
         }),
+        ParsedMessage::Prune(crate::transport::TreeLink {
+            epoch,
+            epoch_seq: 1,
+            writer: 5,
+        }),
+        ParsedMessage::Graft(crate::transport::TreeLink {
+            epoch,
+            epoch_seq: 1,
+            writer: 5,
+        }),
     ]
 }
 
@@ -563,7 +582,7 @@ mod tests {
     #[test]
     fn every_kind_round_trips() {
         let samples = samples();
-        assert_eq!(samples.len(), 14, "a kind is missing from `samples`");
+        assert_eq!(samples.len(), 16, "a kind is missing from `samples`");
         for message in samples {
             let bytes = encode(&message).unwrap();
             let decoded = decode(&bytes).unwrap();