blop

Author: fulmicoton

config/quickwit.yaml

@@ -150,3 +150,9 @@ indexer:
 
 jaeger:
   enable_endpoint: ${QW_ENABLE_JAEGER_ENDPOINT:-true}
+
+license: ${QW_LICENSE}
+
+# authorization:
+#     root_key: ${QW_ROOT_KEY}
+#     node_token: ${QW_NODE_TOKEN}

quickwit/Cargo.lock

@@ -240,6 +240,7 @@ tikv-jemalloc-ctl = "0.5"
 tikv-jemallocator = "0.5"
 time = { version = "0.3", features = ["std", "formatting", "macros"] }
 tokio = { version = "1.40", features = ["full"] }
+tokio-inherit-task-local = "0.2"
 tokio-metrics = { version = "0.3.1", features = ["rt"] }
 tokio-stream = { version = "0.1", features = ["sync"] }
 tokio-util = { version = "0.7", features = ["full"] }

quickwit/Cargo.toml

@@ -9,10 +9,12 @@ authors.workspace = true
 license.workspace = true
 
 [dependencies]
+anyhow = { workspace = true, optional = true }
 tower = { workspace = true}
 biscuit-auth = { workspace = true, optional=true }
 futures = { workspace = true }
 http = { workspace = true }
+tokio-inherit-task-local = { workspace = true }
 serde = { workspace = true }
 thiserror = { workspace = true }
 tonic = { workspace = true }
@@ -23,4 +25,4 @@ pin-project = { workspace = true }
 quickwit-common = { workspace = true }
 
 [features]
-enterprise = ["biscuit-auth"]
+enterprise = ["dep:biscuit-auth", "dep:anyhow"]

quickwit/quickwit-authorize/Cargo.toml

@@ -1,3 +1,22 @@
+// Copyright (C) 2024 Quickwit, Inc.
+//
+// Quickwit is offered under the AGPL v3.0 and as commercial software.
+// For commercial licensing, contact us at [email protected].
+//
+// AGPL:
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as
+// published by the Free Software Foundation, either version 3 of the
+// License, or (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see <http://www.gnu.org/licenses/>.
+
 use std::fmt;
 use std::task::{Context, Poll};
 
@@ -7,6 +26,7 @@ use tower::{Layer, Service};
 
 use crate::AuthorizationError;
 
+#[derive(Clone, Copy, Debug)]
 pub struct AuthorizationLayer;
 
 impl<S: Clone> Layer<S> for AuthorizationLayer {

quickwit/quickwit-authorize/src/community.rs renamed to quickwit/quickwit-authorize/src/community/mod.rs

@@ -0,0 +1,74 @@
+// Copyright (C) 2024 Quickwit, Inc.
+//
+// Quickwit is offered under the AGPL v3.0 and as commercial software.
+// For commercial licensing, contact us at [email protected].
+//
+// AGPL:
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as
+// published by the Free Software Foundation, either version 3 of the
+// License, or (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see <http://www.gnu.org/licenses/>.
+
+use std::task::{Context, Poll};
+
+use futures::future::Either;
+use http::Request;
+use tokio::task::futures::TaskLocalFuture;
+use tokio_inherit_task_local::TaskLocalInheritableTable;
+use tower::{Layer, Service};
+use tracing::debug;
+
+use super::AuthorizationToken;
+
+#[derive(Clone, Copy, Debug)]
+pub struct AuthorizationTokenExtractionLayer;
+
+impl<S: Clone> Layer<S> for AuthorizationTokenExtractionLayer {
+    type Service = AuthorizationTokenExtractionService<S>;
+
+    fn layer(&self, service: S) -> Self::Service {
+        AuthorizationTokenExtractionService { service }
+    }
+}
+
+#[derive(Clone)]
+pub struct AuthorizationTokenExtractionService<S> {
+    service: S,
+}
+
+fn get_authorization_token_opt(headers: &http::HeaderMap) -> Option<AuthorizationToken> {
+    let authorization_header_value = headers.get("Authorization")?;
+    let authorization_header_str = authorization_header_value.to_str().ok()?;
+    crate::get_auth_token_from_str(authorization_header_str).ok()
+}
+
+impl<B, S> Service<Request<B>> for AuthorizationTokenExtractionService<S>
+where S: Service<Request<B>>
+{
+    type Response = S::Response;
+    type Error = S::Error;
+    type Future = Either<S::Future, TaskLocalFuture<TaskLocalInheritableTable, S::Future>>;
+
+    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        self.service.poll_ready(cx)
+    }
+
+    fn call(&mut self, request: Request<B>) -> Self::Future {
+        let authorization_token_opt = get_authorization_token_opt(request.headers());
+        debug!(authorization_token_opt = ?authorization_token_opt, "Authorization token extracted");
+        let fut = self.service.call(request);
+        if let Some(authorization_token) = authorization_token_opt {
+            Either::Right(crate::execute_with_authorization(authorization_token, fut))
+        } else {
+            Either::Left(fut)
+        }
+    }
+}

