use std::convert::Infallible;
use std::sync::LazyLock;
use std::time::Duration;
use axum::http::Uri;
use bytes::Bytes;
use futures::{SinkExt, StreamExt};
use jacquard_common::error::DecodeError;
use miette::Diagnostic;
use smol_str::format_smolstr;
use thiserror::Error;
use tokio_websockets::{ClientBuilder, Connector, Message as WsMsg, WebSocketStream};
use crate::net::{PublicEndpoint, PublicResolver};
use tracing::trace;
use url::Url;
pub mod codec;
pub mod types;
pub use codec::decode_frame;
#[allow(unused_imports)]
pub use types::{
Account, AccountStatus, Commit, Datetime, Identity, Info, InfoName, RepoOp, RepoOpAction,
SubscribeReposMessage, Sync,
};
#[cfg(feature = "relay")]
pub use codec::{encode_error_frame, encode_frame};
#[derive(Debug, Error, Diagnostic)]
pub enum FirehoseError {
#[error("websocket error: {0}")]
WebSocket(#[from] tokio_websockets::Error),
#[error("unknown scheme: {0}")]
UnknownScheme(String),
#[error("invalid websocket uri: {0}")]
InvalidUri(String),
#[error("decode error: {0}")]
Decode(#[from] DecodeError),
#[error("empty frame")]
EmptyFrame,
#[error("relay error {error}: {message:?}")]
RelayError {
error: String,
message: Option<String>,
},
#[error("unknown op: {0}")]
UnknownOp(i64),
#[error("missing type in header")]
MissingType,
#[error("unknown event type: {0}")]
UnknownType(String),
#[error("cbor decode error: {0}")]
Cbor(String),
#[error("stream closed: {code}: {reason}")]
StreamClosed { code: u16, reason: String },
#[error("tcp layer dropped")]
TcpDropped,
#[error("future cursor: {message:?}")]
FutureCursor { message: Option<String> },
#[error("connect timed out")]
ConnectTimeout,
#[error("host went silent: no frame or pong for {0:?}")]
Silent(Duration),
}
impl From<serde_ipld_dagcbor::DecodeError<Infallible>> for FirehoseError {
fn from(e: serde_ipld_dagcbor::DecodeError<Infallible>) -> Self {
Self::Cbor(e.to_string())
}
}
/// shared by every firehose connection: without one, tokio-websockets builds a
/// connector per connect, which re-reads and parses the system root store each
/// time and starves the runtime when thousands of PDS streams start at once.
/// if it can't be built, connects fall back to that path and report its error.
static TLS_CONNECTOR: LazyLock<Option<Connector>> = LazyLock::new(|| Connector::new().ok());
/// a host that takes the tcp connection and never answers the upgrade would park the task forever
const CONNECT_TIMEOUT: Duration = if cfg!(test) {
Duration::from_millis(200)
} else {
Duration::from_secs(30)
};
const PING_INTERVAL: Duration = if cfg!(test) {
Duration::from_millis(100)
} else {
Duration::from_secs(45)
};
/// two ping intervals with nothing heard means the connection is dead, because every host pongs
const SILENCE_DEADLINE: Duration = PING_INTERVAL.saturating_mul(2);
pub struct FirehoseStream {
ws: WebSocketStream<tokio_websockets::MaybeTlsStream<tokio::net::TcpStream>>,
ping_timer: tokio::time::Interval,
last_heard: tokio::time::Instant,
}
impl FirehoseStream {
pub async fn connect(
mut relay: Url,
cursor: Option<i64>,
public: bool,
) -> Result<Self, FirehoseError> {
if public {
let endpoint = PublicEndpoint::parse_websocket(&relay)
.map_err(|error| FirehoseError::InvalidUri(error.to_string()))?;
relay = endpoint.url().clone();
}
let scheme = match relay.scheme() {
"https" | "wss" => "wss",
"http" | "ws" => "ws",
x => return Err(FirehoseError::UnknownScheme(x.to_string())),
};
relay.set_scheme(scheme).expect("to be valid url");
relay.set_path("/xrpc/com.atproto.sync.subscribeRepos");
let cursor = cursor.map(|c| format_smolstr!("cursor={c}"));
relay.set_query(cursor.as_deref());
let uri: Uri = relay
.as_str()
.parse()
.map_err(|e| FirehoseError::InvalidUri(format!("{e}")))?;
let mut builder = ClientBuilder::from_uri(uri);
if scheme == "wss"
&& let Some(connector) = TLS_CONNECTOR.as_ref()
{
builder = builder.connector(connector);
}
let connect = async {
if public {
builder.resolver(PublicResolver::system()).connect().await
} else {
builder.connect().await
}
};
let (ws, _) = tokio::time::timeout(CONNECT_TIMEOUT, connect)
.await
.map_err(|_| FirehoseError::ConnectTimeout)??;
let mut ping_timer = tokio::time::interval(PING_INTERVAL);
ping_timer.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
ping_timer.tick().await; // consume the first tick
Ok(Self {
ws,
ping_timer,
last_heard: tokio::time::Instant::now(),
})
}
/// gets the next message bytes from the firehose
/// none means the stream is closed
pub async fn next(&mut self) -> Result<Bytes, FirehoseError> {
loop {
let res = tokio::select! {
_ = self.ping_timer.tick() => {
let silent = self.last_heard.elapsed();
if silent >= SILENCE_DEADLINE {
return Err(FirehoseError::Silent(silent));
}
self.ping().await?;
continue;
}
res = self.ws.next() => {
res.map(|m| m.map_err(Into::into))
.unwrap_or(Err(FirehoseError::TcpDropped))?
}
};
self.last_heard = tokio::time::Instant::now();
match res {
msg if msg.is_binary() => {
let bytes: Bytes = msg.into_payload().into();
if bytes.is_empty() {
return Err(FirehoseError::EmptyFrame);
}
return Ok(bytes);
}
// if ws closed treat it as an error, since why would a host close the stream??
// TODO: treat hosts that return these as offline ??????
msg if msg.is_close() => {
let (code, reason) = msg.as_close().map_or_else(
|| (0, "no reason".to_string()),
|(code, reason)| (code.into(), reason.to_owned()),
);
return Err(FirehoseError::StreamClosed { code, reason });
}
msg if msg.is_ping() => self.ws.send(WsMsg::pong(msg.into_payload())).await?,
msg if msg.is_pong() => continue,
x => {
trace!(msg = ?x, "host sent unexpected message");
continue;
}
}
}
}
async fn ping(&mut self) -> Result<(), FirehoseError> {
self.ws
.send(WsMsg::ping(Bytes::new()))
.await
.map_err(Into::into)
}
}
#[cfg(test)]
mod test {
#[cfg(feature = "relay")]
use super::FirehoseError;
#[cfg(feature = "relay")]
use super::types::{Datetime, RepoOp};
use super::{SubscribeReposMessage, decode_frame};
#[cfg(feature = "relay")]
use jacquard_common::types::{
cid::CidLink,
string::{Did, Tid},
};
#[cfg(feature = "relay")]
#[derive(serde::Serialize)]
#[serde(rename_all = "camelCase")]
struct CommitWithoutDeprecated<'a> {
blobs: Vec<CidLink>,
#[serde(with = "jacquard_common::serde_bytes_helper")]
blocks: bytes::Bytes,
commit: CidLink,
#[serde(borrow)]
ops: Vec<RepoOp<'a>>,
#[serde(skip_serializing_if = "Option::is_none")]
#[serde(default)]
prev_data: Option<CidLink>,
repo: Did,
rev: Tid,
seq: i64,
since: Option<Tid>,
time: Datetime,
}
#[test]
fn test_decode_account() {
const FRAME: &[u8] = b"omF0aCNhY2NvdW50Ym9wAaRjZGlkeCBkaWQ6cGxjOjNuNDNncWY2YTZua3J3MzU1cnNjNnJ0ZGNzZXEbAAAAAlP5Zt5kdGltZXgbMjAyNi0wMi0yNFQxNDowNToyMC41MjE1NDY5ZmFjdGl2ZfU=";
let bytes = data_encoding::BASE64.decode(FRAME).unwrap();
let SubscribeReposMessage::Account(account) = decode_frame(&bytes).unwrap() else {
panic!("expected account");
};
assert_eq!(account.did.as_str(), "did:plc:3n43gqf6a6nkrw355rsc6rtd");
assert!(account.active);
assert!(account.status.is_none());
assert_eq!(account.seq, 9_998_788_318);
assert_eq!(
account.time.0.to_rfc3339(),
"2026-02-24T14:05:20.521546900+00:00"
);
}
/// regression: some relays send `since: ""` (empty string) instead of null/absent for the initial commit.
/// this should decode cleanly with `since = None`.
/// TODO: is this behaviour we should reject?
#[test]
fn test_decode_commit_empty_since() {
const FRAME: &[u8] = b"omF0ZyNjb21taXRib3ABq2NvcHOAY3Jldm0zbWdkeWNmdWIyeTIyY3NlcRsAAAACZ8sHJmRyZXBveCBkaWQ6cGxjOmpxM3p2cmI1ZXdnMnFndXA3M3Fzb3V6ZWR0aW1leBsyMDI2LTAzLTA1VDE1OjQ3OjU0LjcxNjMyOFplYmxvYnOAZXNpbmNlYGZibG9ja3NZAUk6omVyb290c4HYKlglAAFxEiAlaO/rjabPL4/e2QlxkoxzCCwv69hE4P3Vdxpv7f6uEWd2ZXJzaW9uAeABAXESICVo7+uNps8vj97ZCXGSjHMILC/r2ETg/dV3Gm/t/q4RpmNkaWR4IGRpZDpwbGM6anEzenZyYjVld2cycWd1cDczcXNvdXplY3Jldm0zbWdkeWNmdWIyeTIyY3NpZ1hAwKfrZtwwbN7dW0uSbviOs65NWQRvlS9Qc7oRtiorybMTEYxKGJaFK2kHIMEWIJqumb4751En2aJEpsilWlaQOWRkYXRh2CpYJQABcRIgnf7+Yd126j3K5QI4gLCDedV63yBILW/b4nWSifZHZ3tkcHJldvZndmVyc2lvbgMrAXESIJ3+/mHdduo9yuUCOICwg3nVet8gSC1v2+J1kon2R2d7omFlgGFs9mZjb21taXTYKlglAAFxEiAlaO/rjabPL4/e2QlxkoxzCCwv69hE4P3Vdxpv7f6uEWZyZWJhc2X0ZnRvb0JpZ/Q=";
let bytes = data_encoding::BASE64.decode(FRAME).unwrap();
let msg = decode_frame(&bytes).unwrap();
let SubscribeReposMessage::Commit(c) = msg else {
panic!("expected Commit");
};
assert!(c.since.is_none(), "since should be None for empty string");
let stored = rmp_serde::to_vec_named(c.as_ref()).unwrap();
use sha2::Digest as _;
assert_eq!(
(stored.len(), crate::util::hex::encode(sha2::Sha256::digest(&stored)),),
(
523,
"cb3bb18e50d66f71c22b468f69df1b553922bb1073c2f7841fd3cae00317e260".to_owned(),
)
);
let replayed: super::types::Commit<'_> = rmp_serde::from_slice(&stored).unwrap();
assert!(replayed.since.is_none());
assert_eq!(replayed.repo, c.repo);
assert_eq!(replayed.rev, c.rev);
assert_eq!(replayed.blobs, c.blobs);
assert_eq!(replayed.blocks, c.blocks);
assert_eq!(replayed.commit, c.commit);
assert_eq!(replayed.ops.len(), c.ops.len());
for (actual, expected) in replayed.ops.iter().zip(&c.ops) {
assert_eq!(actual.action, expected.action);
assert_eq!(actual.cid, expected.cid);
assert_eq!(actual.path, expected.path);
assert_eq!(actual.prev, expected.prev);
}
assert_eq!(replayed.prev_data, c.prev_data);
assert_eq!(replayed.rebase, c.rebase);
assert_eq!(replayed.seq, c.seq);
assert_eq!(replayed.time.0, c.time.0);
assert_eq!(replayed.too_big, c.too_big);
}
#[cfg(feature = "relay")]
#[test]
fn test_decode_commit_missing_rebase_and_too_big_defaults_false() {
const FRAME: &[u8] = b"omF0ZyNjb21taXRib3ABq2NvcHOAY3Jldm0zbWdkeWNmdWIyeTIyY3NlcRsAAAACZ8sHJmRyZXBveCBkaWQ6cGxjOmpxM3p2cmI1ZXdnMnFndXA3M3Fzb3V6ZWR0aW1leBsyMDI2LTAzLTA1VDE1OjQ3OjU0LjcxNjMyOFplYmxvYnOAZXNpbmNlYGZibG9ja3NZAUk6omVyb290c4HYKlglAAFxEiAlaO/rjabPL4/e2QlxkoxzCCwv69hE4P3Vdxpv7f6uEWd2ZXJzaW9uAeABAXESICVo7+uNps8vj97ZCXGSjHMILC/r2ETg/dV3Gm/t/q4RpmNkaWR4IGRpZDpwbGM6anEzenZyYjVld2cycWd1cDczcXNvdXplY3Jldm0zbWdkeWNmdWIyeTIyY3NpZ1hAwKfrZtwwbN7dW0uSbviOs65NWQRvlS9Qc7oRtiorybMTEYxKGJaFK2kHIMEWIJqumb4751En2aJEpsilWlaQOWRkYXRh2CpYJQABcRIgnf7+Yd126j3K5QI4gLCDedV63yBILW/b4nWSifZHZ3tkcHJldvZndmVyc2lvbgMrAXESIJ3+/mHdduo9yuUCOICwg3nVet8gSC1v2+J1kon2R2d7omFlgGFs9mZjb21taXTYKlglAAFxEiAlaO/rjabPL4/e2QlxkoxzCCwv69hE4P3Vdxpv7f6uEWZyZWJhc2X0ZnRvb0JpZ/Q=";
let bytes = data_encoding::BASE64.decode(FRAME).unwrap();
let SubscribeReposMessage::Commit(c) = decode_frame(&bytes).unwrap() else {
panic!("expected Commit");
};
let frame = super::encode_frame(
"#commit",
&CommitWithoutDeprecated {
blobs: c.blobs.clone(),
blocks: c.blocks.clone(),
commit: c.commit.clone(),
ops: c.ops.clone(),
prev_data: c.prev_data.clone(),
repo: c.repo.clone(),
rev: c.rev,
seq: c.seq,
since: c.since,
time: c.time,
},
)
.unwrap();
let SubscribeReposMessage::Commit(c) = decode_frame(&frame).unwrap() else {
panic!("expected Commit");
};
assert!(!c.rebase, "missing rebase should default to false");
assert!(!c.too_big, "missing tooBig should default to false");
}
#[cfg(feature = "relay")]
#[test]
fn test_decode_encoded_error_frame() {
let bytes = super::encode_error_frame("ConsumerTooSlow", Some("blocked")).unwrap();
let Err(FirehoseError::RelayError { error, message }) = decode_frame(&bytes) else {
panic!("expected relay error");
};
assert_eq!(error, "ConsumerTooSlow");
assert_eq!(message.as_deref(), Some("blocked"));
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn public_connect_rejects_private_literal_before_socket_creation() {
let relay = Url::parse("wss://127.0.0.1/").unwrap();
let result = FirehoseStream::connect(relay, None, true).await;
assert!(
matches!(result, Err(FirehoseError::InvalidUri(message)) if message.contains("non-public IP"))
);
}
#[tokio::test]
async fn connect_gives_up_on_a_host_that_never_answers_the_upgrade() {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let url = Url::parse(&format!("ws://{}/", listener.local_addr().unwrap())).unwrap();
let held = tokio::spawn(async move {
let _conn = listener.accept().await;
std::future::pending::<()>().await
});
let res = FirehoseStream::connect(url, None, false).await;
assert!(matches!(res, Err(FirehoseError::ConnectTimeout)));
held.abort();
}
#[tokio::test]
async fn a_host_that_stops_answering_pings_is_dropped() {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let url = Url::parse(&format!("ws://{}/", listener.local_addr().unwrap())).unwrap();
// tokio-websockets only pongs while reading, and this server never reads
let held = tokio::spawn(async move {
let (tcp, _) = listener.accept().await.unwrap();
let _ws = tokio_websockets::ServerBuilder::new()
.accept(tcp)
.await
.unwrap();
std::future::pending::<()>().await
});
let mut stream = FirehoseStream::connect(url, None, false).await.unwrap();
let res = tokio::time::timeout(SILENCE_DEADLINE * 5, stream.next()).await;
assert!(matches!(res, Ok(Err(FirehoseError::Silent(_)))));
held.abort();
}
#[tokio::test]
async fn a_host_that_answers_pings_stays_connected() {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let url = Url::parse(&format!("ws://{}/", listener.local_addr().unwrap())).unwrap();
let held = tokio::spawn(async move {
let (tcp, _) = listener.accept().await.unwrap();
let (_, mut ws) = tokio_websockets::ServerBuilder::new()
.accept(tcp)
.await
.unwrap();
while ws.next().await.is_some() {}
});
let mut stream = FirehoseStream::connect(url, None, false).await.unwrap();
let res = tokio::time::timeout(SILENCE_DEADLINE * 5, stream.next()).await;
assert!(
res.is_err(),
"a quiet host that pongs isn't dropped: {res:?}"
);
held.abort();
}
}