graphar-flight 0.1.2

//! Authentication and authorization for the Cypher Flight SQL server.
//!
//! ## Authentication — *who are you?*
//!
//! Three methods, any combination, to cover the clients that matter:
//!
//! - **Bearer token** — a shared secret in `authorization: Bearer <token>`.
//!   More than one token can be valid at once (a **token set**) so a secret can
//!   be **rotated without downtime**: add the new token, let clients migrate,
//!   then drop the old one — both are accepted during the overlap window.
//! - **Basic (username / password)** — `authorization: Basic <base64>`. This
//!   is what generic Flight SQL clients on Windows use: Power BI's ADBC
//!   connector and the Arrow Flight SQL ODBC driver send username/password,
//!   first through the Flight `Handshake` (which issues a bearer token) and
//!   then on every call.
//! - **mutual TLS** — client-certificate auth, configured at the transport in
//!   [`crate::server`] (`TlsOptions::client_ca_pem`); orthogonal to the
//!   header methods here.
//!
//! The same [`AuthConfig`] drives both the per-call interceptor and the
//! handshake, so the two can never disagree about who is allowed in.
//!
//! ## Authorization — *what may you run?*
//!
//! Authentication answers *who*; the [`Authorizer`] answers *what*. After a
//! query is resolved to its registered table name, the authorizer is consulted
//! with the caller's [`Identity`] and may deny it (`permission_denied`). The
//! default is [`Authorizer::AllowAll`] — no behavior change — but
//! [`Authorizer::AllowList`] restricts each identity to a named subset of
//! registered queries, and [`Authorizer::Custom`] takes an arbitrary closure.

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

use base64::Engine;
use base64::engine::general_purpose::STANDARD as BASE64;
use tonic::Status;

/// A username / password pair accepted via Basic auth.
#[derive(Debug, Clone)]
pub struct BasicCredential {
    pub username: String,
    pub password: String,
}

/// The authenticated caller, as resolved from the `authorization` header.
///
/// Carried as a request extension by the interceptor (and returned by
/// [`AuthConfig::authenticate`]) so the per-query [`Authorizer`] can key its
/// decision on *who* is asking. The variants name the credential method:
///
/// - [`Identity::Token`] — a bearer token authenticated; the string is the
///   token's mapped name (from [`AuthConfig::with_named_bearer_tokens`]) or the
///   token value itself when unnamed.
/// - [`Identity::Basic`] — Basic auth authenticated; the string is the username.
/// - [`Identity::Anonymous`] — an open server with no auth configured.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum Identity {
    /// A bearer token authenticated; carries the token's identity name.
    Token(String),
    /// Basic auth authenticated; carries the username.
    Basic(String),
    /// No authentication configured (open server).
    Anonymous,
}

impl Identity {
    /// The identity's name for authz lookups: the token-identity name, the
    /// Basic username, or `"*"` for the anonymous (open-server) identity.
    pub fn name(&self) -> &str {
        match self {
            Identity::Token(name) | Identity::Basic(name) => name,
            Identity::Anonymous => "*",
        }
    }
}

/// Per-query authorization hook, consulted *after* a query resolves to its
/// registered name. Dependency-light by design — no external policy engine.
///
/// The default [`Authorizer::AllowAll`] permits everything (so wiring it in is
/// a no-op for existing servers). [`Authorizer::AllowList`] keys on
/// [`Identity::name`] and permits only the named queries listed for that
/// identity. [`Authorizer::Custom`] defers to a closure for arbitrary policy.
#[derive(Clone, Default)]
pub enum Authorizer {
    /// Permit every identity to run every query (the default).
    #[default]
    AllowAll,
    /// Permit each identity only the registered-query names listed for it.
    /// An identity absent from the map is denied every query; the special key
    /// `"*"` is consulted as a fallback for identities not otherwise listed.
    AllowList(HashMap<String, HashSet<String>>),
    /// Arbitrary policy: `(identity, resolved_query_name) -> allowed`.
    Custom(Arc<dyn Fn(&Identity, &str) -> bool + Send + Sync>),
}

