johnandersen777.bsky.social/knotless
git that lives in your Atmosphere account
Forked from chadtmiller.com/knotless
0
0
//! Keeps mirrors warm from `com.atproto.sync.subscribeRepos`: a commit that
//! touches a knotless collection marks its account's mirrors stale and
//! refreshes them, so the next clone needs no round trip to the PDS.
use crate::App;
use knotless::drisl::{link, record};
use knotless::jetstream::knotless;
pub use knotless::jetstream::{Event, IDLE, Op, Socket, frames};
use knotless::jetstream::{RETRY, RETRY_MAX};
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use tokio_tungstenite::tungstenite::Message;
/// Reads a frame. `None` when it is not one.
pub fn parse(frame: &[u8]) -> Option<Event> {
let mut reader = frame;
let header: ciborium::Value = ciborium::from_reader(&mut reader).ok()?;
let body: ciborium::Value = ciborium::from_reader(&mut reader).ok()?;
let field = |v: &ciborium::Value, name: &str| -> Option<ciborium::Value> {
v.as_map()?
.iter()
.find(|(k, _)| k.as_text() == Some(name))
.map(|(_, v)| v.clone())
};
let text = |v: &ciborium::Value, name: &str| field(v, name)?.as_text().map(str::to_string);
let kind = text(&header, "t")?;
let mut event = Event {
seq: field(&body, "seq")
.and_then(|s| s.as_integer())
.and_then(|i| i64::try_from(i).ok()),
..Event::default()
};
match kind.as_str() {
"#commit" => {
let blocks = field(&body, "blocks");
let blocks = blocks
.as_ref()
.and_then(|b| b.as_bytes())
.map(Vec::as_slice)
.map(car);
for op in field(&body, "ops")?.as_array()? {
let (Some(action), Some(path)) = (text(op, "action"), text(op, "path")) else {
continue;
};
let Some((collection, rkey)) = path.split_once('/') else {
continue;
};
if knotless(collection) {
let cid = field(op, "cid");
let block = cid.as_ref().and_then(link).zip(blocks.as_ref());
let record = block.and_then(|(cid, blocks)| record(blocks.get(cid)?));
event.ops.push(Op {
action,
collection: collection.into(),
rkey: rkey.into(),
record,
});
}
}
if !event.ops.is_empty() {
event.did = Some(text(&body, "repo")?);
event.rev = text(&body, "rev");
}
}
"#sync" | "#account" | "#identity" => {
event.did = Some(text(&body, "did")?);
event.identity = kind == "#identity";
let active = field(&body, "active").and_then(|a| a.as_bool());
let status = knotless::jetstream::status(active, text(&body, "status").as_deref());
event.account = (kind == "#account").then_some(status);
}
_ => {}
}
Some(event)
}
/// The blocks of a CAR file, by their CIDs' bytes: a commit's carry the
/// records it wrote, which every follower of a shared PDS would otherwise
/// fetch, each of them, for every write.
fn car(bytes: &[u8]) -> HashMap<&[u8], &[u8]> {
let mut blocks = HashMap::new();
let mut rest = bytes;
// The header first, then each block: its length, its CID, its data.
let mut sections = std::iter::from_fn(|| {
let len = usize::try_from(varint(&mut rest)?).ok()?;
let section = rest.get(..len)?;
rest = &rest[len..];
Some(section)
});
sections.next();
for section in sections {
let mut cid = section;
// A CIDv1: version, codec, then a multihash's code and length.
let read = (0..4).try_fold(0, |_, _| varint(&mut cid));
let at = read.and_then(|digest| {
let digest = usize::try_from(digest).ok()?;
Some(section.len() - cid.len() + digest).filter(|at| *at <= section.len())
});
if let Some(at) = at {
blocks.insert(§ion[..at], §ion[at..]);
}
}
blocks
}
/// An unsigned LEB128 integer off the front of `bytes`.
fn varint(bytes: &mut &[u8]) -> Option<u64> {
let mut value = 0u64;
for (i, &byte) in bytes.iter().enumerate().take(10) {
value |= u64::from(byte & 0x7f) << (7 * i);
if byte & 0x80 == 0 {
*bytes = &bytes[i + 1..];
return Some(value);
}
}
None
}
/// The subscribeRepos endpoint for a relay or PDS base URL.
pub fn endpoint(base: &str, cursor: Option<i64>) -> String {
let mut url = base.trim_end_matches('/').to_string();
if !url.contains("/xrpc/") {
url.push_str("/xrpc/com.atproto.sync.subscribeRepos");
}
if let Some(c) = cursor {
url.push_str(&format!("?cursor={c}"));
}
url
}
/// Follows the firehose until the process ends, reconnecting from the last
/// sequence number. While disconnected, mirrors fall back to the TTL. A
/// connection that ends is made again after [`RETRY`], and one refused
/// waits twice as long as the last, up to [`RETRY_MAX`], so a relay that is
/// down is not asked twice a second.
pub async fn watch(app: Arc<App>, base: String) {
let mut cursor = None;
let mut wait = RETRY;
app.follows("firehose", &base);
loop {
let connected = tokio_tungstenite::connect_async(endpoint(&base, cursor)).await;
wait = next_wait(wait, connected.is_ok());
if let Ok((mut ws, _)) = connected {
app.connected("firehose", true);
frames(&mut ws, IDLE, |msg| {
if let Message::Binary(frame) = msg
&& let Some(event) = parse(&frame)
{
cursor = event.seq.or(cursor);
app.heard("firehose", event.seq);
app.observe(event);
}
})
.await;
}
app.connected("firehose", false);
tokio::time::sleep(wait).await;
}
}
/// How long to wait before the next attempt: [`RETRY`] after a connection,
/// and after a refusal twice `wait`, up to [`RETRY_MAX`].
fn next_wait(wait: Duration, connected: bool) -> Duration {
if connected {
RETRY
} else {
(wait * 2).min(RETRY_MAX)
}
}
#[cfg(test)]
mod tests {
use super::*;
use ciborium::Value as V;
#[test]
fn a_refused_relay_is_asked_less_often() {
let mut wait = RETRY;
for _ in 0..20 {
wait = next_wait(wait, false);
}
assert_eq!(wait, RETRY_MAX);
assert_eq!(next_wait(RETRY, false), RETRY * 2);
assert_eq!(next_wait(RETRY_MAX, true), RETRY);
}
fn frame(header: V, body: V) -> Vec<u8> {
let mut out = Vec::new();
ciborium::into_writer(&header, &mut out).unwrap();
ciborium::into_writer(&body, &mut out).unwrap();
out
}
fn map(pairs: Vec<(&str, V)>) -> V {
V::Map(
pairs
.into_iter()
.map(|(k, v)| (V::Text(k.into()), v))
.collect(),
)
}
#[test]
fn frames() {
let commit = |ops: Vec<V>| {
frame(
map(vec![
("op", V::Integer(1.into())),
("t", V::Text("#commit".into())),
]),
map(vec![
("seq", V::Integer(7.into())),
("repo", V::Text("did:plc:a".into())),
("ops", V::Array(ops)),
]),
)
};
let op = |action: &str, path: &str| {
map(vec![
("action", V::Text(action.into())),
("path", V::Text(path.into())),
])
};
let event = parse(&commit(vec![
op("create", "dev.knotless.ref/k.x"),
op("delete", "dev.knotless.issue/3"),
op("create", "app.bsky.feed.post/x"),
op("create", "no-slash"),
map(vec![("path", V::Text("dev.knotless.ref/y".into()))]),
]))
.unwrap();
assert_eq!(event.seq, Some(7));
assert_eq!(event.did.as_deref(), Some("did:plc:a"));
assert_eq!(
event.ops,
[
Op {
action: "create".into(),
collection: "dev.knotless.ref".into(),
rkey: "k.x".into(),
record: None
},
Op {
action: "delete".into(),
collection: "dev.knotless.issue".into(),
rkey: "3".into(),
record: None
},
]
);
let elsewhere = parse(&commit(vec![op("create", "app.bsky.feed.post/x")])).unwrap();
assert_eq!((elsewhere.did, elsewhere.ops.len()), (None, 0));
let identity = frame(
map(vec![("t", V::Text("#identity".into()))]),
map(vec![("did", V::Text("did:plc:b".into()))]),
);
let identity = parse(&identity).unwrap();
assert_eq!(
(identity.did.as_deref(), identity.identity),
(Some("did:plc:b"), true)
);
let info = frame(map(vec![("t", V::Text("#info".into()))]), map(vec![]));
assert_eq!(parse(&info), Some(Event::default()));
assert_eq!(parse(b"\xff"), None);
assert_eq!(parse(&frame(map(vec![]), map(vec![]))), None);
}
fn hex(s: &str) -> Vec<u8> {
(0..s.len())
.step_by(2)
.map(|i| u8::from_str_radix(&s[i..i + 2], 16).unwrap())
.collect()
}
/// A record and its CID as the reference implementation writes them,
/// with a blob in it and bytes.
const CID: &str = "01711220e257565846105b66bb046048fe489187b80bdd4ad18dc539082dcb8d5f670092";
const RECORD: &str = "a7616e03626f6bf5636c6f67a463726566d82a58250001551220836ff184e7b41b1e13cb5fd89fa1de98dbbab99e9d2918913ff43b86a5c7c2136473697a650365247479706564626c6f62686d696d65547970656a746578742f706c61696e6372617743010203646e6f6e65f66474616773816161652474797065726465762e6b6e6f746c6573732e636865636b";
#[test]
fn records_come_with_their_commits() {
let section = |bytes: &[u8]| {
let len = bytes.len();
let len = if len < 128 {
vec![len as u8]
} else {
vec![len as u8 | 0x80, (len >> 7) as u8]
};
[len, bytes.to_vec()].concat()
};
let (cid, written) = (hex(CID), hex(RECORD));
// A block cut short is left out, and a CAR cut short ends there.
let short = section(&cid[..20]);
let car = [
section(b"\xa0"),
short,
section(&[cid.clone(), written].concat()),
vec![60],
];
let link = |cid: &[u8]| V::Tag(42, Box::new(V::Bytes([&[0][..], cid].concat())));
let commit = |ops: Vec<V>| {
let body = map(vec![
("repo", V::Text("did:plc:a".into())),
("blocks", V::Bytes(car.concat())),
("ops", V::Array(ops)),
]);
parse(&frame(map(vec![("t", V::Text("#commit".into()))]), body)).unwrap()
};
let op = |path: &str, cid: V| {
map(vec![
("action", V::Text("create".into())),
("path", V::Text(path.into())),
("cid", cid),
])
};
let event = commit(vec![
op("dev.knotless.check/a", link(&cid)),
op("dev.knotless.check/b", link(b"elsewhere")),
op("dev.knotless.check/c", V::Null),
]);
let records: Vec<_> = event.ops.into_iter().map(|o| o.record).collect();
assert_eq!(
records,
[
Some(serde_json::json!({
"$type": "dev.knotless.check",
"n": 3,
"ok": true,
"none": null,
"tags": ["a"],
"raw": {"$bytes": "AQID"},
"log": {
"$type": "blob",
"ref": {"$link": "bafkreiedn7yyjz5udmpbhs273cp2dxuy3o5lthu5femjcp7uhodklr6ccm"},
"mimeType": "text/plain",
"size": 3,
},
})),
None,
None,
]
);
// What no record holds is fetched instead.
let cbor = |v: V| {
let mut out = Vec::new();
ciborium::into_writer(&v, &mut out).unwrap();
record(&out)
};
assert_eq!(cbor(V::Float(1.5)), None);
assert_eq!(cbor(V::Integer(u64::MAX.into())), None);
assert_eq!(cbor(V::Map(vec![(V::Integer(1.into()), V::Null)])), None);
assert_eq!(cbor(V::Tag(42, Box::new(V::Text("x".into())))), None);
assert_eq!(record(b"\xff"), None);
assert_eq!(varint(&mut &[0xff; 11][..]), None);
}
#[test]
fn endpoints() {
assert_eq!(
endpoint("wss://relay/", None),
"wss://relay/xrpc/com.atproto.sync.subscribeRepos"
);
assert_eq!(
endpoint("ws://h/xrpc/com.atproto.sync.subscribeRepos", Some(3)),
"ws://h/xrpc/com.atproto.sync.subscribeRepos?cursor=3"
);
}
}