bitrouter 1.0.0-alpha.31

//! `AuthHook` — the `language_model::PreRequestHook` that authenticates a
//! request against a `brvk_` virtual key.
//!
//! v1 has **no JWT path**: the only credential form is a virtual key,
//! looked up by SHA-256 hash in the `api_keys` table.
//!
//! Relationship to `server.skip_auth`: `skip_auth` is an SDK-level flag
//! handled at the server entry — when it is on, the server synthesises
//! a *local* `CallerContext` for every inbound request, **regardless of
//! any credential header**. The intent is "local-first, fully open" —
//! the same posture the zero-config startup uses to make tools like
//! codex / Claude Code / litellm just work without requiring the user
//! to mint a virtual key first. The four-way truth table:
//!
//! | skip_auth | credential | result                                |
//! |-----------|------------|---------------------------------------|
//! | false     | present    | validated (Allow / Deny)              |
//! | false     | absent     | Deny 401                              |
//! | true      | present    | Allow — any header accepted as local  |
//! | true      | absent     | Allow — local caller passes through   |
//!
//! Previously the `true × present` row also validated — which silently
//! broke Claude Code and litellm because both auto-inject an
//! `Authorization: Bearer …` header that bitrouter saw as a malformed
//! virtual key. Multi-tenant operators set `skip_auth: false` (the SDK
//! default) to get real validation.

use async_trait::async_trait;
use chrono::Utc;
use sea_orm::DatabaseConnection;

use bitrouter_sdk::caller::CallerContext;
use bitrouter_sdk::language_model::{DenyReason, HookDecision, PipelineContext, PreRequestHook};
use bitrouter_sdk::mcp::{McpContext, PreRequestHook as McpPreRequestHook};
use bitrouter_sdk::{PluginId, Result};
use http::HeaderMap;

use crate::auth::db::{self, ApiKeyRecord};
use crate::auth::events::{ApiPrincipalEstablished, Authenticated};
use crate::auth::keys;

/// The auth module id, used as the `PipelineContext` metadata key. The
/// string is preserved as `bitrouter-auth` so policy code that reads
/// metadata under that key continues to work after the move from a
/// shared plugin into a binary module.
pub fn plugin_id() -> PluginId {
    PluginId::new("bitrouter-auth")
}

/// Authenticates a request against the `api_keys` table (a `brvk_`
/// virtual key). Owns no routing or settlement behaviour — it only
/// establishes identity.
pub struct AuthHook {
    db: DatabaseConnection,
}

enum Authentication {
    Local,
    Authenticated {
        caller: CallerContext,
        record: ApiKeyRecord,
        route_scope_id: String,
    },
    Denied(String),
}

impl AuthHook {
    /// Build an `AuthHook` over a database connection. The database must
    /// already carry this module's tables (`crate::db::run_migrations`).
    pub fn new(db: DatabaseConnection) -> Self {
        Self { db }
    }

    /// Turn a validated key record into a `CallerContext`.
    fn caller_from_record(record: &ApiKeyRecord) -> CallerContext {
        CallerContext::new(&record.id, &record.user_id)
    }

    /// Authenticate headers for any protocol pipeline. The language-model and
    /// MCP hooks deliberately share this path so `server.skip_auth` and virtual
    /// key expiry/activation cannot drift between public surfaces.
    async fn authenticate(
        &self,
        headers: &HeaderMap,
        caller: &CallerContext,
    ) -> Result<Authentication> {
        if caller.is_local() {
            return Ok(Authentication::Local);
        }

        let Some(credential) = credential_from_headers(headers) else {
            return Ok(Authentication::Denied("missing API key".to_string()));
        };
        if !keys::looks_like_virtual_key(&credential) {
            return Ok(Authentication::Denied(
                "credential is not a brvk_ virtual key".to_string(),
            ));
        }

        let route_scope_id = keys::hash_key(&credential);
        let record = db::find_key_by_hash(&self.db, &route_scope_id).await?;
        let Some(record) = record else {
            return Ok(Authentication::Denied("unknown API key".to_string()));
        };
        if !record.active {
            return Ok(Authentication::Denied("API key is inactive".to_string()));
        }
        if let Some(expires_at) = record.expires_at
            && expires_at <= Utc::now()
        {
            return Ok(Authentication::Denied("API key has expired".to_string()));
        }

        Ok(Authentication::Authenticated {
            caller: Self::caller_from_record(&record),
            record,
            route_scope_id,
        })
    }
}

/// Shared credential extraction for native inference and app-owned control
/// plane endpoints.
pub(crate) fn credential_from_headers(headers: &HeaderMap) -> Option<String> {
    if let Some(auth) = headers
        .get("authorization")
        .and_then(|value| value.to_str().ok())
    {
        let token = auth.strip_prefix("Bearer ").unwrap_or(auth).trim();
        if !token.is_empty() {
            return Some(token.to_string());
        }
    }
    headers
        .get("x-api-key")
        .and_then(|value| value.to_str().ok())
        .map(str::trim)
        .filter(|value| !value.is_empty())
        .map(ToString::to_string)
}

#[async_trait]
impl PreRequestHook for AuthHook {
    async fn check(&self, ctx: &mut PipelineContext) -> Result<HookDecision> {
        match self.authenticate(ctx.headers(), ctx.caller()).await? {
            Authentication::Local => Ok(HookDecision::Allow),
            Authentication::Denied(message) => {
                Ok(HookDecision::Deny(DenyReason::Unauthorized(message)))
            }
            Authentication::Authenticated {
                caller,
                record,
                route_scope_id,
            } => {
                ctx.set_caller(caller);
                ctx.set_metadata(
                    &plugin_id(),
                    serde_json::json!({
                        "api_key_id": record.id,
                        "user_id": record.user_id,
                        "policy_id": record.policy_id,
                    }),
                );
                ctx.emit(Authenticated {
                    api_key_id: record.id,
                    user_id: record.user_id,
                    policy_id: record.policy_id,
                });
                ctx.emit(ApiPrincipalEstablished { route_scope_id });
                Ok(HookDecision::Allow)
            }
        }
    }
}

#[async_trait]
impl McpPreRequestHook for AuthHook {
    async fn check(&self, ctx: &mut McpContext) -> Result<HookDecision> {
        match self.authenticate(ctx.headers(), ctx.caller()).await? {
            Authentication::Local => Ok(HookDecision::Allow),
            Authentication::Denied(message) => {
                Ok(HookDecision::Deny(DenyReason::Unauthorized(message)))
            }
            Authentication::Authenticated {
                caller,
                record,
                route_scope_id,
            } => {
                ctx.set_caller(caller);
                ctx.emit(Authenticated {
                    api_key_id: record.id,
                    user_id: record.user_id,
                    policy_id: record.policy_id,
                });
                ctx.emit(ApiPrincipalEstablished { route_scope_id });
                Ok(HookDecision::Allow)
            }
        }
    }
}