scout_rs 0.95.95

//! Pub/sub JWT verification and minting.

use std::collections::HashMap;
use std::fmt;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};

use anyhow::{anyhow, Context, Result};
use jsonwebtoken::{decode, decode_header, Algorithm, DecodingKey, Header, Validation};
use serde::Deserialize;

use crate::constants::{env, rpc};
use crate::db_client::DatabaseConfig;

#[derive(Clone)]
pub struct PubsubTokenConfig {
    key_dir: Option<KeyDir>,
    cache: Arc<Mutex<HashMap<String, VerifyingKey>>>,
}

#[derive(Clone)]
enum KeyDir {
    Local(PathBuf),
    Remote(String),
}

#[derive(Clone)]
struct VerifyingKey {
    algorithm: Algorithm,
    decoding_key: DecodingKey,
}

impl PubsubTokenConfig {
    pub fn from_env() -> Result<Self> {
        dotenv::dotenv().ok();
        let key_dir = read_key_dir_from_env()?
            .ok_or_else(|| anyhow!("{} is required", env::SCOUT_PUBSUB_JWT_KEY_DIR))?;
        let config = Self {
            key_dir: Some(key_dir),
            cache: Arc::new(Mutex::new(HashMap::new())),
        };
        if let Some(KeyDir::Local(path)) = &config.key_dir {
            config.refresh_key_dir_local(path)?;
        }
        Ok(config)
    }

    /// Load verifier keys from public JWK JSON (no `d`).
    pub fn from_public_jwks(jwks: &[&str]) -> Result<Self> {
        let mut cache = HashMap::new();
        for jwk in jwks {
            let meta: JwkKid = serde_json::from_str(jwk)?;
            let kid = meta
                .kid
                .ok_or_else(|| anyhow!("public JWK must include kid"))?;
            cache.insert(kid, public_verifying_key_from_jwk(jwk)?);
        }
        Ok(Self {
            key_dir: None,
            cache: Arc::new(Mutex::new(cache)),
        })
    }

    /// Load `{kid}.jwk` files from a local directory.
    pub fn from_local_key_dir(path: impl AsRef<Path>) -> Result<Self> {
        let path = path.as_ref().to_path_buf();
        let config = Self {
            key_dir: Some(KeyDir::Local(path.clone())),
            cache: Arc::new(Mutex::new(HashMap::new())),
        };
        config.refresh_key_dir_local(&path)?;
        Ok(config)
    }

    pub fn refresh_key_dir(&self) -> Result<()> {
        if let Some(KeyDir::Local(path)) = &self.key_dir {
            self.refresh_key_dir_local(path)?;
        }
        Ok(())
    }

    pub fn sync_public_keys_from_database(&self, db: &DatabaseConfig) -> Result<()> {
        let keys = fetch_pubsub_jwt_public_keys(db)?;
        if let Some(KeyDir::Local(path)) = &self.key_dir {
            sync_pubsub_public_keys_to_dir(&keys, path)?;
        }
        self.load_public_keys_into_cache(&keys)
    }

    fn load_public_keys_into_cache(&self, keys: &[PubsubJwkPublicKey]) -> Result<()> {
        let mut loaded = HashMap::new();
        for key in keys {
            if !valid_kid(&key.kid) {
                continue;
            }
            let jwk_json = serde_json::to_string(&key.jwk)?;
            loaded.insert(key.kid.clone(), public_verifying_key_from_jwk(&jwk_json)?);
        }
        *self.cache.lock().expect("pubsub key cache") = loaded;
        Ok(())
    }

    fn refresh_key_dir_local(&self, dir: &Path) -> Result<()> {
        let mut loaded = HashMap::new();
        for entry in
            std::fs::read_dir(dir).with_context(|| format!("read key dir {}", dir.display()))?
        {
            let path = entry?.path();
            if path.extension().and_then(|e| e.to_str()) != Some("jwk") {
                continue;
            }
            let Some(kid) = path.file_stem().and_then(|s| s.to_str()) else {
                continue;
            };
            if !valid_kid(kid) {
                continue;
            }
            let jwk = std::fs::read_to_string(&path)?;
            loaded.insert(kid.to_string(), public_verifying_key_from_jwk(&jwk)?);
        }
        *self.cache.lock().expect("pubsub key cache") = loaded;
        Ok(())
    }

