lightrail / src / identity.rs

//! DID identity resolution — maps DIDs to their home PDS endpoint and signing key.
//!
//! [`build_resolver`] constructs a [`Resolver`] from optional slingshot and
//! PLC directory URLs. The four combinations are:
//!
//! | `slingshot_url` | `plc_url`   | Behaviour                                          |
//! |-----------------|-------------|----------------------------------------------------|
//! | `None`          | `None`      | PLC directory at `https://plc.directory` (default) |
//! | `None`          | `Some(url)` | PLC directory at `url`                             |
//! | `Some(url)`     | `None`      | Slingshot at `url`                                 |
//! | `Some(url)`     | `Some(fb)`  | Slingshot primary, PLC directory fallback          |
//!
//! When slingshot is active it covers both `did:plc` (via the `PlcHttp`
//! resolution step) and `did:web` (via the `PdsResolveDid` → mini-doc
//! fallback path built into jacquard-identity). When a PLC fallback is also
//! configured, any error from the primary resolver retries against it.

use jacquard_common::IntoStatic;
use std::num::NonZeroU32;
use std::sync::Arc;
use std::time::Duration;

use dashmap::DashMap;
use governor::{DefaultDirectRateLimiter, Quota, RateLimiter};
use jacquard_common::types::crypto::PublicKey;
use jacquard_common::types::string::Did;
use jacquard_common::url::Url;
use jacquard_identity::JacquardResolver;
use jacquard_identity::resolver::{
    IdentityError as JacquardIdentityError, IdentityErrorKind as JacquardErrorKind,
    IdentityResolver, PlcSource, ResolverOptions,
};
use mini_moka::sync::Cache;
use tracing::{debug, info, warn};

use crate::util::TokenExt;

#[derive(Debug, thiserror::Error)]
pub enum IdentityError {
    #[error("jacquard: {0}")]
    JacquardIdentity(Box<JacquardIdentityError>),
    #[error("no pds endpoint for: {0}")]
    MissingPdsEndpoint(String),
    #[error("bad signing key: {0}")]
    BadSigningKey(String),
    #[error("identity: {0}")]
    Identity(String),
}

impl From<JacquardIdentityError> for IdentityError {
    fn from(je: JacquardIdentityError) -> Self {
        Self::JacquardIdentity(Box::new(je))
    }
}

/// Resolved identity for a DID: the home PDS endpoint and ATProto signing key.
///
/// Stored in the cache as `Arc<CachedIdentity>` so callers can borrow fields
/// without copying. The PDS URL is an interned `Arc<Url>` (see [`UrlInterner`]);
/// the public key bytes are owned per-entry and borrowed through the outer Arc.
pub struct CachedIdentity {
    pub pds: Arc<Url>,
    pub pubkey: PublicKey<'static>,
}

const DEFAULT_PLC_URL: &str = "https://plc.directory";

/// How long a cached DID→PDS mapping is considered fresh.
///
/// PDS migrations are rare; the firehose delivers `#identity` events when
/// they happen, so proactive [`Resolver::invalidate_did`] handles updates
/// before this TTL expires in practice.
const CACHE_TTL: Duration = Duration::from_secs(24 * 60 * 60);

/// How long a failed resolution is cached before retrying.
///
/// Short enough that transient failures recover quickly, long enough to avoid
/// hammering the identity service for consistently-broken DIDs.
const NEGATIVE_CACHE_TTL: Duration = Duration::from_secs(5 * 60);

/// Maximum retries per identity backend on HTTP 429 responses.
const RATE_LIMIT_RETRIES: u32 = 3;

/// Initial delay between rate-limit retries; doubles each attempt (500ms → 1s → 2s).
const RATE_LIMIT_BASE_DELAY: Duration = Duration::from_millis(500);

/// Check whether a jacquard identity error is an HTTP 429 rate-limit response.
fn is_rate_limited(err: &JacquardIdentityError) -> bool {
    matches!(
        err.kind(),
        JacquardErrorKind::HttpStatus(status) if status.as_u16() == 429
    )
}

// ---------------------------------------------------------------------------
// Resolution throttle
// ---------------------------------------------------------------------------

/// Global identity-resolution rate limiter.
///
/// Backfill tasks call [`Resolver::throttle_wait`] to stay within the budget.
/// All callers (firehose, resync, backfill) consume tokens on cache-miss
/// resolutions via an internal notify, so firehose activity naturally reduces
/// the budget available to backfill.
struct IdentityThrottle {
    limiter: DefaultDirectRateLimiter,
    token_interval: Duration,
}

