graphar-flight 0.1.2

use std::{net::SocketAddr, sync::Arc};

use arrow_flight::flight_service_server::FlightServiceServer;
use tonic::{
    Request, Status,
    service::Interceptor,
    transport::{Certificate, Identity, Server, ServerTlsConfig},
};

use crate::{
    auth::{AuthConfig, Authorizer, BasicCredential},
    error::Result,
    registry::SchemaRegistry,
    service::CypherFlightService,
};

/// TLS and authentication options for [`run_server_with`].
///
/// All three auth methods compose: e.g. mTLS for transport identity *and* a
/// bearer/Basic header for application identity.
#[derive(Debug, Clone, Default)]
pub struct ServerOptions {
    /// PEM-encoded server certificate chain + private key (and, for mTLS, the
    /// client CA). When set, the server only accepts TLS connections.
    pub tls: Option<TlsOptions>,
    /// When set, requests may authenticate with `authorization: Bearer <token>`.
    /// Also the token handed back by the Flight handshake after a Basic login.
    /// For zero-downtime rotation, list more than one in [`bearer_tokens`].
    ///
    /// [`bearer_tokens`]: ServerOptions::bearer_tokens
    pub bearer_token: Option<String>,
    /// A **set** of accepted bearer tokens for zero-downtime rotation — any one
    /// authenticates, the first is issued by the handshake. Merged with
    /// (after) [`bearer_token`]. Empty by default; leave it so to keep the
    /// single-token behavior unchanged.
    ///
    /// [`bearer_token`]: ServerOptions::bearer_token
    pub bearer_tokens: Vec<String>,
    /// When set, requests may authenticate with `authorization: Basic <b64>`
    /// (`username:password`) — what Power BI's ADBC connector and the Flight
    /// SQL ODBC driver send. `(username, password)`.
    pub basic_auth: Option<(String, String)>,
    /// Per-query authorization policy. Default [`Authorizer::AllowAll`] permits
    /// every identity to run every query (no behavior change).
    pub authorizer: Authorizer,
}

impl ServerOptions {
    fn auth_config(&self) -> AuthConfig {
        let mut cfg = AuthConfig::default().with_authorizer(self.authorizer.clone());
        cfg.basic = self.basic_auth.as_ref().map(|(u, p)| BasicCredential {
            username: u.clone(),
            password: p.clone(),
        });
        // `bearer_token` first (so it stays the issued token), then the set.
        if let Some(t) = &self.bearer_token {
            cfg.add_bearer_token(t.clone());
        }
        for t in &self.bearer_tokens {
            cfg.add_bearer_token(t.clone());
        }
        cfg
    }
}

#[derive(Debug, Clone)]
pub struct TlsOptions {
    pub cert_pem: Vec<u8>,
    pub key_pem: Vec<u8>,
    /// PEM-encoded CA certificate used to verify client certificates. When
    /// set, the server requires **mutual TLS**: clients must present a cert
    /// signed by this CA.
    pub client_ca_pem: Option<Vec<u8>>,
}

impl TlsOptions {
    /// Server-side TLS only (no client-certificate verification).
    pub fn from_pem_files(
        cert: impl AsRef<std::path::Path>,
        key: impl AsRef<std::path::Path>,
    ) -> std::io::Result<Self> {
        Ok(Self {
            cert_pem: std::fs::read(cert)?,
            key_pem: std::fs::read(key)?,
            client_ca_pem: None,
        })
    }

    /// Mutual TLS: server cert/key plus the CA that client certs are verified
    /// against. Clients without a CA-signed cert are rejected at the transport.
    pub fn mutual_from_pem_files(
        cert: impl AsRef<std::path::Path>,
        key: impl AsRef<std::path::Path>,
        client_ca: impl AsRef<std::path::Path>,
    ) -> std::io::Result<Self> {
        Ok(Self {
            cert_pem: std::fs::read(cert)?,
            key_pem: std::fs::read(key)?,
            client_ca_pem: Some(std::fs::read(client_ca)?),
        })
    }
}

/// Per-call header authentication: accepts any method configured in
/// [`AuthConfig`] (Bearer token or Basic credentials). The Flight handshake
/// also passes through here carrying its Basic header, so a client can obtain
/// a token and then replay it.
#[derive(Clone)]
struct AuthInterceptor {
    auth: Arc<AuthConfig>,
}