    fn resolve_key(&self, header: &Header) -> Result<VerifyingKey, PubsubTokenError> {
        let kid = header.kid.as_deref().ok_or(PubsubTokenError::Invalid)?;

        if let Some(key) = self
            .cache
            .lock()
            .expect("pubsub key cache")
            .get(kid)
            .cloned()
        {
            return matching_alg(key, header.alg);
        }

        let key_dir = self.key_dir.as_ref().ok_or(PubsubTokenError::Invalid)?;
        let key = fetch_public_key(key_dir, kid).map_err(|_| PubsubTokenError::Invalid)?;
        matching_alg(
            {
                self.cache
                    .lock()
                    .expect("pubsub key cache")
                    .insert(kid.to_string(), key.clone());
                key
            },
            header.alg,
        )
    }
}

fn matching_alg(key: VerifyingKey, alg: Algorithm) -> Result<VerifyingKey, PubsubTokenError> {
    if key.algorithm == alg {
        Ok(key)
    } else {
        Err(PubsubTokenError::Invalid)
    }
}

#[derive(Debug, Clone, Deserialize)]
pub struct PubsubJwkPublicKey {
    pub kid: String,
    pub jwk: serde_json::Value,
}

pub fn fetch_pubsub_jwt_public_keys(db: &DatabaseConfig) -> Result<Vec<PubsubJwkPublicKey>> {
    let api_key = db
        .service_role_key
        .as_deref()
        .filter(|_| db.use_service_key)
        .unwrap_or_else(|| db.get_supabase_api_key());

    let url = format!(
        "{}/rpc/{}",
        db.get_rest_url().trim_end_matches('/'),
        rpc::GET_PUBSUB_JWT_PUBLIC_KEYS
    );

    let response = reqwest::blocking::Client::new()
        .post(url)
        .header("apikey", api_key)
        .header("Authorization", format!("Bearer {api_key}"))
        .header("Content-Type", "application/json")
        .body("{}")
        .send()
        .context("get_pubsub_jwt_public_keys request")?;

    let status = response.status();
    let body = response
        .text()
        .context("get_pubsub_jwt_public_keys response body")?;

    if !status.is_success() {
        return Err(anyhow!(
            "get_pubsub_jwt_public_keys failed: HTTP {status} - {body}"
        ));
    }

    serde_json::from_str(&body).context("parse get_pubsub_jwt_public_keys response")
}

#[derive(Debug, Clone, Deserialize)]
pub struct PubsubTokenMint {
    pub token: String,
    pub put: Vec<String>,
    pub get: Vec<String>,
    pub iat: i64,
    pub exp: i64,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum MintPubsubTokenError {
    MissingDeviceApiKey,
    MissingSupabaseApiKey,
    MissingServiceRoleKey,
    MissingHerdScope,
    Unauthorized,
    Forbidden,
    Http { status: u16, body: String },
    InvalidResponse,
}

impl fmt::Display for MintPubsubTokenError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            MintPubsubTokenError::MissingDeviceApiKey => {
                write!(f, "device API key required to mint pubsub token")
            }
            MintPubsubTokenError::MissingSupabaseApiKey => {
                write!(f, "supabase API key required to mint pubsub token")
            }
            MintPubsubTokenError::MissingServiceRoleKey => {
                write!(f, "service role key required to mint service pubsub token")
            }
            MintPubsubTokenError::MissingHerdScope => {
                write!(
                    f,
                    "at least one herd id required to mint service pubsub token"
                )
            }
            MintPubsubTokenError::Unauthorized => write!(f, "unauthorized"),
            MintPubsubTokenError::Forbidden => write!(f, "forbidden"),
            MintPubsubTokenError::Http { status, body } => {
                write!(f, "mint-pubsub-token failed: HTTP {status} - {body}")
            }
            MintPubsubTokenError::InvalidResponse => {
                write!(f, "mint-pubsub-token returned invalid response")
            }
        }
    }
}