impl IdentityThrottle {
    fn new(rate: NonZeroU32) -> Self {
        Self {
            limiter: RateLimiter::direct(Quota::per_second(rate)),
            token_interval: Duration::from_secs(1) / rate.get(),
        }
    }
}

// ---------------------------------------------------------------------------
// Interned PDS URL cache
// ---------------------------------------------------------------------------

/// Maximum number of distinct PDS URLs to intern.
///
/// The real network has ~2600 PDS hosts online; this cap is 4× that. Beyond
/// it, `intern` returns a fresh `Arc` rather than a shared one, keeping the
/// table bounded against identities with adversarially unique PDS URLs.
const MAX_INTERN_ENTRIES: usize = 10_000;

/// Interns `Url` values so that the ~2600 distinct PDS endpoints across the
/// network are each allocated once, regardless of how many DIDs map to them.
struct UrlInterner {
    map: DashMap<String, Arc<Url>>,
}

impl UrlInterner {
    fn new() -> Self {
        Self {
            map: DashMap::new(),
        }
    }

    fn intern(&self, url: Url) -> Arc<Url> {
        let key = url.as_str().to_owned();
        if let Some(existing) = self.map.get(&key) {
            return Arc::clone(&existing);
        }
        // Past the cap, skip interning to bound memory regardless of
        // how many distinct PDS URLs external parties claim.
        if self.map.len() >= MAX_INTERN_ENTRIES {
            return Arc::new(url);
        }
        let arc = Arc::new(url);
        Arc::clone(
            self.map
                .entry(key)
                .or_insert_with(|| Arc::clone(&arc))
                .value(),
        )
    }
}

/// Capacity-bounded, TTL-expiring cache from DID strings to resolved identities.
///
/// Successful resolutions are cached for [`CACHE_TTL`] (24 h). Failed
/// resolutions are negative-cached for [`NEGATIVE_CACHE_TTL`] (5 min) so that
/// consistently-broken DIDs don't hammer the identity service on every event.
struct IdentityCache {
    entries: Cache<Did<'static>, Arc<CachedIdentity>>,
    /// DIDs whose last resolution failed. Presence means "don't retry yet".
    negative: Cache<Did<'static>, ()>,
    interner: UrlInterner,
}

impl IdentityCache {
    fn new(capacity: u64) -> Self {
        Self {
            entries: Cache::builder()
                .max_capacity(capacity)
                .time_to_live(CACHE_TTL)
                .build(),
            negative: Cache::builder()
                .max_capacity(capacity / 10)
                .time_to_live(NEGATIVE_CACHE_TTL)
                .build(),
            interner: UrlInterner::new(),
        }
    }

    fn get(&self, did: &Did<'_>) -> Option<Arc<CachedIdentity>> {
        self.entries.get(&did.clone().into_static())
    }

    /// Returns `true` if this DID recently failed resolution.
    fn is_negative(&self, did: &Did<'_>) -> bool {
        self.negative.contains_key(&did.clone().into_static())
    }

    fn insert(
        &self,
        did: Did<'static>,
        pds: Url,
        pubkey: PublicKey<'static>,
    ) -> Arc<CachedIdentity> {
        let pds = self.interner.intern(pds);
        let identity = Arc::new(CachedIdentity { pds, pubkey });
        self.negative.invalidate(&did);
        self.entries.insert(did, Arc::clone(&identity));
        identity
    }

    fn insert_negative(&self, did: &Did<'_>) {
        self.negative.insert(did.clone().into_static(), ());
    }

    fn invalidate(&self, did: &Did<'_>) {
        let owned = did.clone().into_static();
        self.entries.invalidate(&owned);
        self.negative.invalidate(&owned);
    }
}

// ---------------------------------------------------------------------------
// Public resolver
// ---------------------------------------------------------------------------

/// DID identity resolver with a DID→identity cache and optional fallback.
///
/// Wrap in [`Arc`] to share across tasks; all cache operations use interior
/// mutability so `&self` is sufficient everywhere.
pub struct Resolver {
    primary: JacquardResolver,
    fallback: Option<JacquardResolver>,
    cache: IdentityCache,
    /// Cancellation token for cooperative shutdown.
    token: tokio_util::sync::CancellationToken,
    /// Optional global resolution rate limit.
    throttle: Option<IdentityThrottle>,
}

impl Resolver {
    /// Wait for a resolution token (backfill / resync).
    ///
    /// Blocks until a token is available or the cancellation token fires.
    /// No-op when no throttle is configured.
    async fn throttle_wait(&self) {
        let Some(t) = &self.throttle else { return };
        while t.limiter.check().is_err() {
            if !self.token.sleep(t.token_interval).await {
                return; // shutting down
            }
        }
    }