quickwit/quickwit-authorize/src/authorization_layer.rs renamed to quickwit/quickwit-authorize/src/enterprise/authorization_layer.rs

@@ -19,15 +19,31 @@
 // components are licensed under the original license provided by the owner of the
 // applicable component.
 
+mod authorization_layer;
+mod authorization_token_extraction_layer;
+
 use std::future::Future;
 use std::str::FromStr;
 use std::sync::{Arc, OnceLock};
 
+use anyhow::Context;
+pub use authorization_layer::AuthorizationLayer;
+pub use authorization_token_extraction_layer::AuthorizationTokenExtractionLayer;
 use biscuit_auth::macros::authorizer;
 use biscuit_auth::{Authorizer, Biscuit, RootKeyProvider};
+use tokio::task::futures::TaskLocalFuture;
+use tokio_inherit_task_local::TaskLocalInheritableTable;
+use tracing::info;
 
 use crate::AuthorizationError;
 
+tokio_inherit_task_local::inheritable_task_local! {
+    pub static AUTHORIZATION_TOKEN: AuthorizationToken;
+}
+
+static ROOT_KEY_PROVIDER: OnceLock<Arc<dyn RootKeyProvider + Sync + Send>> = OnceLock::new();
+static NODE_TOKEN: OnceLock<Arc<AuthorizationToken>> = OnceLock::new();
+
 pub struct AuthorizationToken(Biscuit);
 
 impl AuthorizationToken {
@@ -54,7 +70,22 @@ impl std::fmt::Debug for AuthorizationToken {
     }
 }
 
