use std::collections::HashMap;
use std::sync::Arc;
use std::time::Instant;
use fjall::OwnedWriteBatch;
use jacquard_common::IntoStatic;
use jacquard_common::types::crypto::PublicKey;
use jacquard_common::types::did::Did;
use miette::{IntoDiagnostic, Result};
use smol_str::{SmolStr, ToSmolStr};
use tracing::{debug, trace, warn};
use url::Url;
use crate::db::types::DidKey;
use crate::db::{self, CountDeltas, keys};
use crate::net::PublicHttpClient;
use crate::resolver::NoSigningKeyError;
use crate::state::AppState;
use crate::types::{RepoState, RepoStatus};
use crate::util;
use super::lookup::{AnswerTx, LookupError, LookupKind, Priority, Request, Resolved, Submitted};
use super::{
AuthorityOutcome, WRONG_HOST_AUTHORITY_CACHE_PRUNE_AT, WRONG_HOST_AUTHORITY_RECHECK_INTERVAL,
WorkerMessage,
};
use crate::ingest::stream::AccountStatus;
#[cfg(feature = "indexer")]
use crate::ingest::stream::SubscribeReposMessage;
use crate::ingest::validation::{
CommitValidationError, SyncValidationError, ValidatedCommit, ValidatedSync, ValidationContext,
};
use super::{commit_validation_outcome, sync_validation_outcome};
use crate::ingest::firehose_stats::StatsInstant;
/// how long a doc fetch vouches for the key it put on the row, like indigo's identity cache
pub(crate) const IDENTITY_TTL: i64 = 24 * 60 * 60;
/// the message that found a did's identity due for a refresh, so the refresh can be announced
/// in its place if it moves the identity
#[cfg_attr(
not(any(feature = "indexer_stream", feature = "relay")),
allow(dead_code)
)]
pub(crate) struct RefreshOrigin {
pub(crate) firehose: Url,
pub(crate) seq: i64,
pub(crate) time: crate::ingest::stream::Datetime,
}
/// what a step found out, or that it needs a lookup first. a step that waits has staged nothing
/// durable, and the whole message runs again once the answer is in
pub(crate) enum Checked<T> {
Done(T),
Wait(LookupKind),
}
/// what the shard knows about a did's identity beyond its row
pub(crate) enum Fact<'r> {
/// a lookup for it just answered, and this message is being replayed with the answer
Resolved(&'r Resolved, i64),
/// its lookup just failed, or failed recently enough that asking again is pointless
Failed,
/// nobody asked yet
Unknown,
}
/// a lookup answer the shard is replaying a did's held messages with
pub(crate) struct Answered {
pub(crate) did: Did,
/// the lookup the first held message waited for
pub(crate) kind: LookupKind,
pub(crate) result: Result<Resolved, LookupError>,
pub(crate) resolved_at: i64,
}
#[derive(Default)]
pub(crate) struct IdentityHasher(u64);
impl std::hash::Hasher for IdentityHasher {
fn write(&mut self, _: &[u8]) {
unreachable!("IdentityHasher only hashes u64 keys")
}
fn write_u64(&mut self, i: u64) {
self.0 = i;
}
fn finish(&self) -> u64 {
self.0
}
}
pub(crate) struct WorkerContext<'a> {
pub(crate) verify_signatures: bool,
pub(crate) state: &'a AppState,
pub(crate) vctx: ValidationContext<'a>,
pub(crate) stats: Arc<crate::ingest::firehose_stats::RelayShardStats>,
pub(crate) batch: OwnedWriteBatch,
pub(crate) count_deltas: CountDeltas,
pub(crate) sink: super::sink::EventSink,
#[cfg_attr(all(feature = "relay", not(feature = "indexer")), allow(dead_code))]
pub(crate) http: PublicHttpClient,
pub(crate) error_counts: HashMap<u64, u32, std::hash::BuildHasherDefault<IdentityHasher>>,
pub(crate) wrong_host_authority:
HashMap<u64, Instant, std::hash::BuildHasherDefault<IdentityHasher>>,
/// the repo row as [`Self::load_repo_state`] read it, `None` when it made the row up
pub(crate) loaded: Option<RepoState<'static>>,
/// the repo row this message writes, which the sink stages when the message commits
pub(crate) staged: Option<RepoState<'static>>,
/// where the lookup pool answers this shard's jobs
pub(crate) answers: AnswerTx,
/// see [`RepoState::identity_resolved_or_spread`]
pub(crate) identity_epoch: i64,
pub(crate) refresh_origins: HashMap<Did, RefreshOrigin>,
/// set while a did's held messages replay
pub(crate) answered: Option<Answered>,
pub(crate) parked: super::parked::Parked,
}
impl WorkerContext<'_> {
/// increments host error counter, returns if host is suppressed or not
pub(crate) fn inc_error(&mut self, host: &str) -> bool {
let error_count = self.error_counts.entry(util::hash(&host)).or_default();
let is_suppressed = *error_count > 50;
*error_count += 1;
is_suppressed
}
/// forgets what the current message staged, because it waits for a lookup and runs again
pub(crate) fn discard(&mut self) {
self.batch = self.state.db.inner.batch();
self.count_deltas = CountDeltas::default();
self.loaded = None;
self.staged = None;
self.sink.discard();
}
pub(crate) fn stage_repo_state(&mut self, repo_state: RepoState<'static>) {
self.staged = Some(repo_state);
}
pub(crate) fn reset_error(&mut self, host: &str) {
if let Some(count) = self.error_counts.get_mut(&util::hash(&host)) {
*count = 0;
}
}
fn wrong_host_authority_key(did: &Did, source_host: &str) -> u64 {
util::hash(&(did.as_str(), source_host))
}
fn recently_rejected_wrong_host_authority(&mut self, cache_key: u64) -> bool {
if self
.wrong_host_authority
.get(&cache_key)
.is_some_and(|checked_at| checked_at.elapsed() < WRONG_HOST_AUTHORITY_RECHECK_INTERVAL)
{
return true;
}
self.wrong_host_authority.remove(&cache_key);
false
}
fn remember_wrong_host_authority(&mut self, cache_key: u64) {
if self.wrong_host_authority.len() >= WRONG_HOST_AUTHORITY_CACHE_PRUNE_AT {
self.wrong_host_authority.retain(|_, checked_at| {
checked_at.elapsed() < WRONG_HOST_AUTHORITY_RECHECK_INTERVAL
});
if self.wrong_host_authority.len() >= WRONG_HOST_AUTHORITY_CACHE_PRUNE_AT {
self.wrong_host_authority.clear();
}
}
self.wrong_host_authority.insert(cache_key, Instant::now());
}
pub(crate) fn check_host_authority(
&mut self,
did: &Did,
repo_state: &mut RepoState,
source_host: &str,
) -> Result<Checked<AuthorityOutcome>> {
let pds_host = |pds: &str| {
Url::parse(pds)
.ok()
.and_then(|u| u.host_str().map(SmolStr::new))
};
let cache_key = Self::wrong_host_authority_key(did, source_host);
let expected = repo_state.pds.as_deref().and_then(pds_host);
if expected.as_deref() == Some(source_host) {
self.wrong_host_authority.remove(&cache_key);
return Ok(Checked::Done(AuthorityOutcome::Authorized));
}
if let Some(expected) = expected.clone()
&& self.recently_rejected_wrong_host_authority(cache_key)
{
return Ok(Checked::Done(AuthorityOutcome::WrongHost { expected }));
}
// try again once
let before = repo_state.identity();
if let Checked::Wait(kind) = self.refresh_doc(did, repo_state)? {
return Ok(Checked::Wait(kind));
}
let identity_moved = repo_state.identity() != before;
let Some(expected) = repo_state.pds.as_deref().and_then(pds_host) else {
miette::bail!("can't get pds host???");
};
if expected.as_str() == source_host {
self.wrong_host_authority.remove(&cache_key);
Ok(Checked::Done(AuthorityOutcome::WasStale { identity_moved }))
} else {
self.remember_wrong_host_authority(cache_key);
Ok(Checked::Done(AuthorityOutcome::WrongHost { expected }))
}
}
/// what the shard knows about `did`'s identity right now, see [`Fact`]
pub(crate) fn fact(&self, did: &Did) -> Fact<'_> {
match &self.answered {
Some(answered) if answered.did == *did => match &answered.result {
Ok(resolved) => Fact::Resolved(resolved, answered.resolved_at),
Err(_) => Fact::Failed,
},
_ if self.state.lookups.failing(did) => Fact::Failed,
_ => Fact::Unknown,
}
}
/// a did's sig failures refresh its doc at most once per cooldown. a replay already has
/// the fresh doc and checks with that
fn may_refresh_after_sig_failure(&self, did: &Did) -> bool {
matches!(self.fact(did), Fact::Resolved(..)) || self.state.lookups.start_sig_refresh(did)
}
/// the client a new account's pds probe goes through, only in builds that probe
pub(crate) fn probe_client(&self) -> Option<PublicHttpClient> {
#[cfg(any(not(feature = "relay"), feature = "indexer"))]
return Some(self.http.clone());
#[cfg(all(feature = "relay", not(feature = "indexer")))]
None
}
/// puts a freshly resolved doc on the row, once a lookup for it answered
pub(crate) fn refresh_doc(
&mut self,
did: &Did,
repo_state: &mut RepoState,
) -> Result<Checked<()>> {
let (doc, resolved_at) = match self.fact(did) {
Fact::Resolved(resolved, at) => (resolved.doc.clone(), at),
Fact::Failed => miette::bail!("couldn't resolve the doc of {did}"),
Fact::Unknown => return Ok(Checked::Wait(LookupKind::Refresh)),
};
self.stats.record_refresh_doc();
repo_state.update_from_resolve(doc, resolved_at);
repo_state.touch();
self.stage_repo_state(repo_state.clone().into_static());
Ok(Checked::Done(()))
}
/// starts a doc fetch for `did` that nobody waits on. its answer comes back to this shard,
/// which puts it on the row and announces it from `origin` if the identity moved
pub(crate) fn refresh_in_background(&mut self, did: &Did, origin: RefreshOrigin) {
// the last one's answer may sit in this shard's queue already, out of the pool's in-flight
// set but not on the row yet
if self.refresh_origins.contains_key(did) {
return;
}
let request = Request {
did: did.clone(),
kind: LookupKind::Refresh,
priority: Priority::Background,
resolver: self.state.resolver.clone(),
probe: None,
};
if self.state.lookups.submit(request, &self.answers) == Submitted::Started {
self.stats.record_background_refresh();
self.refresh_origins.insert(did.clone(), origin);
}
}
/// the row of a did this shard still processes, `None` for an untracked or unknown one
pub(crate) fn read_repo_state(&self, did: &Did) -> Result<Option<RepoState<'static>>> {
let db = &self.state.db;
let tracked = db
.repo_metadata
.get(keys::repo_metadata_key(did))
.into_diagnostic()?
.map(|bytes| db::deser_repo_meta(&bytes))
.transpose()?
.is_none_or(|metadata| metadata.tracked);
if !tracked {
return Ok(None);
}
self.current_row(did)
}
/// `did`'s row as this shard left it. in indexer builds that's the durable row with what
/// the indexer hasn't applied yet on top, see [`super::overlay`]
fn current_row(&self, did: &Did) -> Result<Option<RepoState<'static>>> {
let read = || {
self.state
.db
.repos
.get(keys::repo_key(did))
.into_diagnostic()?
.map(|bytes| db::deser_repo_state(bytes.as_ref()).map(|s| s.into_static()))
.transpose()
};
#[cfg(feature = "indexer")]
return self.state.relay_rows.view(did, read);
#[cfg(not(feature = "indexer"))]
read()
}
/// stage an identity learned while authorizing another event. it is committed
/// and flushed through the same sink as the triggering event.
#[cfg(any(feature = "indexer_stream", feature = "relay"))]
pub(crate) fn announce_identity(
&mut self,
did: &Did,
repo_state: &RepoState,
firehose: &Url,
seq: i64,
time: crate::ingest::stream::Datetime,
) -> Result<()> {
self.sink.identity(
self.state,
&mut self.batch,
firehose,
crate::ingest::stream::Identity {
did: did.clone(),
handle: repo_state.handle.clone(),
seq,
time,
},
true,
false,
)
}
pub(crate) fn validate_commit<'c>(
&mut self,
repo_state: &mut RepoState,
firehose: &Url,
commit: &'c crate::ingest::stream::Commit<'c>,
) -> Result<Checked<Option<ValidatedCommit<'c>>>> {
let did = &commit.repo;
let origin = || RefreshOrigin {
firehose: firehose.clone(),
seq: commit.seq,
time: commit.time.clone(),
};
let key = match self.signing_key(did, repo_state, origin)? {
Checked::Done(key) => key,
Checked::Wait(kind) => return Ok(Checked::Wait(kind)),
};
let validate_started = StatsInstant::now();
let validation = self.vctx.validate_commit(commit, repo_state, key.as_ref());
self.stats.record_validate_commit(
validate_started.elapsed(),
commit_validation_outcome(&validation),
);
match validation {
Ok(v) => return Ok(Checked::Done(Some(v))),
Err(CommitValidationError::StaleRev) => {
trace!("skipping replayed commit");
return Ok(Checked::Done(None));
}
// a sig failure may just mean our cached signing key is stale;
// refresh the identity doc and retry before treating it as a real failure.
Err(CommitValidationError::SigFailure) => {
if !self.may_refresh_after_sig_failure(did) {
debug!(%did, "commit signature failed, key refreshed recently");
return Ok(Checked::Done(None));
}
debug!(%did, "commit signature failed, refreshing signing key");
}
Err(e) => {
debug!(err = %e, "commit rejected");
return Ok(Checked::Done(None));
}
}
if let Checked::Wait(kind) = self.refresh_doc(did, repo_state)? {
return Ok(Checked::Wait(kind));
}
let key = repo_state
.signing_key
.as_ref()
.and_then(DidKey::to_public_key);
let validate_started = StatsInstant::now();
let validation = self.vctx.validate_commit(commit, repo_state, key.as_ref());
self.stats.record_validate_commit(
validate_started.elapsed(),
commit_validation_outcome(&validation),
);
match validation {
Ok(v) => Ok(Checked::Done(Some(v))),
Err(CommitValidationError::SigFailure) => {
warn!(%did, "commit signature verification failed after key refresh");
Ok(Checked::Done(None))
}
Err(e) => {
debug!(err = %e, "commit rejected after key refresh");
Ok(Checked::Done(None))
}
}
}
pub(crate) fn validate_sync(
&mut self,
repo_state: &mut RepoState,
firehose: &Url,
sync: &crate::ingest::stream::Sync<'_>,
) -> Result<Checked<Option<ValidatedSync>>> {
let did = &sync.did;
let origin = || RefreshOrigin {
firehose: firehose.clone(),
seq: sync.seq,
time: sync.time.clone(),
};
let key = match self.signing_key(did, repo_state, origin)? {
Checked::Done(key) => key,
Checked::Wait(kind) => return Ok(Checked::Wait(kind)),
};
let validate_started = StatsInstant::now();
let validation = self.vctx.validate_sync(sync, key.as_ref());
self.stats.record_validate_sync(
validate_started.elapsed(),
sync_validation_outcome(&validation),
);
match validation {
Ok(v) => return Ok(Checked::Done(Some(v))),
// a sig failure may just mean our cached signing key is stale;
// refresh the identity doc and retry before treating it as a real failure.
Err(SyncValidationError::SigFailure) => {
if !self.may_refresh_after_sig_failure(did) {
debug!(%did, "sync signature failed, key refreshed recently");
return Ok(Checked::Done(None));
}
debug!(%did, "sync signature failed, refreshing signing key");
}
Err(e) => {
debug!(err = %e, "sync rejected");
return Ok(Checked::Done(None));
}
}
if let Checked::Wait(kind) = self.refresh_doc(did, repo_state)? {
return Ok(Checked::Wait(kind));
}
let key = repo_state
.signing_key
.as_ref()
.and_then(DidKey::to_public_key);
let validate_started = StatsInstant::now();
let validation = self.vctx.validate_sync(sync, key.as_ref());
self.stats.record_validate_sync(
validate_started.elapsed(),
sync_validation_outcome(&validation),
);
match validation {
Ok(v) => Ok(Checked::Done(Some(v))),
Err(SyncValidationError::SigFailure) => {
warn!(%did, "sync signature verification failed after key refresh");
Ok(Checked::Done(None))
}
Err(e) => {
debug!(err = %e, "sync rejected after key refresh");
Ok(Checked::Done(None))
}
}
}
/// the key to check `did`'s signatures with. the row keeps the one from the doc we last
/// resolved, and it's used as is. one older than [`IDENTITY_TTL`] gets a doc fetch in the
/// background, so a rotation nobody announced is picked up within about a day. a row with
/// no key waits for a lookup, and then keeps what it got
fn signing_key(
&mut self,
did: &Did,
repo_state: &mut RepoState,
origin: impl FnOnce() -> RefreshOrigin,
) -> Result<Checked<Option<PublicKey<'static>>>> {
if !self.verify_signatures {
return Ok(Checked::Done(None));
}
if let Some(key) = repo_state
.signing_key
.as_ref()
.and_then(DidKey::to_public_key)
{
self.stats.record_stored_key();
let resolved_at =
repo_state.identity_resolved_or_spread(did, self.identity_epoch, IDENTITY_TTL);
if chrono::Utc::now().timestamp() - resolved_at >= IDENTITY_TTL {
self.refresh_in_background(did, origin());
}
return Ok(Checked::Done(Some(key)));
}
let (doc, resolved_at) = match self.fact(did) {
Fact::Resolved(resolved, at) => (resolved.doc.clone(), at),
Fact::Failed => miette::bail!("couldn't resolve the signing key of {did}"),
Fact::Unknown => return Ok(Checked::Wait(LookupKind::Key)),
};
self.stats.record_fetch_key();
let key = doc.key.clone();
repo_state.update_from_resolve(doc, resolved_at);
repo_state.touch();
self.stage_repo_state(repo_state.clone().into_static());
key.map(|key| Checked::Done(Some(key)))
.ok_or_else(|| NoSigningKeyError(did.clone()).into())
}
/// maps an inactive account status to the corresponding `RepoStatus`.
/// panics on `AccountStatus::Deleted`, caller must handle that
pub(crate) fn inactive_account_repo_status(
&self,
did: &Did,
status: &Option<AccountStatus<'_>>,
) -> RepoStatus {
match status {
Some(AccountStatus::Takendown) => RepoStatus::Takendown,
Some(AccountStatus::Suspended) => RepoStatus::Suspended,
Some(AccountStatus::Deactivated) => RepoStatus::Deactivated,
Some(AccountStatus::Throttled) => RepoStatus::Throttled,
Some(AccountStatus::Desynchronized) => RepoStatus::Desynchronized,
Some(AccountStatus::Other(s)) => {
warn!(did = %did, status = %s, "unknown account status");
RepoStatus::Error(s.to_smolstr())
}
Some(AccountStatus::Deleted) => {
unreachable!("deleted is handled before status mapping")
}
None => {
warn!(did = %did, "account inactive but no status provided");
RepoStatus::Error("unknown".into())
}
}
}
pub(crate) fn load_repo_state(
&mut self,
msg: &WorkerMessage,
) -> Result<Checked<Option<RepoState<'static>>>> {
self.loaded = None;
let db = &self.state.db;
let did = msg.msg.did().expect("we checked if valid");
let metadata_key = keys::repo_metadata_key(did);
let metadata = db
.repo_metadata
.get(&metadata_key)
.into_diagnostic()?
.map(|bytes| db::deser_repo_meta(&bytes))
.transpose()?;
if metadata.is_some_and(|m| !m.tracked) {
trace!(did = %did, "ignoring message, repo is explicitly untracked");
self.stats.record_repo_state_outcome(
crate::ingest::firehose_stats::RepoStateLoadOutcome::Drop,
);
return Ok(Checked::Done(None));
}
if let Some(repo_state) = self.current_row(did)? {
self.stats.record_repo_state_outcome(
crate::ingest::firehose_stats::RepoStateLoadOutcome::Hit,
);
self.loaded = Some(repo_state.clone());
return Ok(Checked::Done(Some(repo_state)));
}
#[cfg(feature = "indexer")]
{
let filter = self.state.filter.load();
if filter.mode == crate::filter::FilterMode::Filter && !filter.signals.is_empty() {
let commit = match &msg.msg {
SubscribeReposMessage::Commit(c) => c,
_ => {
self.stats.record_repo_state_outcome(
crate::ingest::firehose_stats::RepoStateLoadOutcome::Drop,
);
return Ok(Checked::Done(None));
}
};
let touches_signal = commit.ops.iter().any(|op| {
op.path
.split_once('/')
.map(|(col, _)| {
let m = filter.matches_signal(col);
debug!(
did = %did, path = %op.path, col = %col,
signals = ?filter.signals, matched = m,
"signal check"
);
m
})
.unwrap_or(false)
});
if !touches_signal {
trace!(did = %did, "dropping commit, no signal-matching ops");
self.stats.record_repo_state_outcome(
crate::ingest::firehose_stats::RepoStateLoadOutcome::Drop,
);
return Ok(Checked::Done(None));
}
}
}
// a lookup for a known did means the row was there when the message started waiting,
// so it was erased meanwhile, and the message goes with it
if self
.answered
.as_ref()
.is_some_and(|answered| answered.did == *did && answered.kind != LookupKind::NewAccount)
{
trace!(did = %did, "dropping message, repo was erased while it waited");
self.stats.record_repo_state_outcome(
crate::ingest::firehose_stats::RepoStateLoadOutcome::Drop,
);
return Ok(Checked::Done(None));
}
// the doc (and in indexer builds the pds's word on the account) comes from a lookup
let (doc, status, resolved_at) = match self.fact(did) {
Fact::Resolved(resolved, at) => (resolved.doc.clone(), resolved.status.clone(), at),
Fact::Failed => miette::bail!("couldn't resolve the doc of new account {did}"),
Fact::Unknown => return Ok(Checked::Wait(LookupKind::NewAccount)),
};
debug!(did = %did, "discovered new account from firehose, queueing backfill");
let new_account_started = StatsInstant::now();
// if it's a PDS, verify it's the authoritative one
if msg.is_pds {
let pds_host = doc.pds.host_str().map(|h| h.to_string());
if pds_host.as_deref() != msg.firehose.host_str() {
warn!(did = %did, got = ?pds_host, expected = ?msg.firehose.host_str(), "message rejected: wrong host for new account");
self.stats.record_repo_state_outcome(
crate::ingest::firehose_stats::RepoStateLoadOutcome::Drop,
);
return Ok(Checked::Done(None));
}
if let Some(host) = msg.firehose.host_str() {
let count = self
.state
.db
.get_count_sync(&keys::pds_account_count_key(host));
if self.state.is_over_account_limit(host, count) {
warn!(did = %did, host, count, "account limit reached for host, dropping new account");
self.stats.record_repo_state_outcome(
crate::ingest::firehose_stats::RepoStateLoadOutcome::Drop,
);
return Ok(Checked::Done(None));
}
}
}
#[cfg(any(not(feature = "relay"), feature = "indexer"))]
let mut repo_state = status.unwrap_or_else(RepoState::backfilling);
#[cfg(all(feature = "relay", not(feature = "indexer")))]
let mut repo_state = {
let _ = status;
RepoState::synced()
};
repo_state.update_from_resolve(doc, resolved_at);
self.stage_repo_state(repo_state.clone().into_static());
// in indexer builds the indexer counts the repo once it makes the row, and queues its
// backfill
#[cfg(not(feature = "indexer"))]
self.count_deltas.add_repos(1);
self.stats
.record_repo_state_outcome(crate::ingest::firehose_stats::RepoStateLoadOutcome::Miss);
self.stats.record_new_account(new_account_started.elapsed());
Ok(Checked::Done(Some(repo_state)))
}
}