diff --git a/README.md b/README.md index 9d88923..419c2ae 100644 --- a/README.md +++ b/README.md @@ -238,6 +238,7 @@ directory, it will also be loaded automatically. | `ENABLE_CRAWLER` | `true` if full network or crawler sources are configured, `false` otherwise | whether to actively query the network for unknown repositories. | | `CRAWLER_MAX_PENDING_REPOS` | `2000` | max pending repos for crawler. | | `CRAWLER_RESUME_PENDING_REPOS` | `1000` | resume threshold for crawler pending repos. | +| `NEW_HOST_LIMIT` | `50` | in relay mode, decides how many new hosts can be added via `com.atproto.sync.requestCrawl` in a day. | | `RATE_TIERS` | | comma-separated list of named rate tier definitions in `name:base/mul/hourly/daily[/account_limit]` format (e.g. `trusted:5000/10.0/18000000/432000000/10000000`). the optional account limit prevents new accounts from being created on this PDS once reached. built-in tiers (`default`, `trusted`) are always present and can be overridden. | | `TIER_RULES` | | comma-separated ordered list of glob rules in `pattern:tier_name` format (e.g. `*.bsky.network:trusted`). rules are evaluated in order; first match wins. explicit API assignments via `PUT /pds/tiers` take precedence over rules; the `default` tier is the final fallback. uses standard glob wildcards (`*`, `?`) matched against the PDS hostname. | diff --git a/src/api/xrpc/mod.rs b/src/api/xrpc/mod.rs index b96146d..84842d8 100644 --- a/src/api/xrpc/mod.rs +++ b/src/api/xrpc/mod.rs @@ -179,6 +179,23 @@ fn bad_request( } } +#[cfg(feature = "relay")] +fn rate_limited( + nsid: &'static str, + message: impl Display, +) -> XrpcErrorResponse { + XrpcErrorResponse { + status: StatusCode::TOO_MANY_REQUESTS, + error: XrpcError::Generic(GenericXrpcError { + error: "RateLimitExceeded".into(), + message: Some(message.to_smolstr()), + nsid, + method: "POST", + http_status: StatusCode::TOO_MANY_REQUESTS, + }), + } +} + #[cfg(feature = "indexer")] fn upstream_error( nsid: &'static str, diff --git a/src/api/xrpc/request_crawl.rs b/src/api/xrpc/request_crawl.rs index 3384c70..344a1fc 100644 --- a/src/api/xrpc/request_crawl.rs +++ b/src/api/xrpc/request_crawl.rs @@ -1,6 +1,7 @@ use jacquard_api::com_atproto::sync::request_crawl::{ RequestCrawlError, RequestCrawlRequest, RequestCrawlResponse, }; +use miette::IntoDiagnostic; use url::Url; use super::*; @@ -23,6 +24,30 @@ pub async fn handle( }); } + // enforce daily new pds limit on unknown hosts + if !hydrant.firehose.is_source_known(&url) { + let (allowed, to_persist) = hydrant.state.pds_daily_limit.try_increment(); + if !allowed { + return Err(rate_limited( + nsid, + "daily limit for new PDS sources reached", + )); + } + + // persist the new count before returning so a crash cannot reset the counter + // and allow the budget to be replayed. + if let Some((day, count)) = to_persist { + let state = hydrant.state.clone(); + tokio::task::spawn_blocking(move || { + crate::db::save_pds_daily_adds(&state.db, day, count) + }) + .await + .into_diagnostic() + .flatten() + .map_err(|e| internal_error(nsid, e))?; + } + } + hydrant .firehose .add_source(url, true) diff --git a/src/config.rs b/src/config.rs index 7383d58..787931d 100644 --- a/src/config.rs +++ b/src/config.rs @@ -375,6 +375,11 @@ pub struct Config { /// set via `HYDRANT_ONLY_INDEX_LINKS=true`. pub only_index_links: bool, + /// maximum number of new PDS sources that may be added (via seeding or API) in a single + /// UTC calendar day. `None` means unlimited. + /// set via `HYDRANT_NEW_HOST_LIMIT`. + pub new_host_limit: Option, + /// base URL(s) of relay or aggregator services to seed firehose PDS sources from at startup. /// /// hydrant calls `com.atproto.sync.listHosts` on each URL and adds the returned PDSes @@ -495,6 +500,7 @@ impl Default for Config { filter_excludes: None, enable_backlinks: false, only_index_links: false, + new_host_limit: Some(50), tier_rules: vec![], tier_policy: { let mut tiers = HashMap::new(); @@ -671,6 +677,9 @@ impl Config { let enable_backlinks: bool = cfg!("ENABLE_BACKLINKS", defaults.enable_backlinks); let only_index_links: bool = cfg!("ONLY_INDEX_LINKS", defaults.only_index_links); + let max_pds_added_per_day: Option = std::env::var("HYDRANT_NEW_HOST_LIMIT") + .ok() + .and_then(|s| s.parse().ok()); // start with built-in tier definitions, then layer in any env-defined overrides. // format: HYDRANT_RATE_TIERS=name:base/mul/hourly/daily,... @@ -785,6 +794,7 @@ impl Config { filter_excludes, enable_backlinks, only_index_links, + new_host_limit: max_pds_added_per_day, tier_policy, tier_rules, cache_size, @@ -915,6 +925,9 @@ impl fmt::Display for Config { ) )?; } + if let Some(limit) = self.new_host_limit { + config_line!(f, "max pds/day", limit)?; + } Ok(()) } } diff --git a/src/control/crawler.rs b/src/control/crawler.rs index 3b59b85..2d64852 100644 --- a/src/control/crawler.rs +++ b/src/control/crawler.rs @@ -32,9 +32,6 @@ pub(super) struct CrawlerShared { pub struct CrawlerSourceInfo { pub url: Url, pub mode: crate::config::CrawlerMode, - /// whether this source is persisted in the database (i.e. it was dynamically added - /// and will survive restarts). config-sourced entries have `persisted: false`. - pub persisted: bool, } pub(super) fn spawn_crawler_producer( @@ -179,7 +176,6 @@ impl CrawlerHandle { sources.push(CrawlerSourceInfo { url: url.clone(), mode: h.mode, - persisted: self.persisted.contains_sync(url), }); true }) diff --git a/src/control/firehose.rs b/src/control/firehose.rs index 8fd8da9..f39b700 100644 --- a/src/control/firehose.rs +++ b/src/control/firehose.rs @@ -34,8 +34,6 @@ pub(super) struct FirehoseShared { #[derive(Debug, Clone, serde::Serialize)] pub struct FirehoseSourceInfo { pub url: Url, - /// true if added via the API and persisted to the database; false for `RELAY_HOSTS` sources. - pub persisted: bool, /// true when this is a direct PDS connection; enables host authority enforcement. pub is_pds: bool, } @@ -48,8 +46,8 @@ pub struct FirehoseHandle { pub(super) shared: Arc>, /// per-relay running tasks, keyed by url. pub(super) tasks: Arc>, - /// set of urls persisted in the database (dynamically added sources). - pub(super) persisted: Arc>, + /// set of known source urls, includes API-added (db-persisted) and static config sources. + pub(super) known_sources: Arc>, /// ids assigned to spawned tasks next_task_id: Arc, } @@ -60,7 +58,7 @@ impl FirehoseHandle { state, shared: Arc::new(std::sync::OnceLock::new()), tasks: Arc::new(scc::HashMap::new()), - persisted: Arc::new(scc::HashSet::new()), + known_sources: Arc::new(scc::HashSet::new()), next_task_id: Arc::new(AtomicUsize::new(0)), } } @@ -153,6 +151,12 @@ impl FirehoseHandle { *self.state.firehose_enabled.borrow() } + /// returns `true` if this URL is already a known firehose source — either currently + /// running or persisted (e.g. the host is offline but was previously added). + pub fn is_source_known(&self, url: &Url) -> bool { + self.known_sources.contains_sync(url) + } + /// list all currently active firehose sources. pub async fn list_sources(&self) -> Vec { let mut out = Vec::new(); @@ -160,7 +164,6 @@ impl FirehoseHandle { .any_async(|url, handle| { out.push(FirehoseSourceInfo { url: url.clone(), - persisted: self.persisted.contains_sync(url), is_pds: handle.is_pds, }); false @@ -197,7 +200,7 @@ impl FirehoseHandle { .await .into_diagnostic()??; - let _ = self.persisted.insert_async(url.clone()).await; + let _ = self.known_sources.insert_async(url.clone()).await; // reset failure state so the fresh task gets a clean slate. // if the previous task exited after max failures, the failure counter @@ -217,7 +220,7 @@ impl FirehoseHandle { /// if the source was added via the API, it is removed from the database; /// if it came from the static config, only the running task is stopped. pub async fn remove_source(&self, url: &Url) -> Result { - if self.persisted.contains_async(url).await { + if self.known_sources.contains_async(url).await { let url_str = url.to_string(); tokio::task::spawn_blocking({ let state = self.state.clone(); @@ -232,7 +235,7 @@ impl FirehoseHandle { }) .await .into_diagnostic()??; - self.persisted.remove_async(url).await; + self.known_sources.remove_async(url).await; } Ok(self.tasks.remove_async(url).await.is_some()) diff --git a/src/control/mod.rs b/src/control/mod.rs index 4e78ee4..9648c62 100644 --- a/src/control/mod.rs +++ b/src/control/mod.rs @@ -414,6 +414,10 @@ impl Hydrant { "starting firehose ingestor(s)" ); for source in &relay_hosts { + let _ = firehose + .known_sources + .insert_async(source.url.clone()) + .await; firehose .spawn_firehose_ingestor(source, fire_shared, true) .await?; @@ -428,7 +432,10 @@ impl Hydrant { .into_diagnostic()??; for source in &persisted_sources { - let _ = firehose.persisted.insert_async(source.url.clone()).await; + let _ = firehose + .known_sources + .insert_async(source.url.clone()) + .await; if firehose.tasks.contains_async(&source.url).await { continue; } diff --git a/src/db/keys/mod.rs b/src/db/keys/mod.rs index aaf13dc..a98ed55 100644 --- a/src/db/keys/mod.rs +++ b/src/db/keys/mod.rs @@ -100,3 +100,8 @@ pub fn firehose_source_key(url: &str) -> Vec { pub fn pds_account_count_key(host: &str) -> String { format!("p|{host}") } + +/// key for the persisted daily-PDS-add counter in the cursors keyspace. +/// value layout: [day: u64 BE][count: u64 BE] = 16 bytes. +#[cfg(feature = "relay")] +pub const PDS_DAILY_ADDS_KEY: &[u8] = b"pds_daily_adds"; diff --git a/src/db/mod.rs b/src/db/mod.rs index 3fa7416..f983d49 100644 --- a/src/db/mod.rs +++ b/src/db/mod.rs @@ -836,6 +836,36 @@ pub fn persist_counts(db: &Db) -> Result<()> { batch.commit().into_diagnostic() } +/// load the persisted (day, count) pair for the daily PDS add counter, if present. +/// returns `None` if no entry exists or the stored data is malformed. +#[cfg(feature = "relay")] +pub fn load_pds_daily_adds(db: &Db) -> Result> { + let Some(val) = db.cursors.get(keys::PDS_DAILY_ADDS_KEY).into_diagnostic()? else { + return Ok(None); + }; + if val.len() < 16 { + miette::bail!("malformed pds daily limit value"); + } + let day = u64::from_be_bytes(val[..8].try_into().into_diagnostic()?); + let count = u64::from_be_bytes(val[8..].try_into().into_diagnostic()?); + Ok(Some((day, count))) +} + +/// persist the daily PDS add counter (day, count) to the cursors keyspace. +/// value layout: [day: u64 BE][count: u64 BE] = 16 bytes. +/// +/// takes the `cursors` keyspace directly so the caller can clone it into a +/// `spawn_blocking` closure without needing an owned `Db`. +#[cfg(feature = "relay")] +pub fn save_pds_daily_adds(db: &Db, day: u64, count: u64) -> Result<()> { + let mut value = [0u8; 16]; + value[..8].copy_from_slice(&day.to_be_bytes()); + value[8..].copy_from_slice(&count.to_be_bytes()); + db.cursors + .insert(keys::PDS_DAILY_ADDS_KEY, value) + .into_diagnostic() +} + pub fn load_persisted_firehose_sources( db: &crate::db::Db, ) -> Result> { diff --git a/src/lib.rs b/src/lib.rs index 29492b1..7bd8461 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -3,6 +3,8 @@ pub mod config; /// hydrant main api, includes the Hydrant type for programmatic control. pub mod control; pub(crate) mod filter; +#[cfg(feature = "relay")] +pub(crate) mod pds_daily_limit; pub(crate) mod pds_meta; pub mod types; diff --git a/src/pds_daily_limit.rs b/src/pds_daily_limit.rs new file mode 100644 index 0000000..045b33f --- /dev/null +++ b/src/pds_daily_limit.rs @@ -0,0 +1,73 @@ +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::{SystemTime, UNIX_EPOCH}; + +/// per-UTC-day counter for PDS additions via `requestCrawl`. +/// +/// the in-memory state is initialised from the database on startup (see +/// [`crate::db::load_pds_daily_adds`]). the counter resets automatically when the UTC day +/// rolls over. +pub(crate) struct PdsDailyLimit { + limit: Option, + /// current UTC day index (unix seconds / 86400). + day: AtomicU64, + /// requestCrawl calls accepted on the current UTC day. + count: AtomicU64, +} + +impl PdsDailyLimit { + /// construct from the previously-persisted `(day, count)` pair loaded from the database. + /// if the stored day doesn't match today the count is treated as 0. + pub(crate) fn new(limit: Option, stored: Option<(u64, u64)>) -> Self { + let today = utc_day(); + let count = stored + .filter(|(day, _)| *day == today) + .map(|(_, count)| count) + .unwrap_or(0); + Self { + limit, + day: AtomicU64::new(today), + count: AtomicU64::new(count), + } + } + + /// attempt to consume a daily slot. + /// + /// returns `(allowed, to_persist)`: + /// - `allowed`: whether the request is permitted. + /// - `to_persist`: when `Some((day, new_count))`, the caller must persist these values to + /// the database before returning success, so that a process crash cannot reset the counter + /// and allow the budget to be replayed. `None` when no limit is configured. + /// + /// when the UTC day rolls over the counter resets and a fresh quota starts. + pub(crate) fn try_increment(&self) -> (bool, Option<(u64, u64)>) { + let Some(limit) = self.limit else { + return (true, None); + }; + + let today = utc_day(); + if self.day.load(Ordering::Relaxed) != today { + self.count.store(0, Ordering::Relaxed); + self.day.store(today, Ordering::Relaxed); + } + + // fetch_add returns the value *before* the increment + let prev = self.count.fetch_add(1, Ordering::Relaxed); + if prev >= limit { + // undo to avoid the counter drifting upwards on repeated rejections + self.count.fetch_sub(1, Ordering::Relaxed); + return (false, None); + } + + let new_count = prev + 1; + let day = self.day.load(Ordering::Relaxed); + (true, Some((day, new_count))) + } +} + +fn utc_day() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs() + / 86400 +} diff --git a/src/state.rs b/src/state.rs index 796cbff..a8bbd46 100644 --- a/src/state.rs +++ b/src/state.rs @@ -10,6 +10,8 @@ use tokio::sync::Notify; use tokio::sync::watch; use url::Url; +#[cfg(feature = "relay")] +use crate::pds_daily_limit::PdsDailyLimit; use crate::{ config::Config, db::Db, @@ -22,21 +24,23 @@ use crate::{ pub struct AppState { pub db: Db, pub resolver: Resolver, + pub throttler: Throttler, pub(crate) filter: FilterHandle, pub(crate) pds_meta: PdsMetaHandle, + #[cfg(feature = "relay")] + pub(crate) pds_daily_limit: PdsDailyLimit, pub(crate) tier_policy: TierPolicy, pub firehose_cursors: scc::HashIndex, + pub firehose_enabled: watch::Sender, #[cfg(feature = "indexer")] pub backfill_notify: Notify, #[cfg(feature = "indexer")] pub crawler_enabled: watch::Sender, - pub firehose_enabled: watch::Sender, #[cfg(feature = "indexer")] pub backfill_enabled: watch::Sender, pub ephemeral: bool, pub ephemeral_ttl: Duration, pub only_index_links: bool, - pub throttler: Throttler, } impl AppState { @@ -87,6 +91,9 @@ impl AppState { let pds_meta = new_pds_handle(PdsMeta { hosts }); let relay_cursors = scc::HashIndex::new(); + #[cfg(feature = "relay")] + let pds_daily_limit = + PdsDailyLimit::new(config.new_host_limit, crate::db::load_pds_daily_adds(&db)?); #[cfg(feature = "indexer")] let (crawler_enabled, _) = watch::channel(crawler_default); @@ -112,6 +119,8 @@ impl AppState { ephemeral_ttl: config.ephemeral_ttl.clone(), only_index_links: config.only_index_links, throttler: Throttler::new(), + #[cfg(feature = "relay")] + pds_daily_limit, }) } diff --git a/tests/api.nu b/tests/api.nu index f90ff0c..7593fc1 100644 --- a/tests/api.nu +++ b/tests/api.nu @@ -30,10 +30,7 @@ def test-crawler-sources [url: string, pid: int] { if $s.mode != "list_repos" { fail $"expected mode=list_repos, got ($s.mode)" $pid } - if not $s.persisted { - fail "expected persisted=true for dynamically added source" $pid - } - print $" ok: 1 source, url=($s.url), mode=($s.mode), persisted=($s.persisted)" + print $" ok: 1 source, url=($s.url), mode=($s.mode)" # posting the same URL with a different mode replaces the existing entry print " POST /crawler/sources (should override)..." @@ -113,9 +110,6 @@ def test-source-persistence [binary: string, db_path: string, port: int] { fail $"expected 1 source after restart, got ($after | length)" $instance2.pid } let s = ($after | first) - if not $s.persisted { - fail "expected persisted=true after restart" $instance2.pid - } if $s.mode != "by_collection" { fail $"expected mode=by_collection after restart, got ($s.mode)" $instance2.pid } @@ -141,16 +135,13 @@ def test-config-source-not-persisted [binary: string, db_path: string, port: int fail "hydrant did not start" } - # config source should appear, but with persisted=false - print " checking config source appears with persisted=false..." + # config source should appear + print " checking config source appears..." let sources = (http get $"($url)/crawler/sources") if ($sources | length) != 1 { fail $"expected 1 source, got ($sources | length)" $instance.pid } - if ($sources | first).persisted { - fail "expected persisted=false for a CRAWLER_URLS source" $instance.pid - } - print " ok: config source has persisted=false" + print " ok: config source present" # the task can be stopped at runtime print " deleting config source at runtime..." @@ -179,9 +170,6 @@ def test-config-source-not-persisted [binary: string, db_path: string, port: int if ($after_restart | length) != 1 { fail $"expected config source to reappear after restart, got ($after_restart | length)" $instance2.pid } - if ($after_restart | first).persisted { - fail "expected persisted=false after restart" $instance2.pid - } print " ok: config source reappears on restart (not persisted to DB)" kill $instance2.pid @@ -213,10 +201,7 @@ def test-firehose-sources [url: string, pid: int] { fail $"expected 1 source, got ($sources | length)" $pid } let s = ($sources | first) - if not $s.persisted { - fail "expected persisted=true for dynamically added source" $pid - } - print $" ok: 1 source, url=($s.url), persisted=($s.persisted)" + print $" ok: 1 source, url=($s.url)" # posting the same URL replaces the existing entry print " POST /firehose/sources (should override)..."