pub fn mint_pubsub_token(db: &DatabaseConfig) -> Result<PubsubTokenMint, MintPubsubTokenError> {
    let device_api_key = db.get_scout_api_key();
    if device_api_key.is_empty() {
        return Err(MintPubsubTokenError::MissingDeviceApiKey);
    }
    let supabase_api_key = db.get_supabase_api_key();
    if supabase_api_key.is_empty() {
        return Err(MintPubsubTokenError::MissingSupabaseApiKey);
    }

    let url = format!(
        "{}/mint-pubsub-token",
        db.get_functions_url()
            .map_err(|e| MintPubsubTokenError::Http {
                status: 0,
                body: e.to_string(),
            })?
    );

    let response = reqwest::blocking::Client::new()
        .post(url)
        .header("apikey", supabase_api_key)
        .header("Authorization", format!("Bearer {supabase_api_key}"))
        .header("api_key", device_api_key)
        .send()
        .map_err(|e| MintPubsubTokenError::Http {
            status: 0,
            body: e.to_string(),
        })?;

    parse_mint_response(response)
}

/// Mints a publish-only pub/sub token for a trusted service (e.g. a relay) via the
/// `mint-pubsub-token` edge function. Requires a service-role key on `db` and an explicit
/// herd scope; the edge function resolves claims through `get_service_pubsub_token_claims`.
pub fn mint_service_pubsub_token(
    db: &DatabaseConfig,
    herd_ids: &[i64],
) -> Result<PubsubTokenMint, MintPubsubTokenError> {
    let service_role_key = db
        .service_role_key
        .as_deref()
        .filter(|k| !k.is_empty())
        .ok_or(MintPubsubTokenError::MissingServiceRoleKey)?;

    if herd_ids.is_empty() {
        return Err(MintPubsubTokenError::MissingHerdScope);
    }

    let url = format!(
        "{}/mint-pubsub-token",
        db.get_functions_url()
            .map_err(|e| MintPubsubTokenError::Http {
                status: 0,
                body: e.to_string(),
            })?
    );

    let body = serde_json::json!({ "herd_ids": herd_ids }).to_string();

    let response = reqwest::blocking::Client::new()
        .post(url)
        .header("apikey", service_role_key)
        .header("Authorization", format!("Bearer {service_role_key}"))
        .header("Content-Type", "application/json")
        .body(body)
        .send()
        .map_err(|e| MintPubsubTokenError::Http {
            status: 0,
            body: e.to_string(),
        })?;

    parse_mint_response(response)
}

fn parse_mint_response(
    response: reqwest::blocking::Response,
) -> Result<PubsubTokenMint, MintPubsubTokenError> {
    let status = response.status();
    let body = response.text().map_err(|e| MintPubsubTokenError::Http {
        status: status.as_u16(),
        body: e.to_string(),
    })?;

    match status.as_u16() {
        401 => Err(MintPubsubTokenError::Unauthorized),
        403 => Err(MintPubsubTokenError::Forbidden),
        _ if !status.is_success() => Err(MintPubsubTokenError::Http {
            status: status.as_u16(),
            body,
        }),
        _ => {
            let mint: PubsubTokenMint =
                serde_json::from_str(&body).map_err(|_| MintPubsubTokenError::InvalidResponse)?;
            if mint.token.is_empty() {
                Err(MintPubsubTokenError::InvalidResponse)
            } else {
                Ok(mint)
            }
        }
    }
}

/// Writes `{kid}.jwk` files and removes local `*.jwk` files not in `keys`.
pub fn sync_pubsub_public_keys_to_dir(keys: &[PubsubJwkPublicKey], dir: &Path) -> Result<()> {
    std::fs::create_dir_all(dir).with_context(|| format!("create key dir {}", dir.display()))?;

    let active: std::collections::HashSet<&str> = keys.iter().map(|k| k.kid.as_str()).collect();

    for key in keys {
        if !valid_kid(&key.kid) {
            return Err(anyhow!(
                "invalid kid in get_pubsub_jwt_public_keys: {}",
                key.kid
            ));
        }
        let jwk: Jwk = serde_json::from_value(key.jwk.clone())?;
        if jwk.private_scalar.is_some() {
            return Err(anyhow!("public JWK must not include d (kid={})", key.kid));
        }
        let path = dir.join(format!("{}.jwk", key.kid));
        std::fs::write(&path, serde_json::to_string(&key.jwk)?)
            .with_context(|| format!("write {}", path.display()))?;
    }

    for entry in
        std::fs::read_dir(dir).with_context(|| format!("read key dir {}", dir.display()))?
    {
        let path = entry?.path();
        if path.extension().and_then(|e| e.to_str()) != Some("jwk") {
            continue;
        }
        let Some(stem) = path.file_stem().and_then(|s| s.to_str()) else {
            continue;
        };
        if valid_kid(stem) && !active.contains(stem) {
            std::fs::remove_file(&path)
                .with_context(|| format!("remove stale key {}", path.display()))?;
        }
    }

    Ok(())
}

