add chunked download and use it during rejoin

1 files changed, 146 insertions, 0 deletions

diff --git a/src/transport/mycelium.rs b/src/transport/mycelium.rs
index cc7f25f..a8ab104 100644
--- a/src/transport/mycelium.rs
+++ b/src/transport/mycelium.rs
@@ -19,6 +19,7 @@ use tokio::process::{Child, Command};
 use url::Url;
 
 use crate::transport::EncodedMessage;
+use crate::transport::{BlobBody, BlobId, BlobMessage};
 use crate::transport::{Epoch, HelloReply, HelloRequest, SyncEnvelope};
 use crate::transport::{HeadsSummary, Keepalive, MAX_LISTED_HEADS};
 use crate::transport::{MessageHandler, ParsedMessage, Transport};
@@ -31,6 +32,7 @@ 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_BLOB: &str = "continuum.v1.blob";
 
 const EMPTY_QUEUE_DELAY: Duration = Duration::from_millis(25);
 
@@ -747,6 +749,7 @@ fn decode_payload(topic_encoded: &str, payload: &str) -> Result<ParsedMessage, T
         }
         TOPIC_PING => Ok(ParsedMessage::Ping(bytes_to_keepalive(&bytes)?)),
         TOPIC_PONG => Ok(ParsedMessage::Pong(bytes_to_keepalive(&bytes)?)),
+        TOPIC_BLOB => Ok(ParsedMessage::Blob(bytes_to_blob(&bytes)?)),
         t => Err(TransportError::UnknownTopic(t.to_string())),
     }
 }
@@ -774,6 +777,7 @@ 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::Blob(blob) => (TOPIC_BLOB, blob_to_bytes(blob)?),
     };
     let outbound = OutboundMessage {
         dst: Destination {
@@ -1021,6 +1025,107 @@ pub fn bytes_to_sync_envelope(bytes: &[u8]) -> Result<SyncEnvelope, TransportErr
     })
 }
 
+// --- BlobMessage ---
+#[derive(Serialize, Deserialize)]
+enum BlobIdShadow {
+    Change(ByteBuf),
+}
+
+#[derive(Serialize, Deserialize)]
+enum BlobBodyShadow {
+    Request {
+        chunks: Vec<u32>,
+    },
+    Manifest {
+        size: u64,
+        chunk_size: u32,
+        hashes: Vec<ByteBuf>,
+    },
+    Chunk {
+        index: u32,
+        data: ByteBuf,
+    },
+    Missing,
+}
+
+#[derive(Serialize, Deserialize)]
+struct BlobShadow {
+    epoch: EpochShadow,
+    id: BlobIdShadow,
+    body: BlobBodyShadow,
+}
+
+fn blob_to_bytes(msg: &BlobMessage) -> Result<Vec<u8>, rmp_serde::encode::Error> {
+    let id = match msg.id {
+        BlobId::Change(hash) => BlobIdShadow::Change(ByteBuf::from(hash.0.to_vec())),
+    };
+    let body = match &msg.body {
+        BlobBody::Request { chunks } => BlobBodyShadow::Request {
+            chunks: chunks.clone(),
+        },
+        BlobBody::Manifest {
+            size,
+            chunk_size,
+            hashes,
+        } => BlobBodyShadow::Manifest {
+            size: *size,
+            chunk_size: *chunk_size,
+            hashes: hashes
+                .iter()
+                .map(|hash| ByteBuf::from(hash.to_vec()))
+                .collect(),
+        },
+        BlobBody::Chunk { index, data } => BlobBodyShadow::Chunk {
+            index: *index,
+            data: ByteBuf::from(data.to_vec()),
+        },
+        BlobBody::Missing => BlobBodyShadow::Missing,
+    };
+    rmp_serde::to_vec(&BlobShadow {
+        epoch: epoch_to_shadow(&msg.epoch),
+        id,
+        body,
+    })
+}
+
+fn hash32(bytes: &[u8], what: &str) -> Result<[u8; 32], TransportError> {
+    bytes
+        .try_into()
+        .map_err(|_| TransportError::InvalidPayload(format!("{what} must be 32 bytes")))
+}
+
+fn bytes_to_blob(bytes: &[u8]) -> Result<BlobMessage, TransportError> {
+    let shadow: BlobShadow = rmp_serde::from_slice(bytes)?;
+    let id = match shadow.id {
+        BlobIdShadow::Change(hash) => BlobId::Change(ChangeHash(hash32(&hash, "a change hash")?)),
+    };
+    let body = match shadow.body {
+        BlobBodyShadow::Request { chunks } => BlobBody::Request { chunks },
+        BlobBodyShadow::Manifest {
+            size,
+            chunk_size,
+            hashes,
+        } => BlobBody::Manifest {
+            size,
+            chunk_size,
+            hashes: hashes
+                .iter()
+                .map(|hash| hash32(hash, "a chunk hash"))
+                .collect::<Result<_, _>>()?,
+        },
+        BlobBodyShadow::Chunk { index, data } => BlobBody::Chunk {
+            index,
+            data: data.into_vec().into(),
+        },
+        BlobBodyShadow::Missing => BlobBody::Missing,
+    };
+    Ok(BlobMessage {
+        epoch: epoch_from_shadow(shadow.epoch)?,
+        id,
+        body,
+    })
+}
+
 #[derive(Error, Debug)]
 pub enum TransportError {
     #[error("Mycelium delivered a message without a sender public key")]
@@ -1087,6 +1192,47 @@ mod tests {
     }
 
     #[test]
+    fn blob_messages_round_trip() {
+        let epoch = Epoch {
+            generation: 3,
+            id: [7; 32],
+        };
+        let id = BlobId::Change(ChangeHash([9; 32]));
+        let chunk = BlobMessage {
+            epoch,
+            id,
+            body: BlobBody::Chunk {
+                index: 5,
+                data: vec![1u8; 64 * 1024].into(),
+            },
+        };
+        let bytes = blob_to_bytes(&chunk).unwrap();
+        // A 64 KiB chunk costs only a few bytes of framing.
+        assert!(bytes.len() < 64 * 1024 + 128, "{} bytes", bytes.len());
+        let decoded = bytes_to_blob(&bytes).unwrap();
+        assert_eq!(decoded.id, id);
+        assert_eq!(decoded.epoch, epoch);
+        assert!(
+            matches!(decoded.body, BlobBody::Chunk { index: 5, ref data } if data.len() == 64 * 1024)
+        );
+
+        let manifest = BlobMessage {
+            epoch,
+            id,
+            body: BlobBody::Manifest {
+                size: 100,
+                chunk_size: 64,
+                hashes: vec![[1; 32], [2; 32]],
+            },
+        };
+        let decoded = bytes_to_blob(&blob_to_bytes(&manifest).unwrap()).unwrap();
+        assert!(matches!(
+            decoded.body,
+            BlobBody::Manifest { size: 100, chunk_size: 64, ref hashes } if hashes.len() == 2
+        ));
+    }
+
+    #[test]
     fn listed_heads_equal_their_digest() {
         let heads: Vec<_> = (0..3u8).map(|n| ChangeHash([n; 32])).collect();
         let listed = HeadsSummary::of(heads.clone());