2 files changed, 90 insertions, 13 deletions

diff --git a/src/transport/mod.rs b/src/transport/mod.rs
index eeba47f..aee3ee3 100644
--- a/src/transport/mod.rs
+++ b/src/transport/mod.rs
@@ -1,6 +1,8 @@
 use anyhow::Result;
 use async_trait::async_trait;
+use automerge::ChangeHash;
 use automerge::sync::Message;
+use sha2::{Digest, Sha256};
 use std::boxed::Box;
 use std::future::Future;
 use std::pin::Pin;
@@ -10,13 +12,54 @@ pub use crate::document::replica::Epoch;
 pub mod firewall;
 pub mod mycelium;
 
-pub type HeadsDigest = [u8; 32];
+/// Most heads listed in a keepalive; more are sent as a digest. 16 heads are about
+/// 520 bytes, well within one Mycelium chunk (1300 bytes).
+pub const MAX_LISTED_HEADS: usize = 16;
 
+/// A document's heads, compactly: the heads themselves while there are few, otherwise a
+/// digest of them. Either way, two summaries are equal exactly when the heads are (the
+/// number of heads decides the form, so equal heads always take the same one).
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum HeadsSummary {
+    /// Sorted.
+    Heads(Vec<ChangeHash>),
+    /// SHA-256 over the sorted heads.
+    Digest([u8; 32]),
+}
+
+impl HeadsSummary {
+    pub fn of(mut heads: Vec<ChangeHash>) -> Self {
+        heads.sort_unstable();
+        if heads.len() <= MAX_LISTED_HEADS {
+            return Self::Heads(heads);
+        }
+        let mut hasher = Sha256::new();
+        for head in &heads {
+            hasher.update(head.0);
+        }
+        Self::Digest(hasher.finalize().into())
+    }
+
+    /// The heads, if listed.
+    pub fn heads(&self) -> Option<&[ChangeHash]> {
+        match self {
+            Self::Heads(heads) => Some(heads),
+            Self::Digest(_) => None,
+        }
+    }
+}
+
+/// PING and PONG. Addressed to one peer: `syncs_with_you` is about the recipient.
 #[derive(Debug, Clone)]
 pub struct Keepalive {
-    pub heads_digest: HeadsDigest,
-    /// The digest is only comparable between nodes on the same epoch.
+    /// Only comparable between nodes on the same epoch.
+    pub heads: HeadsSummary,
     pub epoch: Epoch,
+    /// The sender is ready (same meaning as `SyncEnvelope::ready`).
+    pub ready: bool,
+    /// The sender would sync with the recipient: counts it as an active member on its
+    /// own epoch.
+    pub syncs_with_you: bool,
 }
 
 #[derive(Debug)]
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,
     })
 }