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

@@ -44,7 +44,7 @@ pub trait StreamAuthorization {
 
 impl<T> StreamAuthorization for T {}
 
-pub fn get_auth_token(
+pub fn extract_auth_token(
     _req_metadata: &tonic::metadata::MetadataMap,
 ) -> Result<AuthorizationToken, AuthorizationError> {
     Ok(())
@@ -63,12 +63,6 @@ pub fn authorize<R: Authorization>(
     Ok(())
 }
 
-pub fn build_tonic_stream_request_with_auth_token<R>(
-    req: R,
-) -> Result<tonic::Request<R>, AuthorizationError> {
-    Ok(tonic::Request::new(req))
-}
-
 pub fn build_tonic_request_with_auth_token<R: Authorization>(
     req: R,
 ) -> Result<tonic::Request<R>, AuthorizationError> {

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

@@ -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/authorization_layer.rs renamed to quickwit/quickwit-authorize/src/enterprise/authorization_layer.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)
+        }
+    }
+}