Something went wrong. Try again.
very fast at protocol indexer with flexible filtering, xrpc queries, cursor-backed event stream, and more, built on fjall
rust fjall at-protocol atproto indexer
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650use miette::{IntoDiagnostic, Result, WrapErr};use smol_str::{SmolStr, ToSmolStr};use std::collections::BTreeSet;use std::net::SocketAddr;use std::path::PathBuf;use std::str::FromStr;use url::Url;
use super::Hydrant;use crate::db::keys;use crate::pds_meta::HostStatus;
/// unix sockets are owner-only unless told otherwise, since the api has no auth of its own.pub const DEFAULT_UNIX_SOCKET_MODE: u32 = 0o600;
/// one address the api listens on.#[derive(Debug, Clone, PartialEq, Eq)]pub enum ApiBind { Tcp(SocketAddr), Unix(PathBuf),}
impl ApiBind { pub fn tcp(&self) -> Option<SocketAddr> { match self { Self::Tcp(addr) => Some(*addr), Self::Unix(_) => None, } }}
impl From<SocketAddr> for ApiBind { fn from(addr: SocketAddr) -> Self { Self::Tcp(addr) }}
impl FromStr for ApiBind { type Err = miette::Report;
/// parses `host:port`, or `unix:` followed by a socket path. fn from_str(s: &str) -> Result<Self> { match s.strip_prefix("unix:") { Some("") => miette::bail!("`unix:` bind is missing a socket path"), Some(path) => Ok(Self::Unix(path.into())), None => s .parse() .map(Self::Tcp) .map_err(|e| miette::miette!("invalid bind address `{s}`: {e}")), } }}
impl std::fmt::Display for ApiBind { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { Self::Tcp(addr) => write!(f, "{addr}"), Self::Unix(path) => write!(f, "unix:{}", path.display()), } }}
#[derive(Debug, Clone)]pub struct ApiBinds { pub(crate) head: ApiBind, pub(crate) tail: Vec<ApiBind>, pub(crate) unix_mode: u32, pub(crate) tcp_management: bool,}
impl ApiBinds { pub fn new(bind: impl Into<ApiBind>) -> Self { Self { head: bind.into(), tail: Vec::new(), unix_mode: DEFAULT_UNIX_SOCKET_MODE, tcp_management: false, } }
pub fn try_from_iter<I>(iter: I) -> Option<Self> where I: IntoIterator, I::Item: Into<ApiBind>, { let mut iter = iter.into_iter().map(Into::into); let head = iter.next()?; Some(Self { head, tail: iter.collect(), unix_mode: DEFAULT_UNIX_SOCKET_MODE, tcp_management: false, }) }
/// sets the permissions unix sockets are created with. only the `0o777` bits are used. pub fn with_unix_mode(mut self, mode: u32) -> Self { self.unix_mode = mode & 0o777; self }
/// serves the management endpoints on tcp binds too. they're unix-only by default, since /// they have no auth and a unix socket's permissions are the only thing guarding them. pub fn with_tcp_management(mut self, enabled: bool) -> Self { self.tcp_management = enabled; self }
/// whether any bind serves the management endpoints. pub fn serves_management(&self) -> bool { self.tcp_management || self.iter().any(|bind| matches!(bind, ApiBind::Unix(_))) }
pub fn iter(&self) -> impl Iterator<Item = &ApiBind> + '_ { std::iter::once(&self.head).chain(self.tail.iter()) }
/// reads `HYDRANT_API_BIND`, `HYDRANT_API_SOCKET_MODE` and `HYDRANT_API_TCP_MANAGEMENT` /// through `lookup`. an unset bind gives `default`, and `none` gives `None`, meaning no api /// at all. pub fn from_lookup( lookup: impl Fn(&str) -> Option<String>, default: Option<ApiBinds>, ) -> Result<Option<Self>> { let binds = match lookup("HYDRANT_API_BIND") { Some(raw) if raw.trim().eq_ignore_ascii_case("none") => None, Some(raw) => Some( raw.parse::<Self>() .map_err(|e| miette::miette!("invalid HYDRANT_API_BIND: {e}"))?, ), None => default, }; let Some(mut binds) = binds else { return Ok(None); }; if let Some(mode) = lookup("HYDRANT_API_SOCKET_MODE") { binds = binds.with_unix_mode(parse_socket_mode(&mode)?); } let tcp_management = crate::config::env::read(&lookup, "HYDRANT_API_TCP_MANAGEMENT", binds.tcp_management)?; Ok(Some(binds.with_tcp_management(tcp_management))) }}
/// octal permission bits like `0660`, with or without the leading `0` or `0o`.fn parse_socket_mode(s: &str) -> Result<u32> { let digits = s.trim(); let digits = digits.strip_prefix("0o").unwrap_or(digits); u32::from_str_radix(digits, 8) .ok() .filter(|mode| !digits.starts_with('+') && *mode <= 0o777) .ok_or_else(|| { miette::miette!( "invalid HYDRANT_API_SOCKET_MODE `{s}`: want octal permissions like 0660" ) })}
impl FromStr for ApiBinds { type Err = miette::Report;
/// parses a comma-separated list of [`ApiBind`]s, the format of `HYDRANT_API_BIND`. fn from_str(s: &str) -> Result<Self> { let binds = s .split(',') .map(str::trim) .filter(|s| !s.is_empty()) .map(str::parse::<ApiBind>) .collect::<Result<Vec<_>>>()?; Self::try_from_iter(binds).ok_or_else(|| miette::miette!("no addresses to bind to")) }}
#[derive(Debug, Clone)]/// information about a host hydrant is consuming from.pub struct Host { /// hostname of the host. pub name: SmolStr, /// latest seq hydrant has processed from this host. pub seq: i64, /// the amount of accounts hydrant has seen from this host. pub account_count: u64, /// the status of this host in hydrant. pub status: crate::pds_meta::HostStatus,}
impl Hydrant { /// get the status of a (firehose) host we are consuming from. /// /// returns the seq we are on for this host. pub async fn get_host_status(&self, hostname: &str) -> Result<Option<Host>> { let state = self.state.clone(); let hostname = hostname.to_smolstr();
let state_closure = state.clone(); state .db .run(move |db| { let key = keys::firehose_cursor_key(&hostname);
let mut seq = 0; if let Some(cursor_bytes) = db.cursors.get(&key).into_diagnostic()? { seq = i64::from_be_bytes(cursor_bytes.as_ref().try_into().into_diagnostic()?); } else { // if it has no cursor, check if it's explicitly tracked in hosts map // or firehose tasks (recently added via API but no messages yet) let meta = state_closure.pds_meta.load(); if !meta.hosts.contains_key(hostname.as_str()) { // we should also allow it if it's an active firehose ingestor let mut found_in_cursors = false; state_closure.firehose_cursors.iter_sync(|u, _| { if u.host_str() == Some(hostname.as_str()) { found_in_cursors = true; } !found_in_cursors // continue if not found });
if !found_in_cursors { return Ok(None); } } }
let account_count = db.get_count_sync(&keys::pds_account_count_key(&hostname)); let status = state_closure.pds_meta.load().status(&hostname); Ok(Some(Host { name: hostname, seq, account_count, status, })) }) .await }
/// enumerates all hosts hydrant is consuming from. /// /// returns hosts enumerated in this pagination and the cursor to paginate from. pub async fn list_hosts( &self, cursor: Option<&str>, limit: usize, ) -> Result<(Vec<Host>, Option<SmolStr>)> { let state = self.state.clone(); let cursor = cursor.map(str::to_string); let state_closure = state.clone();
state .db .run(move |db| { let start_bound = match &cursor { Some(after) => std::ops::Bound::Included(keys::firehose_cursor_key(after)), None => std::ops::Bound::Included(keys::FIREHOSE_CURSOR_PREFIX.to_vec()), };
let prefix_end = { let mut end = keys::FIREHOSE_CURSOR_PREFIX.to_vec(); *end.last_mut().unwrap() += 1; end }; let end_bound = std::ops::Bound::Excluded(prefix_end);
let mut db_hosts = Vec::new(); for item in db.cursors.range((start_bound, end_bound)) { let (k, _) = item.into_inner().into_diagnostic()?; let hostname = std::str::from_utf8(&k[keys::FIREHOSE_CURSOR_PREFIX.len()..]) .into_diagnostic() .wrap_err("firehose cursor key contains non-utf8 hostname")?;
if cursor .as_ref() .is_some_and(|after| hostname <= after.as_str()) { continue; }
db_hosts.push(SmolStr::new(hostname)); if db_hosts.len() > limit { break; } }
let mut meta_hosts = Vec::new(); { let meta = state_closure.pds_meta.load(); for hostname in meta.hosts.keys() { if cursor .as_ref() .is_some_and(|after| hostname.as_str() <= after.as_str()) { continue; } meta_hosts.push(SmolStr::new(hostname)); } } meta_hosts.sort(); meta_hosts.truncate(limit + 1);
let mut merged = BTreeSet::new(); for host in db_hosts { merged.insert(host); } for host in meta_hosts { merged.insert(host); }
let selected: Vec<SmolStr> = merged.into_iter().take(limit + 1).collect();
let mut hosts: Vec<Host> = Vec::with_capacity(selected.len().min(limit)); for hostname in selected.iter().take(limit) { let seq = db .cursors .get(keys::firehose_cursor_key(hostname)) .into_diagnostic()? .map(|v| { v.as_ref() .try_into() .into_diagnostic() .wrap_err("cursor value is not 8 bytes") .map(i64::from_be_bytes) }) .transpose()? .unwrap_or(0); let account_count = db.get_count_sync(&keys::pds_account_count_key(hostname)); let status = state_closure.pds_meta.load().status(hostname); hosts.push(Host { name: hostname.clone(), seq, account_count, status, }); }
let next_cursor = if selected.len() > limit { hosts.last().map(|h| h.name.clone()) } else { None };
Ok((hosts, next_cursor)) }) .await }}
#[cfg(test)]mod api_bind_tests { use super::*;
#[test] fn parses_tcp_and_unix_binds() { assert_eq!( "127.0.0.1:3000".parse::<ApiBind>().unwrap(), ApiBind::Tcp("127.0.0.1:3000".parse().unwrap()) ); assert_eq!( "[::1]:3000".parse::<ApiBind>().unwrap(), ApiBind::Tcp("[::1]:3000".parse().unwrap()) ); assert_eq!( "unix:/run/hydrant/api.sock".parse::<ApiBind>().unwrap(), ApiBind::Unix("/run/hydrant/api.sock".into()) ); assert!("unix:".parse::<ApiBind>().is_err()); assert!("localhost:3000".parse::<ApiBind>().is_err()); assert!("/run/hydrant/api.sock".parse::<ApiBind>().is_err()); }
#[test] fn display_round_trips() { for s in ["127.0.0.1:3000", "[::1]:3000", "unix:/run/hydrant/api.sock"] { assert_eq!(s.parse::<ApiBind>().unwrap().to_string(), s); } }
#[test] fn parses_bind_lists() { let binds: ApiBinds = " 127.0.0.1:3000, unix:/tmp/api.sock ,".parse().unwrap(); assert_eq!( binds.iter().cloned().collect::<Vec<_>>(), [ ApiBind::Tcp("127.0.0.1:3000".parse().unwrap()), ApiBind::Unix("/tmp/api.sock".into()), ] ); assert!("".parse::<ApiBinds>().is_err()); assert!(" , ".parse::<ApiBinds>().is_err()); assert!("127.0.0.1:3000,nope".parse::<ApiBinds>().is_err()); }
fn lookup<'a>(pairs: &'a [(&str, &str)]) -> impl Fn(&str) -> Option<String> + 'a { |key| { pairs .iter() .find(|(k, _)| *k == key) .map(|(_, v)| v.to_string()) } }
#[test] fn from_lookup_reads_bind_and_socket_mode() { let default = ApiBinds::new(ApiBind::Tcp("127.0.0.1:3000".parse().unwrap())); let unset = ApiBinds::from_lookup(lookup(&[]), Some(default.clone())).unwrap(); assert_eq!(unset.unwrap().iter().count(), 1); assert!(ApiBinds::from_lookup(lookup(&[]), None).unwrap().is_none());
for none in ["none", " NONE "] { let binds = ApiBinds::from_lookup(lookup(&[("HYDRANT_API_BIND", none)]), Some(default.clone())); assert!(binds.unwrap().is_none(), "{none:?} should turn the api off"); }
let unix = ApiBinds::from_lookup(lookup(&[("HYDRANT_API_BIND", "unix:/tmp/api.sock")]), None); assert_eq!(unix.unwrap().unwrap().unix_mode, 0o600);
for (mode, want) in [ ("0660", 0o660), ("660", 0o660), ("0o640", 0o640), ("0777", 0o777), ] { let pairs = [ ("HYDRANT_API_BIND", "unix:/tmp/api.sock"), ("HYDRANT_API_SOCKET_MODE", mode), ]; let binds = ApiBinds::from_lookup(lookup(&pairs), None) .unwrap() .unwrap(); assert_eq!(binds.unix_mode, want, "mode {mode}"); } }
#[test] fn management_is_unix_only_unless_enabled_for_tcp() { let tcp = ApiBinds::new(ApiBind::Tcp("127.0.0.1:3000".parse().unwrap())); assert!(!tcp.tcp_management); assert!(!tcp.serves_management()); assert!(tcp.clone().with_tcp_management(true).serves_management()); let both: ApiBinds = "127.0.0.1:3000,unix:/tmp/api.sock".parse().unwrap(); assert!(both.serves_management());
// the flag applies to the default binds too, not just ones set in HYDRANT_API_BIND let enabled = ApiBinds::from_lookup( lookup(&[("HYDRANT_API_TCP_MANAGEMENT", "true")]), Some(tcp.clone()), ); assert!(enabled.unwrap().unwrap().tcp_management); let disabled = ApiBinds::from_lookup( lookup(&[("HYDRANT_API_TCP_MANAGEMENT", "false")]), Some(tcp), ); assert!(!disabled.unwrap().unwrap().tcp_management); }
#[test] fn from_lookup_refuses_bad_values() { let sock = ("HYDRANT_API_BIND", "unix:/tmp/api.sock"); for bad in [ &[("HYDRANT_API_BIND", "")][..], &[("HYDRANT_API_BIND", "none,unix:/tmp/api.sock")], &[sock, ("HYDRANT_API_SOCKET_MODE", "0800")], &[sock, ("HYDRANT_API_SOCKET_MODE", "1777")], &[sock, ("HYDRANT_API_SOCKET_MODE", "+660")], &[sock, ("HYDRANT_API_SOCKET_MODE", "rw")], &[sock, ("HYDRANT_API_TCP_MANAGEMENT", "yes")], ] { assert!( ApiBinds::from_lookup(lookup(bad), None).is_err(), "{bad:?} should be refused" ); } }
#[test] fn with_unix_mode_keeps_only_permission_bits() { let binds = ApiBinds::new(ApiBind::Unix("/tmp/api.sock".into())).with_unix_mode(0o4660); assert_eq!(binds.unix_mode, 0o660); }
#[test] fn binds_can_be_unix_only() { let binds = ApiBinds::new(ApiBind::Unix("/tmp/api.sock".into())); assert_eq!(binds.iter().filter_map(ApiBind::tcp).count(), 0); assert_eq!(binds.iter().count(), 1); }}
#[cfg(test)]mod host_listing_tests { use super::*; use crate::config::Config; use crate::db::set_ks_count; use crate::pds_meta::HostStatus;
fn test_config(path: &std::path::Path) -> Config { Config { database_path: path.to_path_buf(), ..Default::default() } }
#[tokio::test] async fn list_hosts_includes_seeded_hosts_without_cursors() -> Result<()> { let tmp = tempfile::tempdir().into_diagnostic()?; let hydrant = Hydrant::new(test_config(tmp.path())).await?;
{ let state = hydrant.state.clone(); state .db .run(move |db| -> Result<()> { let mut batch = db.inner.batch(); crate::db::pds_meta::set_status( &mut batch, &db.filter, "offline.example", HostStatus::Offline, )?; crate::db::pds_meta::set_status( &mut batch, &db.filter, "active.example", HostStatus::Active, )?; set_ks_count( &mut batch, db, &keys::pds_account_count_key("offline.example"), 7, ); set_ks_count( &mut batch, db, &keys::pds_account_count_key("active.example"), 42, ); db.cursors .insert( keys::firehose_cursor_key("active.example"), 123_i64.to_be_bytes(), ) .into_diagnostic()?; batch.commit().into_diagnostic()?; db.persist() }) .await?;
crate::pds_meta::PdsMeta::update_host( &hydrant.state.pds_meta, "offline.example", |h| { h.status = HostStatus::Offline; }, ); crate::pds_meta::PdsMeta::update_host(&hydrant.state.pds_meta, "active.example", |h| { h.status = HostStatus::Active; }); hydrant.state.db.update_count("p|offline.example", 7); hydrant.state.db.update_count("p|active.example", 42); }
let (hosts, next) = hydrant.list_hosts(None, 10).await?; assert!(next.is_none());
let offline = hosts .iter() .find(|h| h.name == "offline.example") .expect("seeded host without cursor should be listed"); let active = hosts .iter() .find(|h| h.name == "active.example") .expect("host with cursor should be listed");
assert_eq!(offline.seq, 0); assert_eq!(offline.account_count, 7); assert_eq!(offline.status, HostStatus::Offline);
assert_eq!(active.seq, 123); assert_eq!(active.account_count, 42); assert_eq!(active.status, HostStatus::Active);
let status = hydrant.get_host_status("active.example").await?; assert_eq!(status.map(|host| host.seq), Some(123));
Ok(()) }
#[tokio::test] async fn list_hosts_pagination_correctness() -> Result<()> { let tmp = tempfile::tempdir().into_diagnostic()?; let hydrant = Hydrant::new(test_config(tmp.path())).await?;
// Seed some in DB: host2, host4, host6 { let state = hydrant.state.clone(); state .db .run(move |db| -> Result<()> { db.cursors .insert(keys::firehose_cursor_key("host2"), 2_i64.to_be_bytes()) .into_diagnostic()?; db.cursors .insert(keys::firehose_cursor_key("host4"), 4_i64.to_be_bytes()) .into_diagnostic()?; db.cursors .insert(keys::firehose_cursor_key("host6"), 6_i64.to_be_bytes()) .into_diagnostic()?; db.persist() }) .await?; }
// Seed some in memory: host1, host3, host5 crate::pds_meta::PdsMeta::update_host(&hydrant.state.pds_meta, "host1", |h| { h.status = HostStatus::Active; }); crate::pds_meta::PdsMeta::update_host(&hydrant.state.pds_meta, "host3", |h| { h.status = HostStatus::Active; }); crate::pds_meta::PdsMeta::update_host(&hydrant.state.pds_meta, "host5", |h| { h.status = HostStatus::Active; });
// Lexicographically sorted: host1, host2, host3, host4, host5, host6 // Page 1: limit = 2 let (hosts1, next1) = hydrant.list_hosts(None, 2).await?; assert_eq!(hosts1.len(), 2); assert_eq!(hosts1[0].name.as_str(), "host1"); assert_eq!(hosts1[1].name.as_str(), "host2"); assert_eq!(next1.as_deref(), Some("host2"));
// Page 2: limit = 2, cursor = host2 let (hosts2, next2) = hydrant.list_hosts(next1.as_deref(), 2).await?; assert_eq!(hosts2.len(), 2); assert_eq!(hosts2[0].name.as_str(), "host3"); assert_eq!(hosts2[1].name.as_str(), "host4"); assert_eq!(next2.as_deref(), Some("host4"));
// Page 3: limit = 2, cursor = host4 let (hosts3, next3) = hydrant.list_hosts(next2.as_deref(), 2).await?; assert_eq!(hosts3.len(), 2); assert_eq!(hosts3[0].name.as_str(), "host5"); assert_eq!(hosts3[1].name.as_str(), "host6"); assert!(next3.is_none());
Ok(()) }}