#[derive(Debug, Clone)]
pub struct PubsubTokenClaims {
    pub root: String,
    pub publish: Vec<String>,
    pub subscribe: Vec<String>,
    pub issued_at: Option<i64>,
    pub expires_at: Option<i64>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum PubsubTokenError {
    Invalid,
    Expired,
}

impl fmt::Display for PubsubTokenError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            PubsubTokenError::Invalid => write!(f, "invalid pubsub token"),
            PubsubTokenError::Expired => write!(f, "expired pubsub token"),
        }
    }
}

#[derive(Debug, Deserialize)]
struct Jwk {
    kty: String,
    #[serde(rename = "crv")]
    curve: Option<String>,
    alg: Option<String>,
    x: Option<String>,
    y: Option<String>,
    n: Option<String>,
    e: Option<String>,
    #[serde(rename = "d")]
    private_scalar: Option<String>,
}

#[derive(Debug, Deserialize)]
struct JwtClaims {
    #[serde(default)]
    root: String,
    #[serde(default, rename = "put")]
    put: ClaimPaths,
    #[serde(default, rename = "get")]
    get: ClaimPaths,
    iat: Option<i64>,
    exp: Option<i64>,
}

#[derive(Debug, Default, Deserialize)]
#[serde(untagged)]
enum ClaimPaths {
    #[default]
    Absent,
    One(String),
    Many(Vec<String>),
}

impl ClaimPaths {
    fn into_vec(self) -> Vec<String> {
        match self {
            ClaimPaths::Absent => vec![],
            ClaimPaths::One(s) if s.is_empty() => vec![],
            ClaimPaths::One(s) => vec![s],
            ClaimPaths::Many(v) => v,
        }
    }
}

fn read_key_dir_from_env() -> Result<Option<KeyDir>> {
    let Ok(raw) = std::env::var(env::SCOUT_PUBSUB_JWT_KEY_DIR) else {
        return Ok(None);
    };
    let trimmed = raw.trim();
    if trimmed.is_empty() {
        return Ok(None);
    }
    if trimmed.starts_with("http://") || trimmed.starts_with("https://") {
        Ok(Some(KeyDir::Remote(
            trimmed.trim_end_matches('/').to_string(),
        )))
    } else {
        Ok(Some(KeyDir::Local(PathBuf::from(trimmed))))
    }
}

fn valid_kid(kid: &str) -> bool {
    !kid.is_empty()
        && kid
            .chars()
            .all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_')
}

fn fetch_public_key(key_dir: &KeyDir, kid: &str) -> Result<VerifyingKey> {
    if !valid_kid(kid) {
        return Err(anyhow!("invalid key id"));
    }
    let jwk = match key_dir {
        KeyDir::Local(dir) => std::fs::read_to_string(dir.join(format!("{kid}.jwk")))?,
        KeyDir::Remote(base) => {
            let url = format!("{base}/{kid}.jwk");
            reqwest::blocking::get(&url)?
                .error_for_status()?
                .text()
                .with_context(|| format!("GET {url}"))?
        }
    };
    public_verifying_key_from_jwk(&jwk)
}