-static ROOT_KEY_PROVIDER: OnceLock<Arc<dyn RootKeyProvider + Sync + Send>> = OnceLock::new();
+pub fn set_node_token_hex(node_token_hex: &str) -> anyhow::Result<()> {
+    let node_token =
+        AuthorizationToken::from_str(node_token_hex).context("failed to set node token")?;
+    if NODE_TOKEN.set(Arc::new(node_token)).is_err() {
+        tracing::error!("node token was already initialized");
+    }
+    Ok(())
+}
+
+pub fn set_root_public_key(root_key_hex: &str) -> anyhow::Result<()> {
+    let public_key = biscuit_auth::PublicKey::from_bytes_hex(root_key_hex)
+        .context("failed to parse root public key")?;
+    let key_provider: Arc<dyn RootKeyProvider + Sync + Send> = Arc::new(public_key);
+    set_root_key_provider(key_provider);
+    Ok(())
+}
 
 pub fn set_root_key_provider(key_provider: Arc<dyn RootKeyProvider + Sync + Send>) {
     if ROOT_KEY_PROVIDER.set(key_provider).is_err() {
@@ -79,10 +110,6 @@ impl FromStr for AuthorizationToken {
     }
 }
 
-tokio::task_local! {
-    pub static AUTHORIZATION_TOKEN: AuthorizationToken;
-}
-
 const AUTHORIZATION_VALUE_PREFIX: &str = "Bearer ";
 
 fn default_operation_authorizer<T: ?Sized>(
@@ -146,6 +173,16 @@ impl From<biscuit_auth::error::Token> for AuthorizationError {
     }
 }
 
+pub fn get_auth_token_from_str(
+    authorization_header_value: &str,
+) -> Result<AuthorizationToken, AuthorizationError> {
+    let authorization_token_str: &str = authorization_header_value
+        .strip_prefix(AUTHORIZATION_VALUE_PREFIX)
+        .ok_or(AuthorizationError::InvalidToken)?;
+    let biscuit: Biscuit = Biscuit::from_base64(authorization_token_str, get_root_key_provider())?;
+    Ok(AuthorizationToken(biscuit))
+}
+
 pub fn get_auth_token(
     req_metadata: &tonic::metadata::MetadataMap,
 ) -> Result<AuthorizationToken, AuthorizationError> {
@@ -154,11 +191,7 @@ pub fn get_auth_token(
         .ok_or(AuthorizationError::AuthorizationTokenMissing)?
         .to_str()
         .map_err(|_| AuthorizationError::InvalidToken)?;
-    let authorization_token_str: &str = authorization_header_value
-        .strip_prefix(AUTHORIZATION_VALUE_PREFIX)
-        .ok_or(AuthorizationError::InvalidToken)?;
-    let biscuit: Biscuit = Biscuit::from_base64(authorization_token_str, get_root_key_provider())?;
-    Ok(AuthorizationToken(biscuit))
+    get_auth_token_from_str(authorization_header_value)
 }
 
 pub fn set_auth_token(
@@ -216,15 +249,17 @@ pub fn authorize_stream<R: StreamAuthorization>(
 }
 
 pub fn authorize_request<R: Authorization>(req: &R) -> Result<(), AuthorizationError> {
-    AUTHORIZATION_TOKEN
+    let res = AUTHORIZATION_TOKEN
         .try_with(|auth_token| authorize(req, auth_token))
-        .unwrap_or(Err(AuthorizationError::AuthorizationTokenMissing))
+        .unwrap_or(Err(AuthorizationError::AuthorizationTokenMissing));
+    info!("request authorization");
+    res
 }
 
 pub fn execute_with_authorization<F, O>(
     token: AuthorizationToken,
     f: F,
-) -> impl Future<Output = O>
+) -> TaskLocalFuture<TaskLocalInheritableTable, F>
 where
     F: Future<Output = O>,
 {

quickwit/quickwit-authorize/src/enterprise/authorization_token_extraction_layer.rs

@@ -17,14 +17,12 @@
 // You should have received a copy of the GNU Affero General Public License
 // along with this program. If not, see <http://www.gnu.org/licenses/>.
 
-mod authorization_layer;
-
 #[cfg(not(feature = "enterprise"))]
-#[path = "community.rs"]
+#[path = "community/mod.rs"]
 mod implementation;
 
 #[cfg(feature = "enterprise")]
-#[path = "enterprise.rs"]
+#[path = "enterprise/mod.rs"]
 mod implementation;
 
 pub use implementation::*;

quickwit/quickwit-authorize/src/enterprise.rs renamed to quickwit/quickwit-authorize/src/enterprise/mod.rs

@@ -79,7 +79,7 @@ quickwit-metastore = { workspace = true, features = ["testsuite"] }
 quickwit-storage = { workspace = true, features = ["testsuite"] }
 
 [features]
-enterprise = ["quickwit-config/enterprise", "quickwit-ingest/enterprise", "quickwit-proto/enterprise"]
+enterprise = ["quickwit-config/enterprise", "quickwit-ingest/enterprise", "quickwit-proto/enterprise", "quickwit-serve/enterprise"]
 jemalloc = ["dep:tikv-jemalloc-ctl", "dep:tikv-jemallocator"]
 ci-test = []
 pprof = ["quickwit-serve/pprof"]