impl std::fmt::Debug for Authorizer {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Authorizer::AllowAll => f.write_str("Authorizer::AllowAll"),
            Authorizer::AllowList(m) => f.debug_tuple("Authorizer::AllowList").field(m).finish(),
            Authorizer::Custom(_) => f.write_str("Authorizer::Custom(<fn>)"),
        }
    }
}

impl Authorizer {
    /// Build an [`Authorizer::AllowList`] from `(identity_name, &[query_name])`
    /// pairs. Each identity is permitted exactly the listed registered-query
    /// names (table names from `register_named`, or the Cypher key otherwise).
    pub fn allow_list<I, S1, S2>(entries: I) -> Self
    where
        I: IntoIterator<Item = (S1, Vec<S2>)>,
        S1: Into<String>,
        S2: Into<String>,
    {
        let map = entries
            .into_iter()
            .map(|(id, queries)| (id.into(), queries.into_iter().map(Into::into).collect()))
            .collect();
        Authorizer::AllowList(map)
    }

    /// Wrap a closure as an [`Authorizer::Custom`].
    pub fn custom<F>(f: F) -> Self
    where
        F: Fn(&Identity, &str) -> bool + Send + Sync + 'static,
    {
        Authorizer::Custom(Arc::new(f))
    }

    /// Is `identity` allowed to run the query resolved to `query_name`?
    pub fn is_allowed(&self, identity: &Identity, query_name: &str) -> bool {
        match self {
            Authorizer::AllowAll => true,
            Authorizer::AllowList(map) => {
                let allowed = |key: &str| map.get(key).is_some_and(|qs| qs.contains(query_name));
                // The identity's own allowances, then the wildcard fallback.
                allowed(identity.name()) || allowed("*")
            }
            Authorizer::Custom(f) => f(identity, query_name),
        }
    }

    /// Authorize or reject, mapping a denial to `permission_denied`.
    pub fn authorize(&self, identity: &Identity, query_name: &str) -> Result<(), Status> {
        if self.is_allowed(identity, query_name) {
            Ok(())
        } else {
            Err(Status::permission_denied(format!(
                "identity '{}' is not permitted to run query '{query_name}'",
                identity.name()
            )))
        }
    }
}

/// Which credentials the server accepts. An empty config (the default) means
/// the server is open — no interceptor is installed.
#[derive(Debug, Clone, Default)]
pub struct AuthConfig {
    /// Accepted `Bearer` tokens, keyed by token value → identity name. Any
    /// token in the set authenticates (rotation: old + new both valid). The
    /// first inserted token is the one issued by the handshake.
    bearer_tokens: Vec<(String, String)>,
    /// Accepted Basic username/password, if any.
    pub basic: Option<BasicCredential>,
    /// Per-query authorization policy, consulted after a query resolves. The
    /// default [`Authorizer::AllowAll`] is a no-op (no behavior change).
    pub authorizer: Authorizer,
}

impl AuthConfig {
    /// A single accepted bearer token (back-compat shape). Equivalent to
    /// [`with_bearer_tokens`] with one element; the token names itself.
    ///
    /// [`with_bearer_tokens`]: AuthConfig::with_bearer_tokens
    pub fn with_bearer_token(mut self, token: impl Into<String>) -> Self {
        self.add_bearer_token(token);
        self
    }

    /// A **set** of accepted bearer tokens for zero-downtime rotation. Any
    /// listed token authenticates; the first is issued by the handshake. Each
    /// token names itself as its [`Identity`] (use
    /// [`with_named_bearer_tokens`] to map tokens to authz identities).
    ///
    /// [`with_named_bearer_tokens`]: AuthConfig::with_named_bearer_tokens
    pub fn with_bearer_tokens<I, S>(mut self, tokens: I) -> Self
    where
        I: IntoIterator<Item = S>,
        S: Into<String>,
    {
        for t in tokens {
            self.add_bearer_token(t);
        }
        self
    }

