diff --git a/Cargo.lock b/Cargo.lock index 6629fa36..6f5a4c85 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3501,6 +3501,14 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "modelardb_auth" +version = "0.1.0" +dependencies = [ + "async-trait", + "tonic", +] + [[package]] name = "modelardb_bulkloader" version = "0.1.0" @@ -3528,6 +3536,7 @@ dependencies = [ "bytes", "clap", "dirs", + "modelardb_auth", "rustyline", "tokio", "tonic", @@ -3553,6 +3562,7 @@ dependencies = [ "datafusion", "deltalake", "futures", + "modelardb_auth", "modelardb_compression", "modelardb_storage", "modelardb_types", @@ -3575,7 +3585,10 @@ dependencies = [ "datafusion", "deltalake", "futures", + "http 1.4.2", + "http-body-util", "log", + "modelardb_auth", "modelardb_compression", "modelardb_storage", "modelardb_test", @@ -3586,11 +3599,13 @@ dependencies = [ "rand 0.10.1", "serde", "snmalloc-rs", + "sqlparser", "tempfile", "tokio", "tokio-stream", "toml", "tonic", + "tower", "tracing", "tracing-subscriber", ] diff --git a/Cargo.toml b/Cargo.toml index 52524dfe..346c167b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -32,6 +32,8 @@ delta_kernel = { package = "buoyant_kernel", version = "0.22.2" } deltalake = "0.32.4" dirs = "6.0.0" futures = "0.3.32" +http = "1.4.0" +http-body-util = "0.1.3" log = "0.4.29" object_store = "0.13.2" proptest = "1.11.0" @@ -50,6 +52,7 @@ tokio = "1.52.2" tokio-stream = "0.1.18" toml = "1.1.2" tonic = "0.14.4" +tower = "0.5.3" tracing = "0.1.44" tracing-subscriber = "0.3.23" url = "2.5.8" diff --git a/crates/modelardb_auth/Cargo.toml b/crates/modelardb_auth/Cargo.toml new file mode 100644 index 00000000..c36c0927 --- /dev/null +++ b/crates/modelardb_auth/Cargo.toml @@ -0,0 +1,27 @@ +# Copyright 2026 The ModelarDB Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +[package] +name = "modelardb_auth" +version = "0.1.0" +license = "Apache-2.0" +edition = "2024" +authors = ["Soeren Kejser Jensen "] + +[dependencies] +async-trait.workspace = true +tonic.workspace = true + +[features] +testing = [] diff --git a/crates/modelardb_auth/src/authenticator/mock.rs b/crates/modelardb_auth/src/authenticator/mock.rs new file mode 100644 index 00000000..464107c4 --- /dev/null +++ b/crates/modelardb_auth/src/authenticator/mock.rs @@ -0,0 +1,69 @@ +/* Copyright 2026 The ModelarDB Contributors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +//! A mock [`Authenticator`] for use in tests. Always grants access and records every permission +//! that was requested so tests can assert which permission the auth layer required for a given +//! request. + +use std::sync::Mutex; + +use async_trait::async_trait; +use tonic::Status; +use tonic::metadata::MetadataMap; + +use crate::Permission; +use crate::authenticator::Authenticator; + +/// An [`Authenticator`] for use in tests that always grants access and records every +/// [`Permission`] passed to [`authorize`](Authenticator::authorize). +pub struct MockAuthenticator { + /// Every [`Permission`] passed to [`authorize`](Authenticator::authorize). A [`Mutex`] is used + /// because [`authorize`](Authenticator::authorize) records through `&self`, which requires + /// interior mutability. + permissions: Mutex>, +} + +impl MockAuthenticator { + pub fn new() -> Self { + Self { + permissions: Mutex::new(vec![]), + } + } + + /// Return a copy of all permissions passed to [`authorize`](Authenticator::authorize) + /// in the order they were received. + pub fn permissions(&self) -> Vec { + self.permissions.lock().unwrap().clone() + } +} + +impl Default for MockAuthenticator { + // Trait implemented to silence clippy warning. + fn default() -> Self { + Self::new() + } +} + +#[async_trait] +impl Authenticator for MockAuthenticator { + async fn authorize( + &self, + _metadata: &MetadataMap, + required_permission: Permission, + ) -> Result<(), Status> { + self.permissions.lock().unwrap().push(required_permission); + Ok(()) + } +} diff --git a/crates/modelardb_auth/src/authenticator/mod.rs b/crates/modelardb_auth/src/authenticator/mod.rs new file mode 100644 index 00000000..0af878fa --- /dev/null +++ b/crates/modelardb_auth/src/authenticator/mod.rs @@ -0,0 +1,44 @@ +/* Copyright 2026 The ModelarDB Contributors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +//! Defines the [`Authenticator`] trait for validating credentials and authorizing access to +//! ModelarDB. + +#[cfg(feature = "testing")] +pub mod mock; + +use async_trait::async_trait; +use tonic::Status; +use tonic::metadata::MetadataMap; + +use crate::Permission; + +/// Validates the credentials in `metadata` and checks that the caller has the `required_permission`. +/// Implementations must return: +/// +/// - `Ok(())` if authorized. +/// - `Err(Status::unauthenticated(...))` if credentials are missing or invalid. +/// - `Err(Status::permission_denied(...))` if credentials are valid but the caller lacks the +/// required permission. +/// +/// Error messages must not reveal the specific reason for rejection to prevent information leakage. +#[async_trait] +pub trait Authenticator: Send + Sync { + async fn authorize( + &self, + metadata: &MetadataMap, + required_permission: Permission, + ) -> Result<(), Status>; +} diff --git a/crates/modelardb_auth/src/lib.rs b/crates/modelardb_auth/src/lib.rs new file mode 100644 index 00000000..763c383f --- /dev/null +++ b/crates/modelardb_auth/src/lib.rs @@ -0,0 +1,171 @@ +/* Copyright 2026 The ModelarDB Contributors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +//! Types to support authentication and authorization in ModelarDB. + +use tonic::metadata::AsciiMetadataValue; +use tonic::service::Interceptor; +use tonic::{Request, Status}; + +pub mod authenticator; + +/// The permission required to perform an operation in a ModelarDB server. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum Permission { + Read, + Write, + Admin, +} + +impl Permission { + /// Return [`true`] if `granted` satisfies this required permission using the hierarchy + /// Admin ⊇ Write ⊇ Read, otherwise [`false`] is returned. + #[allow(clippy::match_like_matches_macro)] + pub fn is_satisfied_by(&self, granted: &Permission) -> bool { + match (self, granted) { + (Permission::Read, _) => true, + (Permission::Write, Permission::Write | Permission::Admin) => true, + (Permission::Admin, Permission::Admin) => true, + _ => false, + } + } +} + +/// Tonic interceptor that attaches a bearer token to every outgoing request when one is present. +#[derive(Clone)] +pub struct BearerInterceptor { + /// The value of the `authorization` header. This is either `Bearer ` or [`None`] if no + /// token has been provided. The header is only attached to an outgoing request if this is not + /// [`None`]. + pub maybe_authorization: Option, +} + +impl BearerInterceptor { + /// Create a new [`BearerInterceptor`] that attaches the bearer token in `maybe_token` to every + /// outgoing request if `maybe_token` is not [`None`]. If `maybe_token` is [`Some`] but the + /// token is not ASCII, [`Status`] is returned. + pub fn try_new(maybe_token: Option<&str>) -> Result { + let maybe_authorization = maybe_token + .map(|token| { + format!("Bearer {token}") + .parse::() + .map_err(|error| Status::invalid_argument(format!("Token is not ASCII: {error}."))) + }) + .transpose()?; + + Ok(Self { + maybe_authorization, + }) + } +} + +impl Interceptor for BearerInterceptor { + fn call(&mut self, mut request: Request<()>) -> Result, Status> { + if let Some(authorization) = &self.maybe_authorization { + request + .metadata_mut() + .insert("authorization", authorization.clone()); + } + + Ok(request) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + // Tests for Permission. + #[test] + fn test_read_satisfied_by_read() { + assert!(Permission::Read.is_satisfied_by(&Permission::Read)); + } + + #[test] + fn test_read_satisfied_by_write() { + assert!(Permission::Read.is_satisfied_by(&Permission::Write)); + } + + #[test] + fn test_read_satisfied_by_admin() { + assert!(Permission::Read.is_satisfied_by(&Permission::Admin)); + } + + #[test] + fn test_write_not_satisfied_by_read() { + assert!(!Permission::Write.is_satisfied_by(&Permission::Read)); + } + + #[test] + fn test_write_satisfied_by_write() { + assert!(Permission::Write.is_satisfied_by(&Permission::Write)); + } + + #[test] + fn test_write_satisfied_by_admin() { + assert!(Permission::Write.is_satisfied_by(&Permission::Admin)); + } + + #[test] + fn test_admin_not_satisfied_by_read() { + assert!(!Permission::Admin.is_satisfied_by(&Permission::Read)); + } + + #[test] + fn test_admin_not_satisfied_by_write() { + assert!(!Permission::Admin.is_satisfied_by(&Permission::Write)); + } + + #[test] + fn test_admin_satisfied_by_admin() { + assert!(Permission::Admin.is_satisfied_by(&Permission::Admin)); + } + + // Tests for BearerInterceptor. + #[test] + fn test_try_new_bearer_interceptor_with_invalid_token() { + let result = BearerInterceptor::try_new(Some("sec\nret")); + + assert_eq!( + result.err().unwrap().to_string(), + "code: 'Client specified an invalid argument', \ + message: \"Token is not ASCII: failed to parse metadata value.\"" + ); + } + + #[test] + fn test_bearer_interceptor_call_attaches_authorization_header() { + let mut interceptor = BearerInterceptor::try_new(Some("secret")).unwrap(); + let request = interceptor.call(Request::new(())).unwrap(); + + assert_eq!( + request + .metadata() + .get("authorization") + .unwrap() + .to_str() + .unwrap(), + "Bearer secret" + ); + } + + #[test] + fn test_bearer_interceptor_call_without_token_attaches_nothing() { + let mut interceptor = BearerInterceptor::try_new(None).unwrap(); + let request = interceptor.call(Request::new(())).unwrap(); + + assert!(request.metadata().get("authorization").is_none()); + } +} diff --git a/crates/modelardb_client/Cargo.toml b/crates/modelardb_client/Cargo.toml index 5db6419f..e29e5b59 100644 --- a/crates/modelardb_client/Cargo.toml +++ b/crates/modelardb_client/Cargo.toml @@ -29,6 +29,7 @@ arrow = { workspace = true, features = ["prettyprint"] } bytes.workspace = true clap.workspace = true dirs.workspace = true +modelardb_auth = { path = "../modelardb_auth" } rustyline.workspace = true tokio.workspace = true tonic.workspace = true diff --git a/crates/modelardb_client/src/main.rs b/crates/modelardb_client/src/main.rs index 734f2905..0cab02b0 100644 --- a/crates/modelardb_client/src/main.rs +++ b/crates/modelardb_client/src/main.rs @@ -35,8 +35,10 @@ use arrow_flight::flight_service_client::FlightServiceClient; use arrow_flight::{Action, Criteria, FlightData, FlightDescriptor, Ticket, utils}; use bytes::Bytes; use clap::Parser; +use modelardb_auth::BearerInterceptor; use rustyline::Editor; use rustyline::history::FileHistory; +use tonic::codegen::InterceptedService; use tonic::transport::{Channel, Endpoint}; use tonic::{Request, Streaming}; @@ -46,6 +48,10 @@ use crate::helper::ClientHelper; /// Error to emit when the server does not provide a response when one is expected. const TRANSPORT_ERROR: &str = "transport error: no messages received."; +/// [`FlightServiceClient`] with a [`BearerInterceptor`] that attaches an authorization header. +type AuthenticatedFlightClient = + FlightServiceClient>; + /// Command line arguments for the ModelarDB client. #[derive(Parser)] #[command( @@ -64,6 +70,11 @@ struct ClientArgs { #[arg(long, default_value_t = 9999, env = "MODELARDB_PORT")] port: u16, + /// Bearer token for authenticating requests sent to the modelardbd instance. If not provided, + /// requests are sent without an authorization header. + #[arg(long, env = "MODELARDB_TOKEN")] + token: Option, + /// Path to a file containing SQL queries to execute. If not provided, an interactive /// read-eval-print loop is opened. query_file: Option, @@ -78,7 +89,7 @@ async fn main() -> Result<()> { let args = ClientArgs::parse(); // Execute the queries. - let flight_service_client = connect(&args.host, args.port).await?; + let flight_service_client = connect(&args.host, args.port, args.token).await?; if let Some(query_file) = args.query_file { execute_queries_from_a_file(flight_service_client, &query_file).await } else { @@ -86,17 +97,28 @@ async fn main() -> Result<()> { } } -/// Connect to the server at `host`:`port`. Returns [`ModelarDbClientError`] if a connection to the -/// server cannot be established. -async fn connect(host: &str, port: u16) -> Result> { +/// Connect to the server at `host`:`port` with an optional bearer `maybe_token`. Returns +/// [`ModelarDbClientError`] if a connection to the server cannot be established or the token is +/// not a valid ASCII metadata value. +async fn connect( + host: &str, + port: u16, + maybe_token: Option, +) -> Result { + let interceptor = BearerInterceptor::try_new(maybe_token.as_deref())?; + let address = format!("grpc://{host}:{port}"); let connection = Endpoint::new(address)?.connect().await?; - Ok(FlightServiceClient::new(connection)) + + Ok(FlightServiceClient::with_interceptor( + connection, + interceptor, + )) } /// Execute the commands and queries in `query_file`. async fn execute_queries_from_a_file( - mut flight_service_client: FlightServiceClient, + mut flight_service_client: AuthenticatedFlightClient, query_file: &StdPath, ) -> Result<()> { let file = File::open(query_file)?; @@ -123,7 +145,7 @@ async fn execute_queries_from_a_file( /// Execute commands and queries in a read-eval-print loop. async fn execute_queries_from_a_repl( - mut flight_service_client: FlightServiceClient, + mut flight_service_client: AuthenticatedFlightClient, ) -> Result<()> { // Create the read-eval-print loop. let mut editor = Editor::::new()?; @@ -158,7 +180,7 @@ async fn execute_queries_from_a_repl( /// Execute a command or a query. Returns [`ModelarDbClientError`] if the command or query could not /// be executed or their result could not be retrieved. async fn execute_and_print_command_or_query( - flight_service_client: &mut FlightServiceClient, + flight_service_client: &mut AuthenticatedFlightClient, command_or_query: &str, ) { let start_time = Instant::now(); @@ -182,7 +204,7 @@ async fn execute_and_print_command_or_query( /// * The command could not be executed. /// * The result could not be retrieved. async fn execute_command( - flight_service_client: &mut FlightServiceClient, + flight_service_client: &mut AuthenticatedFlightClient, command_and_argument: &str, ) -> Result<()> { let mut command_and_argument = command_and_argument.split(' '); @@ -255,7 +277,7 @@ async fn execute_command( /// Retrieve the names of the tables available on the server. Returns [`ModelarDbClientError`] if /// the request could not be performed or the tables names could not be retrieved. async fn retrieve_table_names( - flight_service_client: &mut FlightServiceClient, + flight_service_client: &mut AuthenticatedFlightClient, ) -> Result> { let criteria = Criteria { expression: Bytes::new(), @@ -286,7 +308,7 @@ async fn retrieve_table_names( /// Execute an action. Returns [`ModelarDbClientError`] if the action could not be executed. async fn execute_action( - flight_service_client: &mut FlightServiceClient, + flight_service_client: &mut AuthenticatedFlightClient, action_type: &str, action_body: &str, ) -> Result<()> { @@ -310,7 +332,7 @@ async fn execute_action( /// Execute a query and print each batch in the result set. Returns [`ModelarDbClientError`] if the /// query could not be executed or the batches in the result set could not be printed. async fn execute_query_and_print_result( - flight_service_client: &mut FlightServiceClient, + flight_service_client: &mut AuthenticatedFlightClient, query: &str, ) -> Result<()> { // Execute the query. diff --git a/crates/modelardb_embedded/Cargo.toml b/crates/modelardb_embedded/Cargo.toml index d3b89eb0..530030c4 100644 --- a/crates/modelardb_embedded/Cargo.toml +++ b/crates/modelardb_embedded/Cargo.toml @@ -31,6 +31,7 @@ async-trait.workspace = true datafusion.workspace = true deltalake = { workspace = true, features = ["datafusion", "s3"] } futures.workspace = true +modelardb_auth = { path = "../modelardb_auth" } modelardb_compression = { path = "../modelardb_compression" } modelardb_storage = { path = "../modelardb_storage" } modelardb_types = { path = "../modelardb_types" } diff --git a/crates/modelardb_embedded/bindings/c/modelardb_embedded.h b/crates/modelardb_embedded/bindings/c/modelardb_embedded.h index f8f5f1b5..9bbb5b7d 100644 --- a/crates/modelardb_embedded/bindings/c/modelardb_embedded.h +++ b/crates/modelardb_embedded/bindings/c/modelardb_embedded.h @@ -95,7 +95,7 @@ void* modelardb_embedded_open_azure(const char* account_name_ptr, const char* container_name_ptr); // Connect to a ModelarDB server at the given Arrow Flight URL. -void* modelardb_embedded_connect(const char* url_ptr); +void* modelardb_embedded_connect(const char* url_ptr, const char* maybe_token_ptr); // Close and deallocate the data folder or client. int modelardb_embedded_close(void* maybe_operations_ptr, diff --git a/crates/modelardb_embedded/bindings/python/modelardb/operations.py b/crates/modelardb_embedded/bindings/python/modelardb/operations.py index b82c8791..d66883ee 100644 --- a/crates/modelardb_embedded/bindings/python/modelardb/operations.py +++ b/crates/modelardb_embedded/bindings/python/modelardb/operations.py @@ -221,17 +221,20 @@ def open_azure(cls, account_name: str, access_key: str, container_name: str): return self @classmethod - def connect(cls, url: str): + def connect(cls, url: str, token: None | str = None): """Create a connection to an :obj:`Operations` node. :param url: The URL of the ModelarDB node to connect to. :type url: str + :param token: A bearer token to attach to every request, or ``None`` if no authentication is required. + :type token: str, optional """ self: Operations = cls() url_ptr = ffi.new("char[]", bytes(url, "UTF-8")) + maybe_token_ptr = ffi.new("char[]", bytes(token, "UTF-8")) if token else ffi.NULL - self.__operations_ptr = self.__library.modelardb_embedded_connect(url_ptr) + self.__operations_ptr = self.__library.modelardb_embedded_connect(url_ptr, maybe_token_ptr) self.__is_data_folder = False if self.__operations_ptr == ffi.NULL: @@ -787,10 +790,12 @@ def open_azure(account_name: str, access_key: str, container_name: str) -> Opera return Operations.open_azure(account_name, access_key, container_name) -def connect(url: str) -> Operations: +def connect(url: str, token: None | str = None) -> Operations: """Create a connection to an :obj:`Operations` node. :param url: The URL of the ModelarDB node to connect to. :type url: str + :param token: A bearer token to attach to every request, or ``None`` if no authentication is required. + :type token: str, optional """ - return Operations.connect(url) + return Operations.connect(url, token) diff --git a/crates/modelardb_embedded/src/capi.rs b/crates/modelardb_embedded/src/capi.rs index 6e782065..eec6d78e 100644 --- a/crates/modelardb_embedded/src/capi.rs +++ b/crates/modelardb_embedded/src/capi.rs @@ -199,18 +199,23 @@ unsafe fn open_azure( /// Creates a [`Client`] that is connected to the Apache Arrow Flight server URL in `url_ptr` /// and returns a pointer to the [`Client`] or a zero-initialized pointer if an error occurs. -/// Assumes `url_ptr` points to a valid C string. +/// Assumes `url_ptr` points to a valid C string and `maybe_token_ptr` points to a valid C string +/// or is null if no bearer token should be attached to requests. #[unsafe(no_mangle)] -pub unsafe extern "C" fn modelardb_embedded_connect(url_ptr: *const c_char) -> *const c_void { - let maybe_client = unsafe { connect(url_ptr) }; +pub unsafe extern "C" fn modelardb_embedded_connect( + url_ptr: *const c_char, + maybe_token_ptr: *const c_char, +) -> *const c_void { + let maybe_client = unsafe { connect(url_ptr, maybe_token_ptr) }; set_error_and_return_value_ptr(maybe_client) } /// See documentation for [`modelardb_embedded_connect()`]. -unsafe fn connect(url_ptr: *const c_char) -> Result { +unsafe fn connect(url_ptr: *const c_char, maybe_token_ptr: *const c_char) -> Result { let url_str = unsafe { c_char_ptr_to_str(url_ptr)? }; + let maybe_token = unsafe { c_char_ptr_to_maybe_str(maybe_token_ptr)? }; - TOKIO_RUNTIME.block_on(Client::connect(url_str)) + TOKIO_RUNTIME.block_on(Client::connect(url_str, maybe_token)) } /// Moves the value in `maybe_value` to a [`Box`] and returns a pointer to it if `maybe_value` is diff --git a/crates/modelardb_embedded/src/operations/client.rs b/crates/modelardb_embedded/src/operations/client.rs index a4d74ed4..b71d93b4 100644 --- a/crates/modelardb_embedded/src/operations/client.rs +++ b/crates/modelardb_embedded/src/operations/client.rs @@ -34,6 +34,8 @@ use datafusion::error::DataFusionError; use datafusion::execution::RecordBatchStream; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use futures::{StreamExt, TryStreamExt, stream}; +use modelardb_auth::BearerInterceptor; +use tonic::codegen::InterceptedService; use tonic::transport::{Channel, Endpoint}; use tonic::{Request, Status}; @@ -48,15 +50,18 @@ use crate::{Aggregate, TableType}; #[derive(Clone)] pub struct Client { /// Apache Arrow Flight client connected to the Apache Arrow Flight server of the ModelarDB node. - pub(crate) flight_client: FlightServiceClient, + pub(crate) flight_client: FlightServiceClient>, } impl Client { - /// Create a new [`Client`] that is connected to the node with `url`. If a connection - /// to the node could not be established, [`ModelarDbEmbeddedError`] is returned. - pub async fn connect(url: &str) -> Result { + /// Create a new [`Client`] that is connected to the node with `url`. If `maybe_token` is + /// provided, it is attached as a bearer token to every request. If a connection to the node + /// could not be established, [`ModelarDbEmbeddedError`] is returned. + pub async fn connect(url: &str, maybe_token: Option<&str>) -> Result { + let interceptor = BearerInterceptor::try_new(maybe_token)?; + let connection = Endpoint::new(url.to_owned())?.connect().await?; - let flight_client = FlightServiceClient::new(connection); + let flight_client = FlightServiceClient::with_interceptor(connection, interceptor); Ok(Client { flight_client }) } diff --git a/crates/modelardb_embedded/src/operations/data_folder.rs b/crates/modelardb_embedded/src/operations/data_folder.rs index 2b45bbfe..bb293dbf 100644 --- a/crates/modelardb_embedded/src/operations/data_folder.rs +++ b/crates/modelardb_embedded/src/operations/data_folder.rs @@ -603,6 +603,7 @@ mod tests { use datafusion::physical_expr::{LexOrdering, PhysicalSortExpr}; use datafusion::physical_plan::expressions::Column; use datafusion::physical_plan::sorts::sort; + use modelardb_auth::BearerInterceptor; use modelardb_types::types::{ ArrowTimestamp, ArrowValue, ErrorBound, GeneratedColumn, TimeSeriesTableMetadata, TimestampArray, ValueArray, @@ -2833,9 +2834,14 @@ mod tests { } fn lazy_modelardb_client() -> Client { + let interceptor = BearerInterceptor { + maybe_authorization: None, + }; + Client { - flight_client: FlightServiceClient::new( + flight_client: FlightServiceClient::with_interceptor( Channel::from_static("localhost").connect_lazy(), + interceptor, ), } } diff --git a/crates/modelardb_server/Cargo.toml b/crates/modelardb_server/Cargo.toml index 6f15b1d0..ca1ba761 100644 --- a/crates/modelardb_server/Cargo.toml +++ b/crates/modelardb_server/Cargo.toml @@ -34,6 +34,9 @@ dashmap.workspace = true datafusion.workspace = true deltalake = { workspace = true, features = ["datafusion"] } futures.workspace = true +http.workspace = true +http-body-util.workspace = true +modelardb_auth = { path = "../modelardb_auth" } modelardb_compression = { path = "../modelardb_compression" } modelardb_storage = { path = "../modelardb_storage" } modelardb_types = { path = "../modelardb_types" } @@ -42,10 +45,12 @@ prost.workspace = true rand.workspace = true serde.workspace = true snmalloc-rs = { workspace = true, features = ["build_cc"] } +sqlparser.workspace = true tokio = { workspace = true, features = ["rt-multi-thread", "signal"] } tokio-stream.workspace = true toml.workspace = true tonic.workspace = true +tower.workspace = true # Log is a dependency so the compile time filters for log and tracing can be set to the same values. log = { workspace = true, features = ["max_level_debug", "release_max_level_info"] } @@ -54,6 +59,7 @@ tracing-subscriber.workspace = true [dev-dependencies] modelardb_test = { path = "../modelardb_test" } +modelardb_auth = { path = "../modelardb_auth", features = ["testing"] } proptest.workspace = true tempfile.workspace = true diff --git a/crates/modelardb_server/src/cluster.rs b/crates/modelardb_server/src/cluster.rs index 8e99d370..10959a02 100644 --- a/crates/modelardb_server/src/cluster.rs +++ b/crates/modelardb_server/src/cluster.rs @@ -30,7 +30,7 @@ use modelardb_types::types::{Node, TimeSeriesTableMetadata}; use rand::rng; use rand::seq::IteratorRandom; use tonic::Request; -use tonic::metadata::{Ascii, MetadataValue}; +use tonic::metadata::AsciiMetadataValue; use tonic::transport::Endpoint; use crate::context::Context; @@ -44,7 +44,7 @@ pub(crate) struct Cluster { node: Node, /// Key identifying the cluster. The key is used to validate communication within the cluster /// between nodes. - key: MetadataValue, + key: AsciiMetadataValue, /// The remote data folder that each node in the cluster should be synchronized with. /// When a table is created, dropped, vacuumed, or truncated, it is done in the /// remote data folder first. @@ -66,7 +66,8 @@ impl Cluster { let key = remote_data_folder.cluster_key().await?.to_string(); // Convert the key to a MetadataValue since it is used in tonic requests. - let key = MetadataValue::from_str(&key).expect("UUID Version 4 should be valid ASCII."); + let key = + AsciiMetadataValue::from_str(&key).expect("UUID Version 4 should be valid ASCII."); Ok(Self { node, @@ -81,7 +82,7 @@ impl Cluster { } /// Return the key identifying the cluster. - pub(crate) fn key(&self) -> &MetadataValue { + pub(crate) fn key(&self) -> &AsciiMetadataValue { &self.key } diff --git a/crates/modelardb_server/src/main.rs b/crates/modelardb_server/src/main.rs index 9ff6c8bd..9a015b3e 100644 --- a/crates/modelardb_server/src/main.rs +++ b/crates/modelardb_server/src/main.rs @@ -159,7 +159,7 @@ async fn main() -> Result<()> { context.replay_write_ahead_log().await?; // Start the Apache Arrow Flight interface. - remote::start_apache_arrow_flight_server(context, args.port).await?; + remote::start_apache_arrow_flight_server(context, None, args.port).await?; Ok(()) } diff --git a/crates/modelardb_server/src/remote/auth_layer.rs b/crates/modelardb_server/src/remote/auth_layer.rs new file mode 100644 index 00000000..0443f927 --- /dev/null +++ b/crates/modelardb_server/src/remote/auth_layer.rs @@ -0,0 +1,606 @@ +/* Copyright 2026 The ModelarDB Contributors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +//! Tower middleware layer that enforces authentication and authorization for all incoming Apache +//! Arrow Flight requests. The layer runs before [`FlightServiceHandler`](super::FlightServiceHandler) +//! and checks the cluster key or bearer token before the request reaches the handler. For DoGet +//! requests the SQL ticket is decoded, and the required permission is determined from the parsed +//! statement. + +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use arrow_flight::Ticket; +use http::{Request, Response}; +use http_body_util::{BodyExt, Full}; +use modelardb_auth::Permission; +use modelardb_auth::authenticator::Authenticator; +use modelardb_storage::parser::{self, ModelarDbStatement}; +use prost::Message; +use sqlparser::ast::Statement; +use tonic::Status; +use tonic::body::Body; +use tonic::metadata::{AsciiMetadataValue, MetadataMap}; +use tower::{Layer, Service}; + +use crate::remote::error_to_status_invalid_argument; + +const LIST_FLIGHTS_PATH: &str = "/arrow.flight.protocol.FlightService/ListFlights"; +const GET_FLIGHT_INFO_PATH: &str = "/arrow.flight.protocol.FlightService/GetFlightInfo"; +const GET_SCHEMA_PATH: &str = "/arrow.flight.protocol.FlightService/GetSchema"; +const DO_GET_PATH: &str = "/arrow.flight.protocol.FlightService/DoGet"; +const DO_PUT_PATH: &str = "/arrow.flight.protocol.FlightService/DoPut"; +const DO_ACTION_PATH: &str = "/arrow.flight.protocol.FlightService/DoAction"; +const LIST_ACTIONS_PATH: &str = "/arrow.flight.protocol.FlightService/ListActions"; + +/// [`Layer`] that enforces authentication and authorization on all incoming Apache Arrow Flight +/// requests. +#[derive(Clone)] +pub(super) struct AuthLayer { + /// The [`Authenticator`] to use once the required permission has been determined, or [`None`] + /// if authentication is disabled and every non-cluster request should be allowed. + maybe_authenticator: Option>, + /// The cluster key to use for validating internal cluster requests or [`None`] if running a + /// single-node server. + maybe_cluster_key: Option, +} + +impl AuthLayer { + pub(super) fn new( + maybe_authenticator: Option>, + maybe_cluster_key: Option, + ) -> Self { + Self { + maybe_authenticator, + maybe_cluster_key, + } + } +} + +impl Layer for AuthLayer { + type Service = AuthService; + + fn layer(&self, inner: S) -> AuthService { + AuthService { + inner, + maybe_authenticator: self.maybe_authenticator.clone(), + maybe_cluster_key: self.maybe_cluster_key.clone(), + } + } +} + +/// Middleware layer that enforces authentication and authorization on all incoming Apache Arrow +/// Flight requests. +#[derive(Clone)] +pub(super) struct AuthService { + /// The [`FlightServiceHandler`](super::FlightServiceHandler) that processes the request after + /// authorization succeeds. + inner: S, + /// The [`Authenticator`] to use once the required permission has been determined, or [`None`] + /// if authentication is disabled and every non-cluster request should be allowed. + maybe_authenticator: Option>, + /// The cluster key to use for validating internal cluster requests or [`None`] if running a + /// single-node server. + maybe_cluster_key: Option, +} + +impl Service> for AuthService +where + S: Service, Response = Response> + Clone + Send + 'static, + S::Future: Send + 'static, +{ + type Response = S::Response; + type Error = S::Error; + type Future = Pin> + Send>>; + + fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { + self.inner.poll_ready(cx) + } + + fn call(&mut self, request: Request) -> Self::Future { + // A cloned service may not be ready and panics if called before poll_ready. Use mem::replace + // to take the ready service out of self.inner and put a fresh clone in its place, which + // will be polled for readiness on the next request. + // + // See https://docs.rs/tower/latest/tower/trait.Service.html#be-careful-when-cloning-inner-services. + let clone = self.inner.clone(); + let mut inner = std::mem::replace(&mut self.inner, clone); + + let maybe_authenticator = self.maybe_authenticator.clone(); + let maybe_cluster_key = self.maybe_cluster_key.clone(); + + Box::pin(async move { + match authorize(request, maybe_authenticator.as_deref(), &maybe_cluster_key).await { + Ok(request) => inner.call(request).await, + Err(status) => Ok(status.into_http()), + } + }) + } +} + +/// Full authorization flow. Return the (possibly reconstructed) request on success, or a +/// [`Status`] on failure. +async fn authorize( + request: Request, + maybe_authenticator: Option<&dyn Authenticator>, + maybe_cluster_key: &Option, +) -> Result, Status> { + let metadata = MetadataMap::from_headers(request.headers().clone()); + + // Cluster key check must happen before the Authenticator since internal cluster requests bypass + // auth entirely. + if let Some(request_key) = metadata.get("x-cluster-key") { + return match maybe_cluster_key { + Some(key) if key == request_key => Ok(request), + Some(_) => Err(Status::internal("Invalid cluster key.")), + None => Err(Status::internal("Cluster key sent to single-node server.")), + }; + } + + // When no Authenticator is configured, authentication is disabled and every non-cluster request + // is allowed. + let Some(authenticator) = maybe_authenticator else { + return Ok(request); + }; + + // Decode the ticket and parse the SQL to determine the required permission. + let path = request.uri().path().to_owned(); + if path == DO_GET_PATH { + return authorize_do_get(request, authenticator, &metadata).await; + } + + // For all other endpoints the path determines the permission. + let required_permission = match path.as_str() { + LIST_FLIGHTS_PATH | GET_FLIGHT_INFO_PATH | GET_SCHEMA_PATH => Permission::Read, + DO_PUT_PATH => Permission::Write, + DO_ACTION_PATH | LIST_ACTIONS_PATH => Permission::Admin, + _ => { + return Err(Status::invalid_argument("Unknown path.")); + } + }; + + authenticator + .authorize(&metadata, required_permission) + .await?; + + Ok(request) +} + +/// Buffer the DoGet body, decode the gRPC [`Ticket`] protobuf, parse the SQL, determine the +/// required permission, authorize, then reconstruct the request byte-for-byte. +async fn authorize_do_get( + request: Request, + authenticator: &dyn Authenticator, + metadata: &MetadataMap, +) -> Result, Status> { + let (parts, body) = request.into_parts(); + + // Collect the full body. + let bytes = body + .collect() + .await + .map_err(|_| Status::invalid_argument("Failed to unpack request body."))? + .to_bytes(); + + // gRPC data frames have a 1-byte compression flag, a 4-byte length, and an N bytes message as + // defined in https://github.com/grpc/grpc/blob/master/doc/PROTOCOL-HTTP2.md. + if bytes.len() < 5 { + return Err(Status::invalid_argument( + "Request body too short to be a valid gRPC message.", + )); + } + + let ticket = Ticket::decode(&bytes[5..]).map_err(error_to_status_invalid_argument)?; + let sql = str::from_utf8(&ticket.ticket).map_err(error_to_status_invalid_argument)?; + + let statement = + parser::tokenize_and_parse_sql_statement(sql).map_err(error_to_status_invalid_argument)?; + + authenticator + .authorize(metadata, permission_for_statement(&statement)) + .await?; + + // Reconstruct the request with the original bytes so the server receives it intact. + Ok(Request::from_parts(parts, Body::new(Full::new(bytes)))) +} + +/// Map a parsed [`ModelarDbStatement`] to the required [`Permission`]. +fn permission_for_statement(statement: &ModelarDbStatement) -> Permission { + match statement { + ModelarDbStatement::CreateNormalTable { .. } => Permission::Admin, + ModelarDbStatement::CreateTimeSeriesTable(_) => Permission::Admin, + ModelarDbStatement::DropTable(_) => Permission::Admin, + ModelarDbStatement::TruncateTable(_, _) => Permission::Admin, + ModelarDbStatement::Vacuum(_, _, _) => Permission::Admin, + ModelarDbStatement::Optimize(_, _, _) => Permission::Admin, + ModelarDbStatement::IncludeSelect(_, _) => Permission::Read, + ModelarDbStatement::Statement(statement) => match statement { + Statement::Insert(_) => Permission::Write, + Statement::Query(_) | Statement::Explain { .. } => Permission::Read, + _ => Permission::Admin, + }, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use modelardb_auth::authenticator::mock::MockAuthenticator; + use modelardb_test::table::{NORMAL_TABLE_SQL, TIME_SERIES_TABLE_SQL}; + + const CLUSTER_KEY: &str = "cluster_key"; + + #[tokio::test] + async fn test_authorize_multi_node_with_valid_cluster_key_bypasses_authenticator() { + let authenticator = Arc::new(MockAuthenticator::new()); + + let request = empty_request_with_cluster_key(DO_PUT_PATH, CLUSTER_KEY); + let cluster_key = Some(AsciiMetadataValue::from_static(CLUSTER_KEY)); + + let result = authorize(request, Some(&*authenticator), &cluster_key).await; + + assert!(result.is_ok()); + assert!(authenticator.permissions().is_empty()); + } + + #[tokio::test] + async fn test_authorize_multi_node_with_invalid_cluster_key() { + let authenticator = Arc::new(MockAuthenticator::new()); + + let request = empty_request_with_cluster_key(LIST_FLIGHTS_PATH, "invalid_key"); + let cluster_key = Some(AsciiMetadataValue::from_static(CLUSTER_KEY)); + + let result = authorize(request, Some(&*authenticator), &cluster_key).await; + + assert_eq!( + result.unwrap_err().to_string(), + "code: 'Internal error', message: \"Invalid cluster key.\"" + ); + } + + #[tokio::test] + async fn test_authorize_single_node_with_cluster_key() { + let authenticator = Arc::new(MockAuthenticator::new()); + let request = empty_request_with_cluster_key(LIST_FLIGHTS_PATH, CLUSTER_KEY); + + let result = authorize(request, Some(&*authenticator), &None).await; + + assert_eq!( + result.unwrap_err().to_string(), + "code: 'Internal error', message: \"Cluster key sent to single-node server.\"" + ); + } + + fn empty_request_with_cluster_key(path: &str, key: &str) -> Request { + Request::builder() + .uri(path) + .header("x-cluster-key", key) + .body(Body::empty()) + .unwrap() + } + + #[tokio::test] + async fn test_authorize_without_authenticator_allows_request_without_parsing() { + // A body that would fail ticket decoding if authorize_do_get() ran. + let request = raw_frame_request(&[0xFF, 0xFF, 0xFF, 0xFF]); + + let result = authorize(request, None, &None).await; + + assert!(result.is_ok()); + } + + #[tokio::test] + async fn test_authorize_list_flights_calls_authenticator_with_read() { + let authenticator = Arc::new(MockAuthenticator::new()); + let request = empty_request(LIST_FLIGHTS_PATH); + + let result = authorize(request, Some(&*authenticator), &None).await; + + assert!(result.is_ok()); + assert_eq!(authenticator.permissions(), vec![Permission::Read]); + } + + #[tokio::test] + async fn test_authorize_get_flight_info_calls_authenticator_with_read() { + let authenticator = Arc::new(MockAuthenticator::new()); + let request = empty_request(GET_FLIGHT_INFO_PATH); + + let result = authorize(request, Some(&*authenticator), &None).await; + + assert!(result.is_ok()); + assert_eq!(authenticator.permissions(), vec![Permission::Read]); + } + + #[tokio::test] + async fn test_authorize_get_schema_calls_authenticator_with_read() { + let authenticator = Arc::new(MockAuthenticator::new()); + let request = empty_request(GET_SCHEMA_PATH); + + let result = authorize(request, Some(&*authenticator), &None).await; + + assert!(result.is_ok()); + assert_eq!(authenticator.permissions(), vec![Permission::Read]); + } + + #[tokio::test] + async fn test_authorize_do_put_calls_authenticator_with_write() { + let authenticator = Arc::new(MockAuthenticator::new()); + let request = empty_request(DO_PUT_PATH); + + let result = authorize(request, Some(&*authenticator), &None).await; + + assert!(result.is_ok()); + assert_eq!(authenticator.permissions(), vec![Permission::Write]); + } + + #[tokio::test] + async fn test_authorize_do_action_calls_authenticator_with_admin() { + let authenticator = Arc::new(MockAuthenticator::new()); + let request = empty_request(DO_ACTION_PATH); + + let result = authorize(request, Some(&*authenticator), &None).await; + + assert!(result.is_ok()); + assert_eq!(authenticator.permissions(), vec![Permission::Admin]); + } + + #[tokio::test] + async fn test_authorize_list_actions_calls_authenticator_with_admin() { + let authenticator = Arc::new(MockAuthenticator::new()); + let request = empty_request(LIST_ACTIONS_PATH); + + let result = authorize(request, Some(&*authenticator), &None).await; + + assert!(result.is_ok()); + assert_eq!(authenticator.permissions(), vec![Permission::Admin]); + } + + #[tokio::test] + async fn test_authorize_do_get_with_create_normal_table_calls_authenticator_with_admin() { + let authenticator = Arc::new(MockAuthenticator::new()); + let request = do_get_request(NORMAL_TABLE_SQL); + + let result = authorize(request, Some(&*authenticator), &None).await; + + assert!(result.is_ok()); + assert_eq!(authenticator.permissions(), vec![Permission::Admin]); + } + + #[tokio::test] + async fn test_authorize_do_get_with_create_time_series_table_calls_authenticator_with_admin() { + let authenticator = Arc::new(MockAuthenticator::new()); + let request = do_get_request(TIME_SERIES_TABLE_SQL); + + let result = authorize(request, Some(&*authenticator), &None).await; + + assert!(result.is_ok()); + assert_eq!(authenticator.permissions(), vec![Permission::Admin]); + } + + #[tokio::test] + async fn test_authorize_do_get_with_drop_table_calls_authenticator_with_admin() { + let authenticator = Arc::new(MockAuthenticator::new()); + let request = do_get_request("DROP TABLE test_table"); + + let result = authorize(request, Some(&*authenticator), &None).await; + + assert!(result.is_ok()); + assert_eq!(authenticator.permissions(), vec![Permission::Admin]); + } + + #[tokio::test] + async fn test_authorize_do_get_with_truncate_table_calls_authenticator_with_admin() { + let authenticator = Arc::new(MockAuthenticator::new()); + let request = do_get_request("TRUNCATE test_table"); + + let result = authorize(request, Some(&*authenticator), &None).await; + + assert!(result.is_ok()); + assert_eq!(authenticator.permissions(), vec![Permission::Admin]); + } + + #[tokio::test] + async fn test_authorize_do_get_with_vacuum_calls_authenticator_with_admin() { + let authenticator = Arc::new(MockAuthenticator::new()); + let request = do_get_request("VACUUM test_table"); + + let result = authorize(request, Some(&*authenticator), &None).await; + + assert!(result.is_ok()); + assert_eq!(authenticator.permissions(), vec![Permission::Admin]); + } + + #[tokio::test] + async fn test_authorize_do_get_with_optimize_calls_authenticator_with_admin() { + let authenticator = Arc::new(MockAuthenticator::new()); + let request = do_get_request("OPTIMIZE test_table"); + + let result = authorize(request, Some(&*authenticator), &None).await; + + assert!(result.is_ok()); + assert_eq!(authenticator.permissions(), vec![Permission::Admin]); + } + + #[tokio::test] + async fn test_authorize_do_get_with_include_select_calls_authenticator_with_read() { + let authenticator = Arc::new(MockAuthenticator::new()); + let request = do_get_request("INCLUDE 'grpc://192.168.1.2:9999' SELECT * FROM test_table"); + + let result = authorize(request, Some(&*authenticator), &None).await; + + assert!(result.is_ok()); + assert_eq!(authenticator.permissions(), vec![Permission::Read]); + } + + #[tokio::test] + async fn test_authorize_do_get_with_insert_calls_authenticator_with_write() { + let authenticator = Arc::new(MockAuthenticator::new()); + let request = do_get_request("INSERT INTO test_table VALUES (1, 2, 3)"); + + let result = authorize(request, Some(&*authenticator), &None).await; + + assert!(result.is_ok()); + assert_eq!(authenticator.permissions(), vec![Permission::Write]); + } + + #[tokio::test] + async fn test_authorize_do_get_with_select_calls_authenticator_with_read() { + let authenticator = Arc::new(MockAuthenticator::new()); + let request = do_get_request("SELECT * FROM test_table"); + + let result = authorize(request, Some(&*authenticator), &None).await; + + assert!(result.is_ok()); + assert_eq!(authenticator.permissions(), vec![Permission::Read]); + } + + #[tokio::test] + async fn test_authorize_do_get_with_explain_calls_authenticator_with_read() { + let authenticator = Arc::new(MockAuthenticator::new()); + let request = do_get_request("EXPLAIN SELECT * FROM test_table"); + + let result = authorize(request, Some(&*authenticator), &None).await; + + assert!(result.is_ok()); + assert_eq!(authenticator.permissions(), vec![Permission::Read]); + } + + #[tokio::test] + async fn test_authorize_do_get_body_is_reconstructed_intact() { + let authenticator = Arc::new(MockAuthenticator::new()); + let sql = "SELECT * FROM test_table"; + + // Capture the original body bytes before calling authorize(). + let original_request = do_get_request(sql); + let (_, original_body) = original_request.into_parts(); + let original_bytes = original_body.collect().await.unwrap().to_bytes(); + + let request = do_get_request(sql); + let result = authorize(request, Some(&*authenticator), &None).await; + + let (_, reconstructed_body) = result.unwrap().into_parts(); + let reconstructed_bytes = reconstructed_body.collect().await.unwrap().to_bytes(); + + assert_eq!(original_bytes, reconstructed_bytes); + } + + #[tokio::test] + async fn test_authorize_do_get_with_body_too_short() { + let authenticator = Arc::new(MockAuthenticator::new()); + let request = Request::builder() + .uri(DO_GET_PATH) + .body(Body::new(Full::new(bytes::Bytes::from(vec![0u8; 4])))) + .unwrap(); + + let result = authorize(request, Some(&*authenticator), &None).await; + + assert_eq!( + result.unwrap_err().to_string(), + "code: 'Client specified an invalid argument', \ + message: \"Request body too short to be a valid gRPC message.\"" + ); + } + + #[tokio::test] + async fn test_authorize_do_get_with_invalid_protobuf() { + let authenticator = Arc::new(MockAuthenticator::new()); + + // Valid 5-byte gRPC frame header but invalid protobuf bytes in the message. + let request = raw_frame_request(&[0xFF, 0xFF, 0xFF, 0xFF]); + + let result = authorize(request, Some(&*authenticator), &None).await; + + assert_eq!( + result.unwrap_err().to_string(), + "code: 'Client specified an invalid argument', \ + message: \"failed to decode Protobuf message: invalid varint\"" + ); + } + + #[tokio::test] + async fn test_authorize_do_get_with_non_utf8_ticket() { + let authenticator = Arc::new(MockAuthenticator::new()); + + // Encode a Ticket with invalid UTF-8 bytes. + let request = ticket_frame_request(vec![0xFF, 0xFE]); + + let result = authorize(request, Some(&*authenticator), &None).await; + + assert_eq!( + result.unwrap_err().to_string(), + "code: 'Client specified an invalid argument', \ + message: \"invalid utf-8 sequence of 1 bytes from index 0\"" + ); + } + + #[tokio::test] + async fn test_authorize_do_get_with_unparseable_sql() { + let authenticator = Arc::new(MockAuthenticator::new()); + let request = do_get_request("invalid sql"); + + let result = authorize(request, Some(&*authenticator), &None).await; + + assert_eq!( + result.unwrap_err().to_string(), + "code: 'Client specified an invalid argument', \ + message: \"Parser Error: sql parser error: Expected: an SQL statement, found: invalid at Line: 1, Column: 1\"" + ); + } + + fn do_get_request(sql: &str) -> Request { + ticket_frame_request(sql.as_bytes().to_vec()) + } + + fn ticket_frame_request(ticket_bytes: Vec) -> Request { + let ticket = Ticket { + ticket: ticket_bytes.into(), + }; + + raw_frame_request(&ticket.encode_to_vec()) + } + + fn raw_frame_request(message_bytes: &[u8]) -> Request { + // Construct a gRPC frame with the 1-byte compression flag, 4-byte message length, and message. + let mut frame = Vec::with_capacity(5 + message_bytes.len()); + frame.push(0u8); + frame.extend_from_slice(&(message_bytes.len() as u32).to_be_bytes()); + frame.extend_from_slice(message_bytes); + + Request::builder() + .uri(DO_GET_PATH) + .body(Body::new(Full::new(bytes::Bytes::from(frame)))) + .unwrap() + } + + #[tokio::test] + async fn test_authorize_unknown_path() { + let authenticator = Arc::new(MockAuthenticator::new()); + let request = empty_request("/unknown/path"); + + let result = authorize(request, Some(&*authenticator), &None).await; + + assert_eq!( + result.unwrap_err().to_string(), + "code: 'Client specified an invalid argument', message: \"Unknown path.\"" + ); + } + + fn empty_request(path: &str) -> Request { + Request::builder().uri(path).body(Body::empty()).unwrap() + } +} diff --git a/crates/modelardb_server/src/remote.rs b/crates/modelardb_server/src/remote/mod.rs similarity index 95% rename from crates/modelardb_server/src/remote.rs rename to crates/modelardb_server/src/remote/mod.rs index a6a0bdb1..7558da13 100644 --- a/crates/modelardb_server/src/remote.rs +++ b/crates/modelardb_server/src/remote/mod.rs @@ -17,6 +17,8 @@ //! [`FlightServiceHandler`]. An Apache Arrow Flight server that process requests //! using [`FlightServiceHandler`] can be started with [`start_apache_arrow_flight_server()`]. +mod auth_layer; + use std::collections::HashMap; use std::error::Error; use std::net::SocketAddr; @@ -43,6 +45,7 @@ use datafusion::physical_plan::{EmptyRecordBatchStream, SendableRecordBatchStrea use deltalake::arrow::datatypes::Schema; use futures::StreamExt; use futures::stream::{self, BoxStream, SelectAll}; +use modelardb_auth::authenticator::Authenticator; use modelardb_storage::parser::{self, ModelarDbStatement}; use modelardb_types::flight::protocol; use modelardb_types::functions; @@ -51,25 +54,39 @@ use prost::Message; use tokio::sync::mpsc::{self, Sender}; use tokio::task; use tokio_stream::wrappers::ReceiverStream; -use tonic::metadata::MetadataMap; +use tonic::metadata::{AsciiMetadataValue, MetadataMap}; use tonic::transport::{Endpoint, Server}; use tonic::{Request, Response, Status, Streaming}; use tracing::{debug, error, info}; use crate::ClusterMode; -use crate::cluster::Cluster; use crate::context::Context; use crate::error::{ModelarDbServerError, Result}; +use crate::remote::auth_layer::AuthLayer; /// Start an Apache Arrow Flight server on 0.0.0.0:`port` that passes `context` to the methods that -/// process the requests through [`FlightServiceHandler`]. -pub async fn start_apache_arrow_flight_server(context: Arc, port: u16) -> Result<()> { +/// process the requests through [`FlightServiceHandler`]. All requests are passed through the +/// [`AuthLayer`], which authenticates them using `maybe_authenticator` before they are passed to +/// the [`FlightServiceHandler`]. If `maybe_authenticator` is [`None`], authentication is disabled, +/// and every request that is not an internal cluster request is allowed. +pub async fn start_apache_arrow_flight_server( + context: Arc, + maybe_authenticator: Option>, + port: u16, +) -> Result<()> { let localhost_with_port = "0.0.0.0:".to_owned() + &port.to_string(); let localhost_with_port: SocketAddr = localhost_with_port.parse().map_err(|error| { ModelarDbServerError::InvalidArgument(format!( "Unable to parse {localhost_with_port}: {error}" )) })?; + + let maybe_cluster_key = match context.configuration_manager.read().await.cluster_mode() { + ClusterMode::MultiNode(cluster) => Some(cluster.key().clone()), + ClusterMode::SingleNode => None, + }; + + let auth_layer = AuthLayer::new(maybe_authenticator, maybe_cluster_key); let handler = FlightServiceHandler::new(context); // Increase the maximum message size from 4 MiB to 16 MiB to allow bulk-loading larger batches. @@ -79,6 +96,7 @@ pub async fn start_apache_arrow_flight_server(context: Arc, port: u16) info!("Starting Apache Arrow Flight on {}.", localhost_with_port); Server::builder() + .layer(auth_layer) .add_service(flight_service_server) .serve(localhost_with_port) .await @@ -91,6 +109,7 @@ pub async fn start_apache_arrow_flight_server(context: Arc, port: u16) async fn execute_query_at_addresses_and_union( sql: &str, addresses: Vec, + maybe_authorization: Option, local_sendable_record_batch_stream: Pin>, ) -> Result>> { // Remove INCLUDE address+ as sqlparser seems incapable of doing so. @@ -107,8 +126,12 @@ async fn execute_query_at_addresses_and_union( unioned_sendable_record_batch_streams.push(local_sendable_record_batch_stream); for address in addresses { - let remote_sendable_record_batch_stream = - execute_query_at_address(sql_select.to_owned(), address.to_owned()).await?; + let remote_sendable_record_batch_stream = execute_query_at_address( + sql_select.to_owned(), + address.to_owned(), + maybe_authorization.clone(), + ) + .await?; unioned_sendable_record_batch_streams.push(remote_sendable_record_batch_stream); } @@ -123,13 +146,22 @@ async fn execute_query_at_addresses_and_union( async fn execute_query_at_address( sql: String, address: String, + maybe_authorization: Option, ) -> Result>> { // Connect and execute query. let connection = Endpoint::new(address.clone())?.connect().await?; let mut flight_client = FlightServiceClient::new(connection); - let ticket = Ticket { ticket: sql.into() }; - let mut stream = flight_client.do_get(ticket).await?.into_inner(); + // Insert the authorization header if it is present. The header is added directly to avoid + // extracting the token from the authorization header. + let mut request = Request::new(Ticket { ticket: sql.into() }); + if let Some(authorization) = maybe_authorization { + request + .metadata_mut() + .insert("authorization", authorization); + } + + let mut stream = flight_client.do_get(request).await?.into_inner(); // Read schema of record batches. let flight_data = stream.message().await?.ok_or_else(|| { @@ -252,26 +284,6 @@ pub fn table_name_from_flight_descriptor( .ok_or_else(|| Status::invalid_argument("No table name in FlightDescriptor.path.")) } -/// Return `true` if the request contains the cluster key and `false` if not. If the request -/// contains a key that does not match the cluster key, return [`Status`]. -#[allow(clippy::result_large_err)] -fn cluster_key_in_request( - cluster: &Cluster, - request_metadata: &MetadataMap, -) -> StdResult { - if let Some(request_key) = request_metadata.get("x-cluster-key") { - if cluster.key() == request_key { - Ok(true) - } else { - Err(Status::invalid_argument( - "The cluster key in the request does not match the cluster key in the configuration.", - )) - } - } else { - Ok(false) - } -} - /// Return an empty stream of [`RecordBatches`](datafusion::arrow::record_batch::RecordBatch) that /// can be returned when a SQL command has been successfully executed but did not produce any rows /// to return. @@ -330,7 +342,7 @@ impl FlightServiceHandler { // If the cluster key is in the request, the request is from a peer node, which means the // table has already been created in the remote data folder and propagated to all nodes. if let ClusterMode::MultiNode(cluster) = configuration_manager.cluster_mode() - && !cluster_key_in_request(cluster, request_metadata)? + && request_metadata.get("x-cluster-key").is_none() { cluster .create_cluster_normal_table(table_name, schema) @@ -359,7 +371,7 @@ impl FlightServiceHandler { // If the cluster key is in the request, the request is from a peer node, which means the // table has already been created in the remote data folder and propagated to all nodes. if let ClusterMode::MultiNode(cluster) = configuration_manager.cluster_mode() - && !cluster_key_in_request(cluster, request_metadata)? + && request_metadata.get("x-cluster-key").is_none() { cluster .create_cluster_time_series_table(time_series_table_metadata) @@ -388,7 +400,7 @@ impl FlightServiceHandler { // If the cluster key is in the request, the request is from a peer node, which means the // tables have already been dropped in the remote data folder and propagated to all nodes. if let ClusterMode::MultiNode(cluster) = configuration_manager.cluster_mode() - && !cluster_key_in_request(cluster, request_metadata)? + && request_metadata.get("x-cluster-key").is_none() { cluster .drop_cluster_tables(table_names) @@ -719,9 +731,12 @@ impl FlightService for FlightServiceHandler { .await .map_err(error_to_status_internal)?; + let maybe_authorization = request.metadata().get("authorization").cloned(); + execute_query_at_addresses_and_union( &sql, addresses, + maybe_authorization, local_sendable_record_batch_stream, ) .await