crates/graphar-flight/src/server.rs
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,
/// How many `UNWIND`-window queries a **DoPut** keeps in flight against
/// FalkorDB. `None` (the default) uses
/// [`default_put_concurrency`](crate::service::default_put_concurrency) —
/// the host's parallelism. Set `Some(1)` to pin the ingest back to the old
/// serial one-round-trip-at-a-time dispatch (a shared/throttled FalkorDB, or
/// an A/B measurement).
pub put_concurrency: Option<usize>,
}
impl ServerOptions {
/// **The ONE mapping from these options onto the acceptance rules the server
/// enforces** — Basic credential, the issued bearer, the rotation set, the
/// per-query authorizer. [`run_server_with`] installs the result as its
/// interceptor.
///
/// Public because it is the seam a HOST service needs. `CypherFlightService` is
/// Cypher-bound by construction (every registered table name is a Cypher string),
/// so a host that must also serve non-graph tables — korp's Flight SQL endpoint
/// over its Iceberg warehouse — wraps this service and builds its own tonic
/// server. Without this method that host would have to re-derive the mapping from
/// the public option fields, i.e. keep a second copy of the auth rules that goes
/// stale the first time a field is added here. One writer, read from anywhere.
pub 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 mut service = CypherFlightService::connect(redis_url, graph, registry)
.await?
.with_auth(auth.clone());
if let Some(n) = options.put_concurrency {
service = service.with_put_concurrency(n);
}
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);
}
}