    /// Try to consume a resolution token without waiting (firehose).
    ///
    /// Subtracts from the shared budget so backfill/resync slow down when the
    /// firehose is doing lots of resolutions. No-op if the rate is already
    /// exceeded or no throttle is configured.
    fn throttle_notify(&self) {
        if let Some(t) = &self.throttle {
            let _ = t.limiter.check();
        }
    }

    /// Resolve the PDS endpoint and ATProto signing key for `did`.
    ///
    /// Returns a cached [`CachedIdentity`] on cache hit; otherwise waits for
    /// a throttle token and resolves via the network. Used by backfill and
    /// resync paths.
    pub async fn resolve(&self, did: &Did<'_>) -> Result<Arc<CachedIdentity>, IdentityError> {
        self.resolve_inner(did, false).await
    }

    /// Like [`resolve`] but never waits on the throttle — only notifies it.
    ///
    /// Used by the firehose path where latency matters. Cache-miss resolutions
    /// still subtract from the shared budget so backfill/resync slow down.
    pub async fn resolve_firehose(
        &self,
        did: &Did<'_>,
    ) -> Result<Arc<CachedIdentity>, IdentityError> {
        self.resolve_inner(did, true).await
    }

    async fn resolve_inner(
        &self,
        did: &Did<'_>,
        firehose: bool,
    ) -> Result<Arc<CachedIdentity>, IdentityError> {
        if let Some(cached) = self.cache.get(did) {
            metrics::counter!("lightrail_identity_resolution_total", "outcome" => "hit")
                .increment(1);
            return Ok(cached);
        }

        if self.cache.is_negative(did) {
            metrics::counter!("lightrail_identity_resolution_total", "outcome" => "negative_hit")
                .increment(1);
            return Err(IdentityError::Identity(format!(
                "recently failed for {}",
                did.as_str()
            )));
        }

        if firehose {
            self.throttle_notify();
        } else {
            self.throttle_wait().await;
        }

        match self.resolve_uncached(did).await {
            Ok((pds, pubkey)) => {
                metrics::counter!("lightrail_identity_resolution_total", "outcome" => "miss")
                    .increment(1);
                Ok(self.cache.insert(did.clone().into_static(), pds, pubkey))
            }
            Err(e) => {
                metrics::counter!("lightrail_identity_resolution_total", "outcome" => "error")
                    .increment(1);
                self.cache.insert_negative(did);
                Err(e)
            }
        }
    }

    /// Look up `did` in the cache without making network calls.
    /// Returns `None` on cache miss.
    pub fn resolve_cached(&self, did: &Did<'_>) -> Option<Arc<CachedIdentity>> {
        self.cache.get(did)
    }

    /// Insert a resolved identity into the cache without a network call.
    ///
    /// Used by the dispatcher to warm the cache for discovery-queue items
    /// whose identity was already resolved during backfill. This avoids a
    /// redundant network resolution in the resync worker.
    pub fn insert_cached(
        &self,
        did: Did<'static>,
        pds: Url,
        pubkey: PublicKey<'static>,
    ) -> Arc<CachedIdentity> {
        self.cache.insert(did, pds, pubkey)
    }

    /// Evict `did` from the identity cache.
    ///
    /// Called when a `#identity` firehose event is received, after all
    /// preceding commits for the same DID have been processed, so the next
    /// resolution fetches the updated key and PDS endpoint.
    pub fn invalidate_did(&self, did: &Did<'_>) {
        self.cache.invalidate(did);
    }

    async fn resolve_uncached(
        &self,
        did: &Did<'_>,
    ) -> Result<(Url, PublicKey<'static>), IdentityError> {
        let doc = match &self.fallback {
            None => self.try_backend(&self.primary, did).await,
            Some(fb) => match self.try_backend(&self.primary, did).await {
                ok @ Ok(_) => ok,
                Err(primary_err) => {
                    debug!(
                        did = did.as_str(),
                        error = %primary_err,
                        "primary identity resolution failed, trying fallback"
                    );
                    self.try_backend(fb, did).await
                }
            },
        }?;

        let pds = doc
            .pds_endpoint()
            .ok_or_else(|| IdentityError::MissingPdsEndpoint(did.to_string()))?;

        let pubkey = doc
            .atproto_public_key()
            .map_err(|e| IdentityError::BadSigningKey(e.to_string()))?
            .ok_or_else(|| {
                IdentityError::BadSigningKey(format!(
                    "no ATProto signing key in DID doc for {}",
                    did.as_str()
                ))
            })?;

        Ok((pds, pubkey))
    }

