use anyhow::{Context, Result};
use bytes::Bytes;
use iroh::{Endpoint, EndpointId};
use iroh_blobs::store::mem::MemStore;
use iroh_blobs::ticket::BlobTicket;
use iroh_blobs::{BlobFormat, BlobsProtocol};
use iroh_gossip::api::{Event, GossipReceiver, GossipSender};
use iroh_gossip::net::Gossip;
use iroh_gossip::proto::TopicId;
use n0_future::StreamExt;
use serde::{Deserialize, Serialize};
pub fn roster_topic_bytes(org_id: &str) -> [u8; 32] {
*blake3::hash(format!("mcpmesh/roster/{org_id}").as_bytes()).as_bytes()
}
pub fn presence_topic_bytes(org_id: &str) -> [u8; 32] {
*blake3::hash(format!("mcpmesh/presence/{org_id}").as_bytes()).as_bytes()
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RosterAnnounce {
pub serial: u64,
pub roster_hash: String, pub blob_ticket: String, }
impl RosterAnnounce {
pub fn to_bytes(&self) -> Vec<u8> {
serde_json::to_vec(self).expect("RosterAnnounce serializes")
}
pub fn from_bytes(b: &[u8]) -> Result<Self> {
serde_json::from_slice(b).context("parse roster announce")
}
}
pub struct RosterGossip {
pub sender: GossipSender,
pub receiver: Option<GossipReceiver>,
}
pub fn spawn_gossip(endpoint: &Endpoint) -> Gossip {
Gossip::builder().spawn(endpoint.clone())
}
pub async fn subscribe(
gossip: &Gossip,
topic: [u8; 32],
bootstrap: Vec<EndpointId>,
) -> Result<RosterGossip> {
let topic_id = TopicId::from_bytes(topic);
let topic = gossip
.subscribe(topic_id, bootstrap)
.await
.context("subscribe to gossip topic")?;
let (sender, receiver) = topic.split();
Ok(RosterGossip {
sender,
receiver: Some(receiver),
})
}
pub async fn broadcast(sender: &GossipSender, payload: Vec<u8>) -> Result<()> {
sender
.broadcast(Bytes::from(payload))
.await
.context("broadcast gossip message")
}
pub async fn next_message(receiver: &mut GossipReceiver) -> Option<Bytes> {
while let Some(ev) = receiver.next().await {
if let Ok(Event::Received(msg)) = ev {
return Some(msg.content);
}
}
None
}
pub const GOSSIP_ALPN: &[u8] = iroh_gossip::ALPN;
pub const BLOB_ALPN: &[u8] = iroh_blobs::ALPN;
#[derive(Clone)]
pub struct RosterBlobs {
store: MemStore,
}
impl RosterBlobs {
pub fn new(_endpoint: &Endpoint) -> Self {
Self {
store: MemStore::new(),
}
}
pub fn protocol(&self) -> BlobsProtocol {
BlobsProtocol::new(&self.store, None)
}
#[cfg(test)]
pub(crate) fn spawn_accept(&self, endpoint: &Endpoint) {
let proto = self.protocol();
let ep = endpoint.clone();
tokio::spawn(async move {
while let Some(incoming) = ep.accept().await {
if let Ok(conn) = incoming.await
&& conn.alpn() == BLOB_ALPN
{
let _ = iroh::protocol::ProtocolHandler::accept(&proto, conn).await;
}
}
});
}
pub async fn publish(&self, doc: &[u8], endpoint: &Endpoint) -> Result<(String, String)> {
let tag = self
.store
.blobs()
.add_bytes(doc.to_vec())
.await
.context("add roster blob")?;
let ticket = BlobTicket::new(endpoint.addr(), tag.hash, BlobFormat::Raw);
let hash_hex = format!("blake3:{}", blake3::hash(doc).to_hex());
Ok((ticket.to_string(), hash_hex))
}
pub async fn fetch(
&self,
ticket_str: &str,
roster_hash: &str,
endpoint: &Endpoint,
) -> Result<Vec<u8>> {
let ticket: BlobTicket = ticket_str.parse().context("parse blob ticket")?;
self.store
.downloader(endpoint)
.download(ticket.hash(), Some(ticket.addr().id))
.await
.context("download roster blob")?;
let bytes = self
.store
.get_bytes(ticket.hash())
.await
.context("read fetched roster blob")?
.to_vec();
let got = format!("blake3:{}", blake3::hash(&bytes).to_hex());
if got != roster_hash {
anyhow::bail!("roster blob hash mismatch: announce {roster_hash}, fetched {got}");
}
Ok(bytes)
}
}
pub struct RosterAddrBook {
lookup: iroh::address_lookup::MemoryLookup,
known: std::sync::Mutex<std::collections::HashSet<[u8; 32]>>,
cap: usize,
}
impl RosterAddrBook {
pub fn register(endpoint: &iroh::Endpoint, cap: usize) -> Self {
let lookup = iroh::address_lookup::MemoryLookup::new();
if let Ok(al) = endpoint.address_lookup() {
al.add(lookup.clone());
}
Self {
lookup,
known: std::sync::Mutex::new(std::collections::HashSet::new()),
cap,
}
}
pub fn note(&self, addr: iroh::EndpointAddr) -> bool {
let id = *addr.id.as_bytes();
let mut known = self.known.lock().expect("roster addr book mutex");
if known.contains(&id) || known.len() >= self.cap {
return false;
}
known.insert(id);
self.lookup.add_endpoint_info(addr);
true
}
pub fn known_len(&self) -> usize {
self.known.lock().expect("roster addr book mutex").len()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn topics_are_deterministic_distinct_and_org_scoped() {
assert_eq!(roster_topic_bytes("acme"), roster_topic_bytes("acme"));
assert_ne!(roster_topic_bytes("acme"), presence_topic_bytes("acme"));
assert_ne!(roster_topic_bytes("acme"), roster_topic_bytes("globex"));
assert_eq!(
roster_topic_bytes("acme"),
*blake3::hash(b"mcpmesh/roster/acme").as_bytes()
);
}
#[test]
fn roster_announce_round_trips_json() {
let a = RosterAnnounce {
serial: 42,
roster_hash: "blake3:deadbeef".into(),
blob_ticket: "blobabc123".into(),
};
let bytes = a.to_bytes();
let back = RosterAnnounce::from_bytes(&bytes).expect("valid announce");
assert_eq!(back, a);
assert!(RosterAnnounce::from_bytes(b"not json").is_err());
assert!(
bytes.len() < 512,
"announce is compact: {} bytes",
bytes.len()
);
}
#[tokio::test]
async fn blob_add_ticket_and_fetch_round_trips_with_hash_check() {
tokio::time::timeout(std::time::Duration::from_secs(30), async {
let provider_ep = iroh::Endpoint::builder(iroh::endpoint::presets::Minimal)
.relay_mode(iroh::RelayMode::Disabled)
.alpns(vec![GOSSIP_ALPN.to_vec(), BLOB_ALPN.to_vec()])
.bind()
.await
.unwrap();
let provider = RosterBlobs::new(&provider_ep);
provider.spawn_accept(&provider_ep);
let doc = br#"{"format":"mcpmesh-roster/1","serial":7}"#.to_vec();
let (ticket, hash_hex) = provider.publish(&doc, &provider_ep).await.unwrap();
let fetcher_ep = iroh::Endpoint::builder(iroh::endpoint::presets::Minimal)
.relay_mode(iroh::RelayMode::Disabled)
.bind()
.await
.unwrap();
let mem = iroh::address_lookup::MemoryLookup::new();
mem.add_endpoint_info(provider_ep.addr());
fetcher_ep
.address_lookup()
.expect("address lookup services")
.add(mem);
let fetcher = RosterBlobs::new(&fetcher_ep);
let fetched = fetcher
.fetch(&ticket, &hash_hex, &fetcher_ep)
.await
.unwrap();
assert_eq!(fetched, doc, "blob fetch returns the exact bytes");
assert!(
fetcher
.fetch(&ticket, "blake3:0000", &fetcher_ep)
.await
.is_err(),
"a wrong roster_hash rejects the fetched blob"
);
})
.await
.expect("blob round-trip timed out");
}
#[tokio::test]
async fn roster_addr_book_dedups_and_caps() {
let ep = iroh::Endpoint::builder(iroh::endpoint::presets::Minimal)
.relay_mode(iroh::RelayMode::Disabled)
.bind()
.await
.unwrap();
let p1 = iroh::Endpoint::builder(iroh::endpoint::presets::Minimal)
.relay_mode(iroh::RelayMode::Disabled)
.bind()
.await
.unwrap();
let p2 = iroh::Endpoint::builder(iroh::endpoint::presets::Minimal)
.relay_mode(iroh::RelayMode::Disabled)
.bind()
.await
.unwrap();
let book = RosterAddrBook::register(&ep, 1);
assert!(book.note(p1.addr()), "first provider added");
assert!(
!book.note(p1.addr()),
"same provider is a no-op (dedup by id)"
);
assert_eq!(book.known_len(), 1);
assert!(
!book.note(p2.addr()),
"second distinct provider refused at the cap (bounded)"
);
assert_eq!(book.known_len(), 1, "the known set never exceeds the cap");
}
}