1 files changed, 59 insertions, 1 deletions

diff --git a/src/mycelium.rs b/src/mycelium.rs
index 8d0821b..22d49c1 100644
--- a/src/mycelium.rs
+++ b/src/mycelium.rs
@@ -19,7 +19,9 @@ use url::Url;
 
 use crate::transport::EncodedMessage;
 use crate::transport::Keepalive;
-use crate::transport::{Epoch, HelloReply, HelloRequest, SyncEnvelope};
+use crate::transport::{
+    Epoch, HelloReply, HelloRequest, SnapshotReply, SnapshotRequest, SyncEnvelope,
+};
 use crate::transport::{MessageHandler, ParsedMessage, Transport};
 
 const MYCELIUM_POLL_TIMEOUT_SECS: u64 = 300;
@@ -30,6 +32,8 @@ const TOPIC_HELLO_REPLY: &str = "continuum.v1.hello-reply";
 const TOPIC_SYNC: &str = "continuum.v1.sync";
 const TOPIC_PING: &str = "continuum.v1.ping";
 const TOPIC_PONG: &str = "continuum.v1.pong";
+const TOPIC_SNAPSHOT_REQUEST: &str = "continuum.v1.snapshot-request";
+const TOPIC_SNAPSHOT_REPLY: &str = "continuum.v1.snapshot-reply";
 
 const EMPTY_QUEUE_DELAY: Duration = Duration::from_millis(25);
 
@@ -696,6 +700,12 @@ fn decode_payload(topic: &str, payload: &str) -> Result<ParsedMessage, Transport
         }
         TOPIC_PING => Ok(ParsedMessage::Ping(bytes_to_keepalive(&bytes)?)),
         TOPIC_PONG => Ok(ParsedMessage::Pong(bytes_to_keepalive(&bytes)?)),
+        TOPIC_SNAPSHOT_REQUEST => Ok(ParsedMessage::SnapshotRequest(bytes_to_snapshot_request(
+            &bytes,
+        )?)),
+        TOPIC_SNAPSHOT_REPLY => Ok(ParsedMessage::SnapshotReply(bytes_to_snapshot_reply(
+            &bytes,
+        )?)),
         t => Err(TransportError::UnknownTopic(t.to_string())),
     }
 }
@@ -723,6 +733,10 @@ fn format_message(
         ParsedMessage::Sync(env) => (TOPIC_SYNC, sync_envelope_to_bytes(env)?),
         ParsedMessage::Ping(k) => (TOPIC_PING, keepalive_to_bytes(k)?),
         ParsedMessage::Pong(k) => (TOPIC_PONG, keepalive_to_bytes(k)?),
+        ParsedMessage::SnapshotRequest(req) => {
+            (TOPIC_SNAPSHOT_REQUEST, snapshot_request_to_bytes(req)?)
+        }
+        ParsedMessage::SnapshotReply(res) => (TOPIC_SNAPSHOT_REPLY, snapshot_reply_to_bytes(res)?),
     };
     Ok(OutboundMessage {
         dst: Destination {
@@ -900,6 +914,50 @@ pub fn bytes_to_sync_envelope(bytes: &[u8]) -> Result<SyncEnvelope, TransportErr
     })
 }
 
+// --- Snapshot transfer ---
+#[derive(Serialize, Deserialize)]
+struct SnapshotRequestShadow {
+    epoch: EpochShadow,
+    wanted: EpochShadow,
+}
+
+#[derive(Serialize, Deserialize)]
+struct SnapshotReplyShadow {
+    epoch: EpochShadow,
+    document: ByteBuf,
+}
+
+fn snapshot_request_to_bytes(msg: &SnapshotRequest) -> Result<Vec<u8>, rmp_serde::encode::Error> {
+    rmp_serde::to_vec(&SnapshotRequestShadow {
+        epoch: epoch_to_shadow(&msg.epoch),
+        wanted: epoch_to_shadow(&msg.wanted),
+    })
+}
+
+fn bytes_to_snapshot_request(bytes: &[u8]) -> Result<SnapshotRequest, TransportError> {
+    let shadow: SnapshotRequestShadow = rmp_serde::from_slice(bytes)?;
+    Ok(SnapshotRequest {
+        epoch: epoch_from_shadow(shadow.epoch)?,
+        wanted: epoch_from_shadow(shadow.wanted)?,
+    })
+}
+
+fn snapshot_reply_to_bytes(msg: &SnapshotReply) -> Result<Vec<u8>, rmp_serde::encode::Error> {
+    rmp_serde::to_vec(&SnapshotReplyShadow {
+        epoch: epoch_to_shadow(&msg.epoch),
+        // One copy of the document; replies are rare.
+        document: ByteBuf::from(msg.document.clone()),
+    })
+}
+
+fn bytes_to_snapshot_reply(bytes: &[u8]) -> Result<SnapshotReply, TransportError> {
+    let shadow: SnapshotReplyShadow = rmp_serde::from_slice(bytes)?;
+    Ok(SnapshotReply {
+        epoch: epoch_from_shadow(shadow.epoch)?,
+        document: shadow.document.into_vec(),
+    })
+}
+
 #[derive(Error, Debug)]
 pub enum TransportError {
     #[error("Mycelium delivered a message without a sender public key")]