ic_auth/portable/

revocation.rs

//! Bounded volatile revocation bookkeeping, not an authentication endpoint.

use super::{VerifiedPortableClaims, browser_principal, is_login_caller, validate_shape};
use crate::admission_budget::{AdmissionBudgetError, SessionAdmissionBudget};
use ic_auth_protocol_types::{AudienceId, PortableClaims, Principal};
use std::collections::{BTreeMap, BTreeSet};
use thiserror::Error;

#[cfg(feature = "sessions")]
const MAX_RECORD_BYTES: usize = 1024;

/// A refusal leaves records unchanged; trusted clock observations are retained.
#[derive(Clone, Copy, Debug, Error, Eq, PartialEq)]
pub enum LoginRevocationError {
    #[error("logout requires an authenticated self-authenticating caller")]
    InvalidCaller,
    #[error("invalid revocation limits or authenticated login deadline")]
    InvalidPolicy,
    #[error("login is locally revoked")]
    Revoked,
    #[error("revocation capacity is exhausted")]
    Capacity,
    #[error("revocation clock rolled back or retention arithmetic overflowed")]
    Clock,
    #[error("host session/replay transaction failed")]
    Storage,
}

#[derive(Clone, Copy)]
#[cfg_attr(feature = "sessions", derive(serde::Serialize))]
#[cfg_attr(feature = "session-restore", derive(serde::Deserialize))]
struct Record {
    until_ns: u64,
    state: LoginRecordState,
}

// A reservation always owns admitted capacity; a tombstone retains that pool
// exactly when it retains an authenticated login binding.
/// Ledger state projected without inferring state from independent flags.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
#[cfg_attr(feature = "sessions", derive(serde::Serialize))]
#[cfg_attr(feature = "session-restore", derive(serde::Deserialize))]
pub enum LoginRecordState {
    Reserved(LoginBinding),
    Revoked(Option<LoginBinding>),
}

impl Record {
    fn login(self) -> Option<LoginBinding> {
        match self.state {
            LoginRecordState::Reserved(login) => Some(login),
            LoginRecordState::Revoked(login) => login,
        }
    }

    fn revoked(self) -> bool {
        matches!(self.state, LoginRecordState::Revoked(_))
    }