    /// A set of accepted bearer tokens mapped to authz **identity names**, so
    /// several tokens can share one identity (e.g. an old and a new token for
    /// the same `analyst` role rotate transparently). Pairs are `(token,
    /// identity_name)`; the first pair's token is issued by the handshake.
    pub fn with_named_bearer_tokens<I, S1, S2>(mut self, tokens: I) -> Self
    where
        I: IntoIterator<Item = (S1, S2)>,
        S1: Into<String>,
        S2: Into<String>,
    {
        for (token, name) in tokens {
            self.add_named_bearer_token(token, name);
        }
        self
    }

    /// Add one accepted bearer token (self-named). Keeps any already added —
    /// the building block for rotation. Returns `&mut Self` for chaining.
    pub fn add_bearer_token(&mut self, token: impl Into<String>) -> &mut Self {
        let token = token.into();
        let name = token.clone();
        self.add_named_bearer_token(token, name)
    }

    /// Add one accepted bearer token mapped to an authz identity name.
    pub fn add_named_bearer_token(
        &mut self,
        token: impl Into<String>,
        name: impl Into<String>,
    ) -> &mut Self {
        self.bearer_tokens.push((token.into(), name.into()));
        self
    }

    /// Accept Basic `username` / `password`.
    pub fn with_basic(mut self, username: impl Into<String>, password: impl Into<String>) -> Self {
        self.basic = Some(BasicCredential {
            username: username.into(),
            password: password.into(),
        });
        self
    }

    /// Attach a per-query [`Authorizer`]. Default is [`Authorizer::AllowAll`].
    pub fn with_authorizer(mut self, authorizer: Authorizer) -> Self {
        self.authorizer = authorizer;
        self
    }

    /// Authorize an already-resolved query for an identity (delegates to the
    /// configured [`Authorizer`]). Denial → `permission_denied`.
    pub fn authorize(&self, identity: &Identity, query_name: &str) -> Result<(), Status> {
        self.authorizer.authorize(identity, query_name)
    }

    /// All accepted bearer token values, in insertion order.
    pub fn bearer_tokens(&self) -> impl Iterator<Item = &str> {
        self.bearer_tokens.iter().map(|(t, _)| t.as_str())
    }

    /// Any credential method configured? When false, the server requires no
    /// authentication and the interceptor is skipped entirely.
    pub fn is_enabled(&self) -> bool {
        !self.bearer_tokens.is_empty() || self.basic.is_some()
    }

    /// The token issued to a client after a successful handshake — the first
    /// configured bearer token, which the client then replays on each call.
    pub fn issued_token(&self) -> Option<&str> {
        self.bearer_tokens.first().map(|(t, _)| t.as_str())
    }

    /// Validate an `authorization` header value against every configured
    /// method. Open servers accept anything (including a missing header).
    ///
    /// A thin wrapper over [`authenticate`](Self::authenticate) that discards
    /// the resolved [`Identity`]; kept for call sites that only gate access.
    pub fn check_header(&self, header: Option<&str>) -> Result<(), Status> {
        self.authenticate(header).map(|_| ())
    }

    /// Validate an `authorization` header and return the resolved [`Identity`].
    ///
    /// Open servers (no method configured) authenticate anything as
    /// [`Identity::Anonymous`]. Otherwise a valid `Bearer` token resolves to its
    /// mapped [`Identity::Token`] and valid Basic creds to [`Identity::Basic`];
    /// anything else is `unauthenticated`.
    pub fn authenticate(&self, header: Option<&str>) -> Result<Identity, Status> {
        if !self.is_enabled() {
            return Ok(Identity::Anonymous);
        }
        let value = header.ok_or_else(|| {
            Status::unauthenticated(
                "missing 'authorization' header (Bearer token or Basic credentials)",
            )
        })?;

        if let Some(token) = value.strip_prefix("Bearer ") {
            // Constant-time compare against every accepted token (rotation).
            for (expected, name) in &self.bearer_tokens {
                if constant_time_eq(token.as_bytes(), expected.as_bytes()) {
                    return Ok(Identity::Token(name.clone()));
                }
            }
        } else if let Some(user) = value
            .strip_prefix("Basic ")
            .and_then(|b64| self.check_basic(b64))
        {
            return Ok(Identity::Basic(user));
        }
        Err(Status::unauthenticated("invalid credentials"))
    }

