1 files changed, 44 insertions, 10 deletions

diff --git a/src/transport/mycelium.rs b/src/transport/mycelium.rs
index 193beaa..9d1fb90 100644
--- a/src/transport/mycelium.rs
+++ b/src/transport/mycelium.rs
@@ -1,5 +1,6 @@
 use anyhow::{Context, Result, anyhow, bail};
 use async_trait::async_trait;
+use automerge::ChangeHash;
 use base64::engine::{Engine, general_purpose::STANDARD as BASE64};
 use reqwest::StatusCode;
 use serde::{Deserialize, Deserializer, Serialize, Serializer};
@@ -18,8 +19,8 @@ use tokio::process::{Child, Command};
 use url::Url;
 
 use crate::transport::EncodedMessage;
-use crate::transport::Keepalive;
 use crate::transport::{Epoch, HelloReply, HelloRequest, SyncEnvelope};
+use crate::transport::{HeadsSummary, Keepalive, MAX_LISTED_HEADS};
 use crate::transport::{MessageHandler, ParsedMessage, Transport};
 
 const MYCELIUM_POLL_TIMEOUT_SECS: u64 = 300;
@@ -764,30 +765,63 @@ fn epoch_from_shadow(shadow: EpochShadow) -> Result<Epoch, TransportError> {
     })
 }
 
+/// Either `heads` (listed) or `digest` (32 bytes, `heads` empty) is set.
 #[derive(Serialize, Deserialize)]
 struct KeepaliveShadow {
-    heads_digest: ByteBuf,
-    #[serde(default)]
+    heads: Vec<ByteBuf>,
+    digest: ByteBuf,
     epoch: EpochShadow,
+    ready: bool,
+    syncs_with_you: bool,
 }
 
 fn keepalive_to_bytes(msg: &Keepalive) -> Result<Vec<u8>, rmp_serde::encode::Error> {
+    let (heads, digest) = match &msg.heads {
+        HeadsSummary::Heads(heads) => (
+            heads
+                .iter()
+                .map(|head| ByteBuf::from(head.0.to_vec()))
+                .collect(),
+            ByteBuf::new(),
+        ),
+        HeadsSummary::Digest(digest) => (Vec::new(), ByteBuf::from(digest.to_vec())),
+    };
     rmp_serde::to_vec(&KeepaliveShadow {
-        heads_digest: ByteBuf::from(msg.heads_digest.to_vec()),
+        heads,
+        digest,
         epoch: epoch_to_shadow(&msg.epoch),
+        ready: msg.ready,
+        syncs_with_you: msg.syncs_with_you,
     })
 }
 
 fn bytes_to_keepalive(bytes: &[u8]) -> Result<Keepalive, TransportError> {
     let shadow: KeepaliveShadow = rmp_serde::from_slice(bytes)?;
-    let heads_digest = shadow
-        .heads_digest
-        .as_slice()
-        .try_into()
-        .map_err(|_| TransportError::InvalidPayload("heads digest must be 32 bytes".into()))?;
+    let heads = if shadow.digest.is_empty() {
+        if shadow.heads.len() > MAX_LISTED_HEADS {
+            return Err(TransportError::InvalidPayload(format!(
+                "more than {MAX_LISTED_HEADS} heads listed"
+            )));
+        }
+        let heads = shadow
+            .heads
+            .iter()
+            .map(|head| ChangeHash::try_from(head.as_slice()))
+            .collect::<Result<Vec<_>, _>>()
+            .map_err(|_| TransportError::InvalidPayload("a head must be 32 bytes".into()))?;
+        HeadsSummary::Heads(heads)
+    } else {
+        let digest =
+            shadow.digest.as_slice().try_into().map_err(|_| {
+                TransportError::InvalidPayload("heads digest must be 32 bytes".into())
+            })?;
+        HeadsSummary::Digest(digest)
+    };
     Ok(Keepalive {
-        heads_digest,
+        heads,
         epoch: epoch_from_shadow(shadow.epoch)?,
+        ready: shadow.ready,
+        syncs_with_you: shadow.syncs_with_you,
     })
 }