    fn admitted_pool(self) -> bool {
        self.login().is_some()
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
#[cfg_attr(feature = "sessions", derive(serde::Serialize))]
#[cfg_attr(feature = "session-restore", derive(serde::Deserialize))]
#[cfg_attr(feature = "sessions", serde(deny_unknown_fields))]
/// Authenticated login association retained by a reservation or known tombstone.
/// This value alone is not authentication or admission evidence.
pub struct LoginBinding {
    pub issuer: Principal,
    pub subject: Principal,
    pub started_at_ns: u64,
    pub deadline_ns: u64,
}

/// Proposed ledger record from a serialized commit callback or trusted storage.
/// A durable host must persist this with its session/replay changes before
/// returning success. This is not a credential or reusable authorization result.
pub struct LoginRevocationChange {
    audience: AudienceId,
    caller: Principal,
    record: Record,
}

impl LoginRevocationChange {
    /// Encode the proposed record for the host's atomic durable commit.
    #[cfg(feature = "sessions")]
    pub fn encode(&self) -> Result<Vec<u8>, LoginRevocationError> {
        crate::storage_codec::encode(
            &(self.audience, self.caller, self.record),
            crate::storage_codec::LOGIN,
            MAX_RECORD_BYTES,
            LoginRevocationError::Storage,
            LoginRevocationError::InvalidPolicy,
        )
    }

    /// Decode trusted storage only; never accept these bytes from a client.
    /// This checks shape, not provenance, completeness, rollback or live authority.
    #[cfg(feature = "session-restore")]
    pub fn decode(bytes: &[u8]) -> Result<Self, LoginRevocationError> {
        let (audience, caller, record): (_, Principal, Record) = crate::storage_codec::decode(
            bytes,
            crate::storage_codec::LOGIN,
            MAX_RECORD_BYTES,
            LoginRevocationError::InvalidPolicy,
            LoginRevocationError::InvalidPolicy,
        )?;
        if !is_login_caller(caller) {
            return Err(LoginRevocationError::InvalidCaller);
        }
        if let Some(login) = record.login()
            && (login.issuer == Principal::anonymous()
                || login.subject == Principal::anonymous()
                || login.started_at_ns >= login.deadline_ns
                || login.deadline_ns - login.started_at_ns > crate::limits::MAX_LOGIN_NS)
        {
            return Err(LoginRevocationError::InvalidPolicy);
        }
        let change = Self {
            audience,
            caller,
            record,
        };
        Ok(change)
    }

    /// Protected trust domain owning the record.
    pub fn audience(&self) -> AudienceId {
        self.audience
    }
    /// Authenticated runtime caller whose record changes.
    pub fn caller(&self) -> Principal {
        self.caller
    }
    /// Exclusive retention deadline, including the configured clock margin.
    pub fn until_ns(&self) -> u64 {
        self.record.until_ns
    }
    /// Typed state from this checked ledger record. Storage codecs retain their
    /// existing frozen representation independently of this Rust projection.
    pub fn state(&self) -> LoginRecordState {
        self.record.state
    }
}

/// Bounded reference ledger for one host-installed trust domain. Callers are
/// runtime principals, never request fields; this ledger cannot authenticate them.
/// One admitted slot survives lease expiry until the login deadline plus skew.
/// Unknown callers have a separate pool and cannot consume admitted capacity.
///
/// The host MUST serialize these operations with proof/session/replay state and
/// authority changes. Commit callbacks receive the exact proposed login record
/// and must atomically write it with that other state in durable backends, or
/// return an error without effects. No await/reentrancy may split the operation.
/// This ledger commits its staged record only after callback success, with no
/// further fallible policy/storage checks. Callbacks must not return success
/// before their commit, or commit and then return an error.
///
/// This is not portable admission or a logout endpoint. No proof verification,
/// startup/enrollment readiness or durable storage is supplied here. Hosts MUST
/// apply those gates and a qualified clock-skew bound separately. Fresh unknown
/// revocations consume an explicit host-retained budget before the commit callback.
/// Losing this ledger requires the protocol's fresh login/reset floors; restoring
/// sessions alone is unsafe. Never recreate it per request or share across domains.
pub struct MemoryLoginRevocations {
    audience: AudienceId,
    records: BTreeMap<Principal, Record>,
    expiry: BTreeSet<(u64, Principal)>,
    max_admitted: usize,
    max_unknown: usize,
    admitted: usize,
    unknown: usize,
    skew_ns: u64,
    observed_ns: u64,
}

impl MemoryLoginRevocations {
    /// Stream the complete physical ledger under a serialized host checkpoint.
    /// Persist these records with occupancy, clock and session/replay/authority
    /// state atomically; expired entries count until explicit pruning.
    #[cfg(feature = "session-restore")]
    pub fn records(&self) -> impl ExactSizeIterator<Item = LoginRevocationChange> + '_ {
        self.records
            .iter()
            .map(|(&caller, &record)| LoginRevocationChange {
                audience: self.audience,
                caller,
                record,
            })
    }

    /// Rebuild a fresh ledger from a complete, trusted atomic storage generation.
    /// `self` must be newly constructed at current trusted time. The host supplies
    /// that generation's persisted clock high-water mark and physical pool counts,
    /// including expired records, from the same transaction as the records.
    /// Failure consumes the partial ledger; no usable partial state is returned.
    /// Counts detect missing entries, not malicious rollback: hosts MUST protect
    /// snapshot provenance/generation and restore reset/authority/session/replay
    /// state together, keeping admission closed until fresh enrollment is checked.
    #[cfg(feature = "session-restore")]
    pub fn restore_records(
        mut self,
        observed_ns: u64,
        expected_counts: (usize, usize),
        records: impl IntoIterator<Item = LoginRevocationChange>,
    ) -> Result<Self, LoginRevocationError> {
        if !self.records.is_empty()
            || observed_ns > self.observed_ns
            || expected_counts.0 > self.max_admitted
            || expected_counts.1 > self.max_unknown
        {
            return Err(LoginRevocationError::InvalidPolicy);
        }
        let max_until = retention(observed_ns, self.skew_ns)?;
        for change in records {
            if change.audience != self.audience || self.records.contains_key(&change.caller) {
                return Err(LoginRevocationError::InvalidPolicy);
            }
            let record = change.record;
            if record.until_ns == 0 || record.until_ns > max_until {
                return Err(LoginRevocationError::InvalidPolicy);
            }
            if let Some(login) = record.login() {
                let minimum = login
                    .deadline_ns
                    .checked_add(self.skew_ns)
                    .ok_or(LoginRevocationError::Clock)?;
                if login.started_at_ns > observed_ns
                    || record.until_ns < minimum
                    || (!record.revoked() && record.until_ns != minimum)
                {
                    return Err(LoginRevocationError::InvalidPolicy);
                }
            }
            // Stop at the persisted counts, even for an unbounded input iterator.
            let (used, expected) = if record.admitted_pool() {
                (self.admitted, expected_counts.0)
            } else {
                (self.unknown, expected_counts.1)
            };
            if used >= expected {
                return Err(LoginRevocationError::Capacity);
            }
            self.install(change.caller, record);
        }
        if self.occupancy() != expected_counts {
            return Err(LoginRevocationError::InvalidPolicy);
        }
        Ok(self)
    }

    /// Fixed record counts bound memory; no caller-supplied blobs are retained.
    /// `skew_ns` is deployment-qualified protected configuration, not a default.
    pub fn new(
        audience: AudienceId,
        max_admitted: usize,
        max_unknown: usize,
        skew_ns: u64,
        now_ns: u64,
    ) -> Result<Self, LoginRevocationError> {
        if max_admitted == 0 || max_unknown == 0 {
            return Err(LoginRevocationError::InvalidPolicy);
        }
        retention(now_ns, skew_ns)?;
        Ok(Self {
            audience,
            records: BTreeMap::new(),
            expiry: BTreeSet::new(),
            max_admitted,
            max_unknown,
            admitted: 0,
            unknown: 0,
            skew_ns,
            observed_ns: now_ns,
        })
    }

    #[cfg(feature = "sessions")]
    pub(crate) fn audience(&self) -> AudienceId {
        self.audience
    }

    // One observation for a serialized session operation. Missing reservation is
    // checked after the session's ordinary liveness/authority refusals; denial
    // still precedes those refusals. This is never retained between operations.
    #[cfg(feature = "sessions")]
    pub(crate) fn observe_login(
        &mut self,
        caller: Principal,
        now_ns: u64,
    ) -> Result<Option<LoginBinding>, LoginRevocationError> {
        self.observe(caller, now_ns)?;
        let record = self.records.get(&caller).filter(|r| r.until_ns > now_ns);
        if record.is_some_and(|r| r.revoked()) {
            return Err(LoginRevocationError::Revoked);
        }
        Ok(record.and_then(|r| r.login()))
    }

    /// Physical pool occupancy, including expired entries until bounded pruning.
    pub fn occupancy(&self) -> (usize, usize) {
        (self.admitted, self.unknown)
    }

    /// Clock high-water mark for an atomic host checkpoint, including refusals.
    /// Durable hosts retain this with pool counts and the complete record set.
    pub fn observed_at_ns(&self) -> u64 {
        self.observed_ns
    }

    /// Cheap refusal before proof verification. This does not authenticate claims,
    /// reserve capacity or permit a session commit; success is only permission to
    /// continue checking. Untrusted claims can cause refusal, never admission.
    /// Records and pool occupancy stay unchanged; trusted clock observations remain.
    /// The host must serialize this check, verification and `with_reservation`
    /// with authority/session/replay changes. The final reservation rechecks state.
    pub fn check_reservation(
        &mut self,
        caller: Principal,
        claims: &PortableClaims,
        now_ns: u64,
    ) -> Result<(), LoginRevocationError> {
        self.reservation(caller, claims, now_ns).map(|_| ())
    }

    #[cfg(all(
        feature = "sessions",
        any(
            feature = "portable-signature-verification",
            feature = "certified-proof-verification"
        )
    ))]
    pub(crate) fn check_decoded_reservation(
        &mut self,
        caller: Principal,
        decoded: &super::DecodedPortableClaims,
        now_ns: u64,
    ) -> Result<(), LoginRevocationError> {
        self.observe(caller, now_ns)?;
        self.reservation_checked(caller, &decoded.claims, now_ns)
            .map(|_| ())
    }

    /// Reserve before committing a fully verified admission's session/replay pair.
    /// The issuer/account and original login window stay unchanged on renewal.
    /// Verified claims must match this ledger domain, caller and current clock.
    /// The host still enforces current authority and replay under the same lock.
    /// A live tombstone rejects even an exact admission retry before its callback.
    /// No reservation is installed if session/replay storage in `commit` fails.
    /// The callback receives the proposed record for that same atomic commit.
    /// Its error is returned unchanged; `E` also carries ledger refusals through
    /// `From<LoginRevocationError>`. No separate storage-error channel is needed.
    pub fn with_reservation<T, E: From<LoginRevocationError>>(
        &mut self,
        caller: Principal,
        verified: &VerifiedPortableClaims,
        now_ns: u64,
        commit: impl FnOnce(&LoginRevocationChange) -> Result<T, E>,
    ) -> Result<T, E> {
        self.observe(caller, now_ns)?;
        let record = self.reservation_checked(caller, verified.claims(), now_ns)?;
        if verified.verified_at_ns() != now_ns {
            return Err(LoginRevocationError::InvalidPolicy.into());
        }
        let value = commit(&LoginRevocationChange {
            audience: self.audience,
            caller,
            record,
        })?;
        self.install(caller, record);
        Ok(value)
    }

    /// Stage local denial, then atomically remove this caller's host sessions in
    /// `commit`. Return the denial deadline only after that callback commits.
    /// An existing live tombstone keeps its deadline; the callback still runs so
    /// retry never falsely acknowledges a failed host operation. No proof needed.
    /// An admitted reservation converts in its own pool even if unknowns fill theirs.
    /// A caller without a live record consumes one unknown-logout budget credit
    /// before the callback; exhaustion returns Capacity, never false success.
    /// Keep that receiver-wide budget separate from fresh proof verification so
    /// unknown logout traffic cannot spend admission credit. Failure retains the
    /// charge; hosts must return normal errors rather than trap/roll back credit.
    /// Live reservations and tombstone retries do not compete for this budget.
    /// The callback receives the proposed tombstone, including retained binding.
    /// Its error is returned unchanged; `E` also carries ledger refusals through
    /// `From<LoginRevocationError>`. A failed callback never installs denial.
    pub fn with_revocation<E: From<LoginRevocationError>>(
        &mut self,
        unknown_budget: &mut SessionAdmissionBudget,
        caller: Principal,
        now_ns: u64,
        commit: impl FnOnce(&LoginRevocationChange) -> Result<(), E>,
    ) -> Result<u64, E> {
        self.observe(caller, now_ns)?;
        let old = self
            .records
            .get(&caller)
            .copied()
            .filter(|r| r.until_ns > now_ns);
        let record = if let Some(old) = old.filter(|r| r.revoked()) {
            old
        } else {
            let until_ns = retention(now_ns, self.skew_ns)?;
            if old.is_some_and(|r| r.until_ns > until_ns) {
                return Err(LoginRevocationError::Clock.into());
            }
            Record {
                until_ns,
                state: LoginRecordState::Revoked(old.and_then(Record::login)),
            }
        };
        self.capacity(caller, record)?;
        if old.is_none() {
            unknown_budget
                .consume(now_ns)
                .map_err(|error| match error {
                    AdmissionBudgetError::Exhausted => LoginRevocationError::Capacity,
                    AdmissionBudgetError::ClockRollback => LoginRevocationError::Clock,
                    AdmissionBudgetError::InvalidPolicy => LoginRevocationError::InvalidPolicy,
                })?;
        }
        commit(&LoginRevocationChange {
            audience: self.audience,
            caller,
            record,
        })?;
        self.install(caller, record);
        Ok(record.until_ns)
    }

    /// Consult local denial before any session use, including exact retries.
    /// Absence is not authentication or authorization; other gates still apply.
    pub fn is_revoked(
        &mut self,
        caller: Principal,
        now_ns: u64,
    ) -> Result<bool, LoginRevocationError> {
        self.observe(caller, now_ns)?;
        Ok(self
            .records
            .get(&caller)
            .is_some_and(|r| r.revoked() && r.until_ns > now_ns))
    }

    /// Explicit maintenance, at most 128 expired entries; never evict live state.
    /// The host clock must be monotonic, including across failed operations.
    pub fn prune(&mut self, now_ns: u64, budget: usize) -> Result<usize, LoginRevocationError> {
        self.prune_with_commit(now_ns, budget, |_, _| Ok(()))
    }
    /// Commit bounded expired-login deletion and resulting admitted/unknown counts
    /// before heap changes. Empty work skips the callback. Failure preserves all
    /// records but retains the trusted clock observation, as other refusals do.
    /// Hosts persist clock/authority metadata separately even when no rows expire.
    pub fn prune_with_commit(
        &mut self,
        now_ns: u64,
        budget: usize,
        commit: impl FnOnce(&[Principal], (usize, usize)) -> Result<(), LoginRevocationError>,
    ) -> Result<usize, LoginRevocationError> {
        self.clock(now_ns)?;
        let callers: Vec<_> = self
            .expiry
            .iter()
            .take_while(|(until, _)| *until <= now_ns)
            .take(budget.min(128))
            .map(|(_, caller)| *caller)
            .collect();
        if callers.is_empty() {
            return Ok(0);
        }
        let mut counts = self.occupancy();
        for caller in &callers {
            if self.records[caller].admitted_pool() {
                counts.0 -= 1;
            } else {
                counts.1 -= 1;
            }
        }
        commit(&callers, counts)?;
        for caller in &callers {
            self.remove(*caller);
        }
        Ok(callers.len())
    }

    fn reservation(
        &mut self,
        caller: Principal,
        claims: &PortableClaims,
        now_ns: u64,
    ) -> Result<Record, LoginRevocationError> {
        self.observe(caller, now_ns)?;
        validate_shape(claims).map_err(|_| LoginRevocationError::InvalidPolicy)?;
        self.reservation_checked(caller, claims, now_ns)
    }

    // Only canonical decoded or verified claims can skip static shape checks.
    // Caller/clock and all mutable ledger state are checked on every operation.
    fn reservation_checked(
        &self,
        caller: Principal,
        claims: &PortableClaims,
        now_ns: u64,
    ) -> Result<Record, LoginRevocationError> {
        let login_deadline_ns = claims.login_deadline_ns;
        if claims.audience != self.audience
            || browser_principal(&claims.browser_public_key) != caller
        {
            return Err(LoginRevocationError::InvalidPolicy);
        }
        let until_ns = login_deadline_ns
            .checked_add(self.skew_ns)
            .ok_or(LoginRevocationError::Clock)?;
        let login = LoginBinding {
            issuer: claims.issuer,
            subject: claims.account_subject,
            started_at_ns: claims.login_started_at_ns,
            deadline_ns: claims.login_deadline_ns,
        };
        if let Some(old) = self.records.get(&caller).filter(|r| r.until_ns > now_ns) {
            if old.revoked() {
                return Err(LoginRevocationError::Revoked);
            }
            if old.login() != Some(login) {
                return Err(LoginRevocationError::InvalidPolicy);
            }
        }
        let record = Record {
            until_ns,
            state: LoginRecordState::Reserved(login),
        };
        self.capacity(caller, record)?;
        Ok(record)
    }

    fn clock(&mut self, now_ns: u64) -> Result<(), LoginRevocationError> {
        if now_ns < self.observed_ns {
            return Err(LoginRevocationError::Clock);
        }
        self.observed_ns = now_ns;
        Ok(())
    }

    fn observe(&mut self, caller: Principal, now_ns: u64) -> Result<(), LoginRevocationError> {
        if !is_login_caller(caller) {
            return Err(LoginRevocationError::InvalidCaller);
        }
        self.clock(now_ns)
    }

    fn capacity(&self, caller: Principal, record: Record) -> Result<(), LoginRevocationError> {
        let reuses_pool = self
            .records
            .get(&caller)
            .is_some_and(|old| old.admitted_pool() == record.admitted_pool());
        let (count, limit) = if record.admitted_pool() {
            (self.admitted, self.max_admitted)
        } else {
            (self.unknown, self.max_unknown)
        };
        if !reuses_pool && count >= limit {
            return Err(LoginRevocationError::Capacity);
        }
        Ok(())
    }

    fn remove(&mut self, caller: Principal) {
        if let Some(old) = self.records.remove(&caller) {
            self.expiry.remove(&(old.until_ns, caller));
            if old.admitted_pool() {
                self.admitted -= 1;
            } else {
                self.unknown -= 1;
            }
        }
    }

    fn install(&mut self, caller: Principal, record: Record) {
        self.remove(caller);
        if record.admitted_pool() {
            self.admitted += 1;
        } else {
            self.unknown += 1;
        }
        self.expiry.insert((record.until_ns, caller));
        self.records.insert(caller, record);
    }
}

fn retention(now_ns: u64, skew_ns: u64) -> Result<u64, LoginRevocationError> {
    skew_ns
        .checked_mul(2)
        .and_then(|skew| crate::limits::MAX_LOGIN_NS.checked_add(skew))
        .and_then(|ttl| now_ns.checked_add(ttl))
        .ok_or(LoginRevocationError::Clock)
}