    /// Validate a base64-encoded `user:pass` Basic payload, returning the
    /// username on success.
    pub fn check_basic(&self, b64: &str) -> Option<String> {
        let expected = self.basic.as_ref()?;
        let decoded = BASE64.decode(b64).ok()?;
        let text = std::str::from_utf8(&decoded).ok()?;
        let (user, pass) = text.split_once(':')?;
        // Compare both fields in constant time; `&` (not `&&`) so timing does
        // not reveal which field differed.
        let ok = constant_time_eq(user.as_bytes(), expected.username.as_bytes())
            & constant_time_eq(pass.as_bytes(), expected.password.as_bytes());
        ok.then(|| user.to_string())
    }
}

/// Length-aware constant-time byte comparison.
pub(crate) fn constant_time_eq(a: &[u8], b: &[u8]) -> bool {
    if a.len() != b.len() {
        return false;
    }
    a.iter().zip(b).fold(0u8, |acc, (x, y)| acc | (x ^ y)) == 0
}

#[cfg(test)]
mod tests {
    use super::*;

    fn basic_header(user: &str, pass: &str) -> String {
        format!("Basic {}", BASE64.encode(format!("{user}:{pass}")))
    }

    #[test]
    fn open_server_accepts_anything() {
        let auth = AuthConfig::default();
        assert!(!auth.is_enabled());
        assert!(auth.check_header(None).is_ok());
        assert!(auth.check_header(Some("garbage")).is_ok());
        assert_eq!(auth.authenticate(None).unwrap(), Identity::Anonymous);
    }

    #[test]
    fn bearer_token_accepted_and_rejected() {
        let auth = AuthConfig::default().with_bearer_token("s3cret");
        assert!(auth.check_header(Some("Bearer s3cret")).is_ok());
        assert!(auth.check_header(Some("Bearer nope")).is_err());
        assert!(auth.check_header(None).is_err());
        assert_eq!(
            auth.authenticate(Some("Bearer s3cret")).unwrap(),
            Identity::Token("s3cret".into())
        );
    }

    #[test]
    fn multiple_tokens_all_accepted_for_rotation() {
        // Old + new token both valid during the overlap window.
        let auth = AuthConfig::default().with_bearer_tokens(["old-tok", "new-tok"]);
        assert!(auth.check_header(Some("Bearer old-tok")).is_ok());
        assert!(auth.check_header(Some("Bearer new-tok")).is_ok());
        // A stale/never-valid token is still rejected.
        assert!(auth.check_header(Some("Bearer retired-tok")).is_err());
        // The handshake issues the first (primary) token.
        assert_eq!(auth.issued_token(), Some("old-tok"));
    }

    #[test]
    fn add_bearer_token_grows_the_set() {
        let mut auth = AuthConfig::default().with_bearer_token("t1");
        auth.add_bearer_token("t2");
        assert!(auth.check_header(Some("Bearer t1")).is_ok());
        assert!(auth.check_header(Some("Bearer t2")).is_ok());
        assert_eq!(auth.bearer_tokens().collect::<Vec<_>>(), vec!["t1", "t2"]);
    }