impl Interceptor for AuthInterceptor {
    fn call(&mut self, mut request: Request<()>) -> std::result::Result<Request<()>, Status> {
        let header = request
            .metadata()
            .get("authorization")
            .and_then(|v| v.to_str().ok());
        // Authenticate and stash the resolved identity so the per-query
        // authorization hook in the service can key its decision on *who* asked.
        let identity = self.auth.authenticate(header)?;
        request.extensions_mut().insert(identity);
        Ok(request)
    }
}

/// Start the Cypher Flight SQL server and block until it shuts down.
///
/// ```no_run
/// use std::{net::SocketAddr, sync::Arc};
/// use graphar_flight::{SchemaRegistry, run_server};
/// use arrow_schema::{DataType, Field, Schema};
///
/// #[tokio::main]
/// async fn main() {
///     let registry = Arc::new(SchemaRegistry::new());
///     registry.register(
///         "MATCH (n:Person) RETURN n.id, n.name",
///         Arc::new(Schema::new(vec![
///             Field::new("n.id",   DataType::Int64, false),
///             Field::new("n.name", DataType::Utf8,  true),
///         ])),
///     );
///     run_server("redis://127.0.0.1:6379", "mygraph", registry, "0.0.0.0:50051".parse().unwrap()).await.unwrap();
/// }
/// ```
pub async fn run_server(
    redis_url: &str,
    graph: &str,
    registry: Arc<SchemaRegistry>,
    addr: SocketAddr,
) -> Result<()> {
    run_server_with(redis_url, graph, registry, addr, ServerOptions::default()).await
}

/// [`run_server`] with TLS and/or header authentication.
///
/// This example pairs mutual TLS (transport identity) with Basic + bearer
/// header auth (application identity) — the belt-and-suspenders setup a
/// Power BI / ADBC client on Windows connects to.
///
/// ```no_run
/// use std::sync::Arc;
/// use graphar_flight::{SchemaRegistry, run_server_with, ServerOptions, TlsOptions};
///
/// #[tokio::main]
/// async fn main() {
///     let opts = ServerOptions {
///         tls: Some(
///             TlsOptions::mutual_from_pem_files("server.crt", "server.key", "client-ca.crt")
///                 .unwrap(),
///         ),
///         bearer_token: Some("s3cret".into()),
///         basic_auth: Some(("analyst".into(), "p@ss".into())),
///         ..Default::default()
///     };
///     run_server_with(
///         "redis://127.0.0.1:6379",
///         "mygraph",
///         Arc::new(SchemaRegistry::new()),
///         "0.0.0.0:50051".parse().unwrap(),
///         opts,
///     ).await.unwrap();
/// }
/// ```
pub async fn run_server_with(
    redis_url: &str,
    graph: &str,
    registry: Arc<SchemaRegistry>,
    addr: SocketAddr,
    options: ServerOptions,
) -> Result<()> {
    let auth = Arc::new(options.auth_config());
    let service = CypherFlightService::connect(redis_url, graph, registry)
        .await?
        .with_auth(auth.clone());

    let mut builder = Server::builder();
    if let Some(tls) = &options.tls {
        let identity = Identity::from_pem(&tls.cert_pem, &tls.key_pem);
        let mut tls_config = ServerTlsConfig::new().identity(identity);
        if let Some(ca) = &tls.client_ca_pem {
            // mTLS: require and verify a client certificate.
            tls_config = tls_config.client_ca_root(Certificate::from_pem(ca));
        }
        builder = builder.tls_config(tls_config)?;
    }

    // Install the header-auth interceptor only when a header method is
    // configured; otherwise the server is open (transport-only / mTLS-only).
    let router = if auth.is_enabled() {
        let interceptor = AuthInterceptor { auth };
        builder.add_service(FlightServiceServer::with_interceptor(service, interceptor))
    } else {
        builder.add_service(FlightServiceServer::new(service))
    };

    router.serve(addr).await?;
    Ok(())
}

#[cfg(test)]
mod auth_tests {
    use super::*;
    use crate::auth::Identity;
    use base64::Engine;
    use base64::engine::general_purpose::STANDARD as BASE64;

    fn request_with_auth(value: Option<&str>) -> Request<()> {
        let mut req = Request::new(());
        if let Some(v) = value {
            req.metadata_mut()
                .insert("authorization", v.parse().unwrap());
        }
        req
    }