fn public_verifying_key_from_jwk(jwk_json: &str) -> Result<VerifyingKey> {
    let jwk: Jwk = serde_json::from_str(jwk_json.trim())?;
    if jwk.private_scalar.is_some() {
        return Err(anyhow!("public JWK must not include d"));
    }

    match jwk.kty.as_str() {
        "EC" => {
            let crv = jwk.curve.as_deref().unwrap_or("P-256");
            let algorithm = match crv {
                "P-256" => Algorithm::ES256,
                "P-384" => Algorithm::ES384,
                _ => return Err(anyhow!("unsupported EC curve: {crv}")),
            };
            Ok(VerifyingKey {
                algorithm,
                decoding_key: DecodingKey::from_ec_components(
                    jwk.x.as_ref().context("EC JWK missing x")?,
                    jwk.y.as_ref().context("EC JWK missing y")?,
                )?,
            })
        }
        "RSA" => {
            let algorithm = match jwk.alg.as_deref() {
                Some("RS384") => Algorithm::RS384,
                Some("RS512") => Algorithm::RS512,
                _ => Algorithm::RS256,
            };
            Ok(VerifyingKey {
                algorithm,
                decoding_key: DecodingKey::from_rsa_components(
                    jwk.n.as_ref().context("RSA JWK missing n")?,
                    jwk.e.as_ref().context("RSA JWK missing e")?,
                )?,
            })
        }
        "OKP" => {
            if jwk.curve.as_deref() != Some("Ed25519") {
                return Err(anyhow!("unsupported OKP curve"));
            }
            Ok(VerifyingKey {
                algorithm: Algorithm::EdDSA,
                decoding_key: DecodingKey::from_ed_components(
                    jwk.x.as_ref().context("OKP JWK missing x")?,
                )?,
            })
        }
        other => Err(anyhow!("unsupported JWK kty: {other}")),
    }
}

pub fn verify_pubsub_token(
    config: &PubsubTokenConfig,
    token: &str,
) -> Result<PubsubTokenClaims, PubsubTokenError> {
    let token = token.trim();
    let header = decode_header(token).map_err(|_| PubsubTokenError::Invalid)?;
    let key = config.resolve_key(&header)?;

    let mut validation = Validation::new(key.algorithm);
    validation.validate_exp = true;
    validation.required_spec_claims.clear();

    let data = decode::<JwtClaims>(token, &key.decoding_key, &validation).map_err(|e| {
        if matches!(e.kind(), jsonwebtoken::errors::ErrorKind::ExpiredSignature) {
            PubsubTokenError::Expired
        } else {
            PubsubTokenError::Invalid
        }
    })?;

    Ok(PubsubTokenClaims {
        root: data.claims.root,
        publish: data.claims.put.into_vec(),
        subscribe: data.claims.get.into_vec(),
        issued_at: data.claims.iat,
        expires_at: data.claims.exp,
    })
}

pub fn can_publish(claims: &PubsubTokenClaims, path: &str) -> bool {
    path_allowed(&claims.root, &claims.publish, path)
}

pub fn can_subscribe(claims: &PubsubTokenClaims, path: &str) -> bool {
    path_allowed(&claims.root, &claims.subscribe, path)
}

fn path_allowed(root: &str, patterns: &[String], path: &str) -> bool {
    if patterns.is_empty() {
        return false;
    }
    let path = normalize_path(path);
    patterns
        .iter()
        .any(|p| pattern_matches_path(root, p, &path))
}

fn pattern_matches_path(root: &str, pattern: &str, path: &str) -> bool {
    let prefix = normalize_pattern(&join_root(root, pattern));
    if prefix.is_empty() {
        return true;
    }
    path == prefix || path.starts_with(&format!("{prefix}/"))
}

fn join_root(root: &str, pattern: &str) -> String {
    let root = normalize_path(root);
    let pattern = pattern.trim().trim_start_matches('/');
    if pattern.is_empty() {
        return root;
    }
    if root.is_empty() {
        return normalize_pattern(pattern);
    }
    format!("{}/{}", root, normalize_pattern(pattern))
}

fn normalize_pattern(s: &str) -> String {
    let s = normalize_path(s);
    s.strip_suffix("/*").unwrap_or(&s).to_string()
}

fn normalize_path(s: &str) -> String {
    s.trim()
        .trim_matches('/')
        .split('/')
        .filter(|p| !p.is_empty())
        .collect::<Vec<_>>()
        .join("/")
}

#[derive(Debug, Deserialize)]
struct JwkKid {
    kid: Option<String>,
}