    #[test]
    fn named_tokens_share_an_identity() {
        // Two tokens (rotation) that both authenticate as "analyst".
        let auth = AuthConfig::default()
            .with_named_bearer_tokens([("tok-a", "analyst"), ("tok-b", "analyst")]);
        assert_eq!(
            auth.authenticate(Some("Bearer tok-a")).unwrap(),
            Identity::Token("analyst".into())
        );
        assert_eq!(
            auth.authenticate(Some("Bearer tok-b")).unwrap(),
            Identity::Token("analyst".into())
        );
    }

    #[test]
    fn basic_credentials_accepted_and_rejected() {
        let auth = AuthConfig::default().with_basic("admin", "pw");
        assert!(
            auth.check_header(Some(&basic_header("admin", "pw")))
                .is_ok()
        );
        assert!(
            auth.check_header(Some(&basic_header("admin", "wrong")))
                .is_err()
        );
        assert!(auth.check_header(Some(&basic_header("eve", "pw"))).is_err());
        assert!(auth.check_header(Some("Basic !!!notbase64")).is_err());
        assert_eq!(
            auth.authenticate(Some(&basic_header("admin", "pw")))
                .unwrap(),
            Identity::Basic("admin".into())
        );
    }

    #[test]
    fn both_methods_accepted_when_both_configured() {
        let auth = AuthConfig::default()
            .with_bearer_token("tok")
            .with_basic("u", "p");
        assert!(auth.check_header(Some("Bearer tok")).is_ok());
        assert!(auth.check_header(Some(&basic_header("u", "p"))).is_ok());
        assert!(auth.check_header(Some("Bearer bad")).is_err());
    }

    #[test]
    fn issued_token_is_the_first_bearer_token() {
        let auth = AuthConfig::default().with_bearer_tokens(["tok", "tok2"]);
        assert_eq!(auth.issued_token(), Some("tok"));
        assert_eq!(AuthConfig::default().issued_token(), None);
    }

    #[test]
    fn authorizer_allow_all_permits_everything() {
        let authz = Authorizer::AllowAll;
        assert!(authz.is_allowed(&Identity::Token("anyone".into()), "q1"));
        assert!(authz.is_allowed(&Identity::Anonymous, "q2"));
        assert!(
            authz
                .authorize(&Identity::Basic("u".into()), "anything")
                .is_ok()
        );
    }

    #[test]
    fn authorizer_allow_list_allows_and_denies_per_identity() {
        // analyst may run q1 but not q2.
        let authz = Authorizer::allow_list([("analyst", vec!["q1"])]);
        let analyst = Identity::Token("analyst".into());
        assert!(authz.is_allowed(&analyst, "q1"));
        assert!(!authz.is_allowed(&analyst, "q2"));
        assert!(authz.authorize(&analyst, "q1").is_ok());
        let denied = authz.authorize(&analyst, "q2").unwrap_err();
        assert_eq!(denied.code(), tonic::Code::PermissionDenied);
        // An identity absent from the list is denied everything.
        assert!(!authz.is_allowed(&Identity::Token("stranger".into()), "q1"));
    }

    #[test]
    fn authorizer_allow_list_wildcard_fallback() {
        // "*" grants q_public to every identity not otherwise listed.
        let authz = Authorizer::allow_list([("analyst", vec!["q1"]), ("*", vec!["q_public"])]);
        let analyst = Identity::Token("analyst".into());
        assert!(authz.is_allowed(&analyst, "q1"));
        assert!(authz.is_allowed(&analyst, "q_public")); // wildcard fallback
        let other = Identity::Token("other".into());
        assert!(authz.is_allowed(&other, "q_public"));
        assert!(!authz.is_allowed(&other, "q1"));
    }

    #[test]
    fn authorizer_custom_closure() {
        let authz = Authorizer::custom(|id, q| id.name() == "root" || q == "q_open");
        assert!(authz.is_allowed(&Identity::Token("root".into()), "anything"));
        assert!(authz.is_allowed(&Identity::Token("nobody".into()), "q_open"));
        assert!(!authz.is_allowed(&Identity::Token("nobody".into()), "q_secret"));
    }
}