    fn interceptor(bearer: Option<&str>, basic: Option<(&str, &str)>) -> AuthInterceptor {
        let opts = ServerOptions {
            bearer_token: bearer.map(Into::into),
            basic_auth: basic.map(|(u, p)| (u.into(), p.into())),
            ..ServerOptions::default()
        };
        AuthInterceptor {
            auth: Arc::new(opts.auth_config()),
        }
    }

    #[test]
    fn bearer_accepts_correct_and_rejects_wrong() {
        let mut auth = interceptor(Some("s3cret"), None);
        assert!(auth.call(request_with_auth(Some("Bearer s3cret"))).is_ok());
        let err = auth
            .call(request_with_auth(Some("Bearer nope")))
            .unwrap_err();
        assert_eq!(err.code(), tonic::Code::Unauthenticated);
        assert!(auth.call(request_with_auth(None)).is_err());
    }

    #[test]
    fn basic_credentials_accepted() {
        let mut auth = interceptor(None, Some(("analyst", "p@ss")));
        let header = format!("Basic {}", BASE64.encode("analyst:p@ss"));
        assert!(auth.call(request_with_auth(Some(&header))).is_ok());
        let wrong = format!("Basic {}", BASE64.encode("analyst:nope"));
        assert!(auth.call(request_with_auth(Some(&wrong))).is_err());
    }

    #[test]
    fn options_map_to_auth_config() {
        let opts = ServerOptions {
            bearer_token: Some("t".into()),
            basic_auth: Some(("u".into(), "p".into())),
            ..ServerOptions::default()
        };
        let cfg = opts.auth_config();
        assert_eq!(cfg.issued_token(), Some("t"));
        assert_eq!(cfg.basic.as_ref().map(|b| b.username.as_str()), Some("u"));
        assert!(cfg.is_enabled());
    }

    #[test]
    fn options_token_set_maps_to_rotating_config() {
        // `bearer_token` stays the issued (primary) token; `bearer_tokens`
        // are additional valid tokens for the rotation overlap window.
        let opts = ServerOptions {
            bearer_token: Some("primary".into()),
            bearer_tokens: vec!["secondary".into()],
            ..ServerOptions::default()
        };
        let cfg = opts.auth_config();
        assert_eq!(cfg.issued_token(), Some("primary"));
        assert!(cfg.check_header(Some("Bearer primary")).is_ok());
        assert!(cfg.check_header(Some("Bearer secondary")).is_ok());
        assert!(cfg.check_header(Some("Bearer stale")).is_err());
    }

    #[test]
    fn rotation_accepts_old_and_new_rejects_stale_through_interceptor() {
        let opts = ServerOptions {
            bearer_tokens: vec!["old".into(), "new".into()],
            ..ServerOptions::default()
        };
        let mut auth = AuthInterceptor {
            auth: Arc::new(opts.auth_config()),
        };
        assert!(auth.call(request_with_auth(Some("Bearer old"))).is_ok());
        assert!(auth.call(request_with_auth(Some("Bearer new"))).is_ok());
        assert!(
            auth.call(request_with_auth(Some("Bearer retired")))
                .is_err()
        );
    }

    #[test]
    fn interceptor_stashes_identity_extension() {
        let opts = ServerOptions {
            bearer_tokens: vec!["tok".into()],
            ..ServerOptions::default()
        };
        let mut auth = AuthInterceptor {
            auth: Arc::new(opts.auth_config()),
        };
        let req = auth.call(request_with_auth(Some("Bearer tok"))).unwrap();
        assert_eq!(
            req.extensions().get::<Identity>(),
            Some(&Identity::Token("tok".into())),
        );
    }

    #[test]
    fn options_carry_authorizer_into_config() {
        // A token allowed `q1` but denied `q2`.
        let opts = ServerOptions {
            bearer_tokens: vec!["tok".into()],
            authorizer: Authorizer::allow_list([("tok", vec!["q1"])]),
            ..ServerOptions::default()
        };
        let cfg = opts.auth_config();
        let id = Identity::Token("tok".into());
        assert!(cfg.authorize(&id, "q1").is_ok());
        let denied = cfg.authorize(&id, "q2").unwrap_err();
        assert_eq!(denied.code(), tonic::Code::PermissionDenied);
    }
}