arcly-http 0.2.2

//! Multi-tenant separation layer.
//!
//! Tenant identity is resolved **once per request** inside
//! `web::boundary::assemble_context` — the same single pipeline that handles
//! credentials and tracing — so every entry point (macro routes, plugin
//! routes) sees the same tenant with zero per-handler boilerplate.
//!
//! ## Hot-path guarantees
//!
//! `TenantRegistry` is a *frozen* map built once at boot and provided into the
//! DI container: resolution is a header read + immutable `HashMap` lookup —
//! no locks, no allocation beyond the interned `SmolStr` key.
//!
//! ## Defense in depth
//!
//! A client-supplied header is never trusted on its own. When the request also
//! carries a JWT with a `tenant` claim, [`TenantGuard`] enforces that both
//! agree — a token minted for tenant A can never act on tenant B by forging
//! `X-Tenant-Id`.
//!
//! ## Usage
//!
//! ```ignore
//! // boot (plugin on_init):
//! ctx.provide(TenantRegistry::new(
//!     TenantStrategy::Header("x-tenant-id"),
//!     vec![
//!         TenantConfig::new("acme",  "Acme Corp", "acme_primary"),
//!         TenantConfig::new("globex", "Globex",    "globex_primary"),
//!     ],
//!     Some(TenantId::new("public")),       // fallback tenant (None = strict)
//! ));
//!
//! // handler:
//! static TENANT: TenantGuard = TenantGuard;
//! TENANT.check(&ctx)?;
//! let t = ctx.tenant().expect("guarded");
//! ```

use std::collections::{HashMap, HashSet};
use std::sync::Arc;

use arc_swap::ArcSwap;
use axum::http::HeaderMap;
use smol_str::SmolStr;

use crate::auth::guards::Guard;
use crate::web::{Error, RequestContext};

/// Interned tenant identifier (no heap allocation for IDs ≤ 23 bytes).
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct TenantId(pub SmolStr);

impl TenantId {
    pub fn new(id: impl AsRef<str>) -> Self {
        Self(SmolStr::new(id.as_ref()))
    }
    pub fn as_str(&self) -> &str {
        &self.0
    }
}

/// How the tenant identifier is carried on the wire.
#[non_exhaustive]
pub enum TenantStrategy {
    /// e.g. `X-Tenant-Id: acme`
    Header(std::borrow::Cow<'static, str>),
    /// e.g. `acme.api.example.com` with `base_domain = "api.example.com"` → `acme`
    Subdomain {
        base_domain: std::borrow::Cow<'static, str>,
    },
    /// Header takes precedence; subdomain is the fallback.
    HeaderThenSubdomain {
        header: std::borrow::Cow<'static, str>,
        base_domain: std::borrow::Cow<'static, str>,
    },
}

impl TenantStrategy {
    /// Header strategy from a static or runtime-config header name.
    pub fn header(name: impl Into<std::borrow::Cow<'static, str>>) -> Self {
        Self::Header(name.into())
    }
    pub fn subdomain(base_domain: impl Into<std::borrow::Cow<'static, str>>) -> Self {
        Self::Subdomain {
            base_domain: base_domain.into(),
        }
    }
}

/// Per-tenant static configuration.
#[non_exhaustive]
pub struct TenantConfig {
    pub id: TenantId,
    pub display_name: String,
    /// Logical datasource name — consumed by `data::DataSourceRegistry` to
    /// select this tenant's connection pool. Loose coupling by name keeps the
    /// tenant layer independent of any database crate.
    pub datasource: String,
}

impl TenantConfig {
    pub fn new(
        id: impl AsRef<str>,
        display_name: impl Into<String>,
        datasource: impl Into<String>,
    ) -> Self {
        Self {
            id: TenantId::new(id),
            display_name: display_name.into(),
            datasource: datasource.into(),
        }
    }
}

/// Tenant registry with **dynamic provisioning**.
///
/// Hot path: `resolve()` is one `ArcSwap` pointer load + an immutable map
/// read — exactly the zero-lock profile of the old frozen map. Control
/// plane: [`Self::upsert`]/[`Self::suspend`]/[`Self::resume`] copy-on-write a new snapshot and
/// swap it atomically, so onboarding a customer (or cutting off a delinquent
/// one) takes effect on the next request — **no redeploy, no restart**.
///
/// Not `#[Injectable]` — provide via `ctx.provide(TenantRegistry::new(...))`.
pub struct TenantRegistry {
    strategy: TenantStrategy,
    snapshot: ArcSwap<TenantSnapshot>,
    fallback: Option<TenantId>,
}

struct TenantSnapshot {
    known: HashMap<TenantId, Arc<TenantConfig>>,
    suspended: HashSet<TenantId>,
}

impl TenantRegistry {
    pub fn new(
        strategy: TenantStrategy,
        tenants: Vec<TenantConfig>,
        fallback: Option<TenantId>,
    ) -> Self {
        let known = tenants
            .into_iter()
            .map(|t| (t.id.clone(), Arc::new(t)))
            .collect();
        Self {
            strategy,
            snapshot: ArcSwap::from_pointee(TenantSnapshot {
                known,
                suspended: HashSet::new(),
            }),
            fallback,
        }
    }

    // ── Control plane (copy-on-write; never touches the hot path) ──────────

