graphar_flight/
server.rs

1use std::{net::SocketAddr, sync::Arc};
2
3use arrow_flight::flight_service_server::FlightServiceServer;
4use tonic::{
5    Request, Status,
6    service::Interceptor,
7    transport::{Certificate, Identity, Server, ServerTlsConfig},
8};
9
10use crate::{
11    auth::{AuthConfig, Authorizer, BasicCredential},
12    error::Result,
13    registry::SchemaRegistry,
14    service::CypherFlightService,
15};
16
17/// TLS and authentication options for [`run_server_with`].
18///
19/// All three auth methods compose: e.g. mTLS for transport identity *and* a
20/// bearer/Basic header for application identity.
21#[derive(Debug, Clone, Default)]
22pub struct ServerOptions {
23    /// PEM-encoded server certificate chain + private key (and, for mTLS, the
24    /// client CA). When set, the server only accepts TLS connections.
25    pub tls: Option<TlsOptions>,
26    /// When set, requests may authenticate with `authorization: Bearer <token>`.
27    /// Also the token handed back by the Flight handshake after a Basic login.
28    /// For zero-downtime rotation, list more than one in [`bearer_tokens`].
29    ///
30    /// [`bearer_tokens`]: ServerOptions::bearer_tokens
31    pub bearer_token: Option<String>,
32    /// A **set** of accepted bearer tokens for zero-downtime rotation — any one
33    /// authenticates, the first is issued by the handshake. Merged with
34    /// (after) [`bearer_token`]. Empty by default; leave it so to keep the
35    /// single-token behavior unchanged.
36    ///
37    /// [`bearer_token`]: ServerOptions::bearer_token
38    pub bearer_tokens: Vec<String>,
39    /// When set, requests may authenticate with `authorization: Basic <b64>`
40    /// (`username:password`) — what Power BI's ADBC connector and the Flight
41    /// SQL ODBC driver send. `(username, password)`.
42    pub basic_auth: Option<(String, String)>,
43    /// Per-query authorization policy. Default [`Authorizer::AllowAll`] permits
44    /// every identity to run every query (no behavior change).
45    pub authorizer: Authorizer,
46}
47
48impl ServerOptions {
49    fn auth_config(&self) -> AuthConfig {
50        let mut cfg = AuthConfig::default().with_authorizer(self.authorizer.clone());
51        cfg.basic = self.basic_auth.as_ref().map(|(u, p)| BasicCredential {
52            username: u.clone(),
53            password: p.clone(),
54        });
55        // `bearer_token` first (so it stays the issued token), then the set.
56        if let Some(t) = &self.bearer_token {
57            cfg.add_bearer_token(t.clone());
58        }
59        for t in &self.bearer_tokens {
60            cfg.add_bearer_token(t.clone());
61        }
62        cfg
63    }
64}
65
66#[derive(Debug, Clone)]
67pub struct TlsOptions {
68    pub cert_pem: Vec<u8>,
69    pub key_pem: Vec<u8>,
70    /// PEM-encoded CA certificate used to verify client certificates. When
71    /// set, the server requires **mutual TLS**: clients must present a cert
72    /// signed by this CA.
73    pub client_ca_pem: Option<Vec<u8>>,
74}
75
76impl TlsOptions {
77    /// Server-side TLS only (no client-certificate verification).
78    pub fn from_pem_files(
79        cert: impl AsRef<std::path::Path>,
80        key: impl AsRef<std::path::Path>,
81    ) -> std::io::Result<Self> {
82        Ok(Self {
83            cert_pem: std::fs::read(cert)?,
84            key_pem: std::fs::read(key)?,
85            client_ca_pem: None,
86        })
87    }
88
89    /// Mutual TLS: server cert/key plus the CA that client certs are verified
90    /// against. Clients without a CA-signed cert are rejected at the transport.
91    pub fn mutual_from_pem_files(
92        cert: impl AsRef<std::path::Path>,
93        key: impl AsRef<std::path::Path>,
94        client_ca: impl AsRef<std::path::Path>,
95    ) -> std::io::Result<Self> {
96        Ok(Self {
97            cert_pem: std::fs::read(cert)?,
98            key_pem: std::fs::read(key)?,
99            client_ca_pem: Some(std::fs::read(client_ca)?),
100        })
101    }
102}
103
104/// Per-call header authentication: accepts any method configured in
105/// [`AuthConfig`] (Bearer token or Basic credentials). The Flight handshake
106/// also passes through here carrying its Basic header, so a client can obtain
107/// a token and then replay it.
108#[derive(Clone)]
109struct AuthInterceptor {
110    auth: Arc<AuthConfig>,
111}
112
113impl Interceptor for AuthInterceptor {
114    fn call(&mut self, mut request: Request<()>) -> std::result::Result<Request<()>, Status> {
115        let header = request
116            .metadata()
117            .get("authorization")
118            .and_then(|v| v.to_str().ok());
119        // Authenticate and stash the resolved identity so the per-query
120        // authorization hook in the service can key its decision on *who* asked.
121        let identity = self.auth.authenticate(header)?;
122        request.extensions_mut().insert(identity);
123        Ok(request)
124    }
125}
126
127/// Start the Cypher Flight SQL server and block until it shuts down.
128///
129/// ```no_run
130/// use std::{net::SocketAddr, sync::Arc};
131/// use graphar_flight::{SchemaRegistry, run_server};
132/// use arrow_schema::{DataType, Field, Schema};
133///
134/// #[tokio::main]
135/// async fn main() {
136///     let registry = Arc::new(SchemaRegistry::new());
137///     registry.register(
138///         "MATCH (n:Person) RETURN n.id, n.name",
139///         Arc::new(Schema::new(vec![
140///             Field::new("n.id",   DataType::Int64, false),
141///             Field::new("n.name", DataType::Utf8,  true),
142///         ])),
143///     );
144///     run_server("redis://127.0.0.1:6379", "mygraph", registry, "0.0.0.0:50051".parse().unwrap()).await.unwrap();
145/// }
146/// ```
147pub async fn run_server(
148    redis_url: &str,
149    graph: &str,
150    registry: Arc<SchemaRegistry>,
151    addr: SocketAddr,
152) -> Result<()> {
153    run_server_with(redis_url, graph, registry, addr, ServerOptions::default()).await
154}
155
156/// [`run_server`] with TLS and/or header authentication.
157///
158/// This example pairs mutual TLS (transport identity) with Basic + bearer
159/// header auth (application identity) — the belt-and-suspenders setup a
160/// Power BI / ADBC client on Windows connects to.
161///
162/// ```no_run
163/// use std::sync::Arc;
164/// use graphar_flight::{SchemaRegistry, run_server_with, ServerOptions, TlsOptions};
165///
166/// #[tokio::main]
167/// async fn main() {
168///     let opts = ServerOptions {
169///         tls: Some(
170///             TlsOptions::mutual_from_pem_files("server.crt", "server.key", "client-ca.crt")
171///                 .unwrap(),
172///         ),
173///         bearer_token: Some("s3cret".into()),
174///         basic_auth: Some(("analyst".into(), "p@ss".into())),
175///         ..Default::default()
176///     };
177///     run_server_with(
178///         "redis://127.0.0.1:6379",
179///         "mygraph",
180///         Arc::new(SchemaRegistry::new()),
181///         "0.0.0.0:50051".parse().unwrap(),
182///         opts,
183///     ).await.unwrap();
184/// }
185/// ```
186pub async fn run_server_with(
187    redis_url: &str,
188    graph: &str,
189    registry: Arc<SchemaRegistry>,
190    addr: SocketAddr,
191    options: ServerOptions,
192) -> Result<()> {
193    let auth = Arc::new(options.auth_config());
194    let service = CypherFlightService::connect(redis_url, graph, registry)
195        .await?
196        .with_auth(auth.clone());
197
198    let mut builder = Server::builder();
199    if let Some(tls) = &options.tls {
200        let identity = Identity::from_pem(&tls.cert_pem, &tls.key_pem);
201        let mut tls_config = ServerTlsConfig::new().identity(identity);
202        if let Some(ca) = &tls.client_ca_pem {
203            // mTLS: require and verify a client certificate.
204            tls_config = tls_config.client_ca_root(Certificate::from_pem(ca));
205        }
206        builder = builder.tls_config(tls_config)?;
207    }
208
209    // Install the header-auth interceptor only when a header method is
210    // configured; otherwise the server is open (transport-only / mTLS-only).
211    let router = if auth.is_enabled() {
212        let interceptor = AuthInterceptor { auth };
213        builder.add_service(FlightServiceServer::with_interceptor(service, interceptor))
214    } else {
215        builder.add_service(FlightServiceServer::new(service))
216    };
217
218    router.serve(addr).await?;
219    Ok(())
220}
221
222#[cfg(test)]
223mod auth_tests {
224    use super::*;
225    use crate::auth::Identity;
226    use base64::Engine;
227    use base64::engine::general_purpose::STANDARD as BASE64;
228
229    fn request_with_auth(value: Option<&str>) -> Request<()> {
230        let mut req = Request::new(());
231        if let Some(v) = value {
232            req.metadata_mut()
233                .insert("authorization", v.parse().unwrap());
234        }
235        req
236    }
237
238    fn interceptor(bearer: Option<&str>, basic: Option<(&str, &str)>) -> AuthInterceptor {
239        let opts = ServerOptions {
240            bearer_token: bearer.map(Into::into),
241            basic_auth: basic.map(|(u, p)| (u.into(), p.into())),
242            ..ServerOptions::default()
243        };
244        AuthInterceptor {
245            auth: Arc::new(opts.auth_config()),
246        }
247    }
248
249    #[test]
250    fn bearer_accepts_correct_and_rejects_wrong() {
251        let mut auth = interceptor(Some("s3cret"), None);
252        assert!(auth.call(request_with_auth(Some("Bearer s3cret"))).is_ok());
253        let err = auth
254            .call(request_with_auth(Some("Bearer nope")))
255            .unwrap_err();
256        assert_eq!(err.code(), tonic::Code::Unauthenticated);
257        assert!(auth.call(request_with_auth(None)).is_err());
258    }
259
260    #[test]
261    fn basic_credentials_accepted() {
262        let mut auth = interceptor(None, Some(("analyst", "p@ss")));
263        let header = format!("Basic {}", BASE64.encode("analyst:p@ss"));
264        assert!(auth.call(request_with_auth(Some(&header))).is_ok());
265        let wrong = format!("Basic {}", BASE64.encode("analyst:nope"));
266        assert!(auth.call(request_with_auth(Some(&wrong))).is_err());
267    }
268
269    #[test]
270    fn options_map_to_auth_config() {
271        let opts = ServerOptions {
272            bearer_token: Some("t".into()),
273            basic_auth: Some(("u".into(), "p".into())),
274            ..ServerOptions::default()
275        };
276        let cfg = opts.auth_config();
277        assert_eq!(cfg.issued_token(), Some("t"));
278        assert_eq!(cfg.basic.as_ref().map(|b| b.username.as_str()), Some("u"));
279        assert!(cfg.is_enabled());
280    }
281
282    #[test]
283    fn options_token_set_maps_to_rotating_config() {
284        // `bearer_token` stays the issued (primary) token; `bearer_tokens`
285        // are additional valid tokens for the rotation overlap window.
286        let opts = ServerOptions {
287            bearer_token: Some("primary".into()),
288            bearer_tokens: vec!["secondary".into()],
289            ..ServerOptions::default()
290        };
291        let cfg = opts.auth_config();
292        assert_eq!(cfg.issued_token(), Some("primary"));
293        assert!(cfg.check_header(Some("Bearer primary")).is_ok());
294        assert!(cfg.check_header(Some("Bearer secondary")).is_ok());
295        assert!(cfg.check_header(Some("Bearer stale")).is_err());
296    }
297
298    #[test]
299    fn rotation_accepts_old_and_new_rejects_stale_through_interceptor() {
300        let opts = ServerOptions {
301            bearer_tokens: vec!["old".into(), "new".into()],
302            ..ServerOptions::default()
303        };
304        let mut auth = AuthInterceptor {
305            auth: Arc::new(opts.auth_config()),
306        };
307        assert!(auth.call(request_with_auth(Some("Bearer old"))).is_ok());
308        assert!(auth.call(request_with_auth(Some("Bearer new"))).is_ok());
309        assert!(
310            auth.call(request_with_auth(Some("Bearer retired")))
311                .is_err()
312        );
313    }
314
315    #[test]
316    fn interceptor_stashes_identity_extension() {
317        let opts = ServerOptions {
318            bearer_tokens: vec!["tok".into()],
319            ..ServerOptions::default()
320        };
321        let mut auth = AuthInterceptor {
322            auth: Arc::new(opts.auth_config()),
323        };
324        let req = auth.call(request_with_auth(Some("Bearer tok"))).unwrap();
325        assert_eq!(
326            req.extensions().get::<Identity>(),
327            Some(&Identity::Token("tok".into())),
328        );
329    }
330
331    #[test]
332    fn options_carry_authorizer_into_config() {
333        // A token allowed `q1` but denied `q2`.
334        let opts = ServerOptions {
335            bearer_tokens: vec!["tok".into()],
336            authorizer: Authorizer::allow_list([("tok", vec!["q1"])]),
337            ..ServerOptions::default()
338        };
339        let cfg = opts.auth_config();
340        let id = Identity::Token("tok".into());
341        assert!(cfg.authorize(&id, "q1").is_ok());
342        let denied = cfg.authorize(&id, "q2").unwrap_err();
343        assert_eq!(denied.code(), tonic::Code::PermissionDenied);
344    }
345}
graphar_flight/server.rs

graphar_flight/
server.rs