    /// Try resolving a DID document against a single backend, retrying on 429.
    ///
    /// Non-rate-limit errors return immediately. After exhausting retries the
    /// final 429 error is returned so the caller can fall through to a fallback
    /// backend or propagate the error. Sleeps are cancellation-aware when a
    /// token has been provided via [`Resolver::set_cancellation_token`].
    async fn try_backend(
        &self,
        backend: &JacquardResolver,
        did: &Did<'_>,
    ) -> std::result::Result<
        jacquard_common::types::did_doc::DidDocument<'static>,
        JacquardIdentityError,
    > {
        for attempt in 0..=RATE_LIMIT_RETRIES {
            match backend.resolve_did_doc_owned(did).await {
                Ok(doc) => return Ok(doc),
                Err(e) if is_rate_limited(&e) && attempt < RATE_LIMIT_RETRIES => {
                    metrics::counter!("lightrail_identity_rate_limited_total").increment(1);
                    let delay = RATE_LIMIT_BASE_DELAY * 2u32.pow(attempt);
                    metrics::counter!("lightrail_identity_rate_limit_backoff_ms")
                        .increment(delay.as_millis() as u64);
                    debug!(
                        did = did.as_str(),
                        attempt = attempt + 1,
                        delay_ms = delay.as_millis() as u64,
                        "identity service rate-limited, retrying"
                    );
                    if !self.token.sleep(delay).await {
                        return Err(e); // shutting down
                    }
                }
                Err(e) => {
                    if is_rate_limited(&e) {
                        metrics::counter!("lightrail_identity_rate_limited_total").increment(1);
                        warn!(
                            did = did.as_str(),
                            retries = RATE_LIMIT_RETRIES,
                            "identity resolution rate limit persisted after retries"
                        );
                    }
                    return Err(e);
                }
            }
        }
        unreachable!()
    }
}

// ---------------------------------------------------------------------------
// Construction
// ---------------------------------------------------------------------------

/// Build the identity resolver from the provided source URLs.
///
/// Emits an `INFO` trace on startup describing the active resolution behaviour.
pub fn build_resolver(
    slingshot_url: Option<Url>,
    plc_url: Option<Url>,
    cache_size: u64,
    token: tokio_util::sync::CancellationToken,
    identity_qps: Option<NonZeroU32>,
) -> Resolver {
    let cache = IdentityCache::new(cache_size);
    let throttle = identity_qps.map(|rate| {
        info!(rate = rate.get(), "identity resolution throttle enabled");
        IdentityThrottle::new(rate)
    });

    match (slingshot_url, plc_url) {
        (None, plc) => {
            let base = plc.unwrap_or_else(|| DEFAULT_PLC_URL.parse().expect("valid url"));
            info!(
                source = "plc",
                url = %base,
                "did resolution: plc directory; \
                 did:plc via plc directory, did:web via https well-known"
            );
            Resolver {
                primary: make_resolver(PlcSource::PlcDirectory { base }),
                fallback: None,
                cache,
                token: token.clone(),
                throttle,
            }
        }
        (Some(slingshot), None) => {
            info!(
                source = "slingshot",
                url = %slingshot,
                "did resolution: slingshot; \
                 did:plc via slingshot, \
                 did:web via https well-known then slingshot mini-doc"
            );
            Resolver {
                primary: make_resolver(PlcSource::Slingshot { base: slingshot }),
                fallback: None,
                cache,
                token: token.clone(),
                throttle,
            }
        }
        (Some(slingshot), Some(plc)) => {
            info!(
                slingshot_url = %slingshot,
                plc_fallback_url = %plc,
                "did resolution: slingshot primary with plc fallback; \
                 did:plc via slingshot (plc on error), \
                 did:web via https well-known then slingshot mini-doc"
            );
            Resolver {
                primary: make_resolver(PlcSource::Slingshot { base: slingshot }),
                fallback: Some(make_resolver(PlcSource::PlcDirectory { base: plc })),
                cache,
                token,
                throttle,
            }
        }
    }
}

fn make_resolver(plc_source: PlcSource) -> JacquardResolver {
    let opts = ResolverOptions {
        plc_source,
        ..Default::default()
    };
    let client = reqwest::Client::builder()
        .user_agent(concat!("microcosm lightrail/v", env!("CARGO_PKG_VERSION")))
        .build()
        .expect("failed to build resolver HTTP client");
    JacquardResolver::new(client, opts)
}