    /// Add or replace a tenant — visible to the very next request.
    pub fn upsert(&self, cfg: TenantConfig) {
        self.snapshot.rcu(|cur| {
            let mut known = cur.known.clone();
            known.insert(
                cfg.id.clone(),
                Arc::new(TenantConfig {
                    id: cfg.id.clone(),
                    display_name: cfg.display_name.clone(),
                    datasource: cfg.datasource.clone(),
                }),
            );
            TenantSnapshot {
                known,
                suspended: cur.suspended.clone(),
            }
        });
        tracing::info!("tenant upserted (live, no restart)");
    }

    /// Suspend a tenant: `resolve()` returns `None` → guards reject with 401.
    /// Deliberately does NOT fall through to the fallback tenant — a
    /// suspended customer must not silently ride the shared pool.
    pub fn suspend(&self, id: &TenantId) {
        self.snapshot.rcu(|cur| {
            let mut suspended = cur.suspended.clone();
            suspended.insert(id.clone());
            TenantSnapshot {
                known: cur.known.clone(),
                suspended,
            }
        });
    }

    pub fn resume(&self, id: &TenantId) {
        self.snapshot.rcu(|cur| {
            let mut suspended = cur.suspended.clone();
            suspended.remove(id);
            TenantSnapshot {
                known: cur.known.clone(),
                suspended,
            }
        });
    }

    fn raw_id(&self, headers: &HeaderMap) -> Option<SmolStr> {
        fn from_header(headers: &HeaderMap, name: &str) -> Option<SmolStr> {
            headers
                .get(name)
                .and_then(|v| v.to_str().ok())
                .map(str::trim)
                .filter(|s| !s.is_empty())
                .map(SmolStr::new)
        }
        fn from_subdomain(headers: &HeaderMap, base: &str) -> Option<SmolStr> {
            let host = headers.get("host")?.to_str().ok()?;
            // Strip an optional port before comparing against the base domain.
            let host = host.rsplit_once(':').map_or(host, |(h, _)| h);
            host.strip_suffix(base)?
                .strip_suffix('.')
                .filter(|s| !s.is_empty())
                .map(SmolStr::new)
        }

        match &self.strategy {
            TenantStrategy::Header(h) => from_header(headers, h),
            TenantStrategy::Subdomain { base_domain } => from_subdomain(headers, base_domain),
            TenantStrategy::HeaderThenSubdomain {
                header,
                base_domain,
            } => from_header(headers, header).or_else(|| from_subdomain(headers, base_domain)),
        }
    }

    /// Resolve the request's tenant. Called once per request by the boundary.
    ///
    /// One atomic snapshot load + immutable map reads — zero locks.
    /// Unknown IDs resolve to the fallback (a typo'd header cannot select an
    /// unconfigured tenant); **suspended** IDs resolve to `None` outright.
    pub fn resolve(&self, headers: &HeaderMap) -> Option<Arc<TenantConfig>> {
        let snap = self.snapshot.load();

        if let Some(id) = self.raw_id(headers).map(TenantId) {
            if snap.suspended.contains(&id) {
                return None; // hard cut-off, no fallback
            }
            if let Some(cfg) = snap.known.get(&id) {
                return Some(cfg.clone()); // Arc clone — no data copy
            }
        }
        self.fallback
            .as_ref()
            .filter(|id| !snap.suspended.contains(id))
            .and_then(|id| snap.known.get(id).cloned())
    }

    pub fn get(&self, id: &TenantId) -> Option<Arc<TenantConfig>> {
        self.snapshot.load().known.get(id).cloned()
    }

    /// Resolve a tenant by its raw id — the non-HTTP twin of
    /// [`resolve`](Self::resolve), used by the consumer mesh where the id
    /// rides a message envelope instead of a header. Same semantics:
    /// **suspended → `None` outright** (a suspended tenant's queued events
    /// must stop processing, exactly like its HTTP traffic). Unknown ids get
    /// no fallback: an envelope names its tenant explicitly, and falling
    /// back would process data under the wrong tenant.
    pub fn resolve_by_id(&self, id: &str) -> Option<Arc<TenantConfig>> {
        let snap = self.snapshot.load();
        let id = TenantId::new(id);
        if snap.suspended.contains(&id) {
            return None;
        }
        snap.known.get(&id).cloned()
    }
}

// ─── TenantGuard ──────────────────────────────────────────────────────────────

/// Requires a resolved tenant AND cross-checks it against the JWT.
///
/// - No tenant on the context → `401 Unauthorized` (strategy failed to
///   resolve and no fallback is configured).
/// - JWT carries a `tenant` claim that differs from the resolved tenant →
///   `403 Forbidden` (forged header / cross-tenant token reuse).
/// - JWT without a `tenant` claim passes — single-tenant tokens stay valid.
pub struct TenantGuard;

/// Ready-to-use singleton.
pub static TENANT: TenantGuard = TenantGuard;

impl Guard for TenantGuard {
    fn check(&self, ctx: &RequestContext) -> Result<(), Error> {
        let tenant = ctx.tenant().ok_or(Error::Unauthorized)?;

        if let Some(claim) = ctx
            .claims()
            .and_then(|c| c.get("tenant"))
            .and_then(|v| v.as_str())
        {
            if claim != tenant.id.as_str() {
                return Err(Error::Forbidden);
            }
        }
        Ok(())
    }
}