added check

Author: fulmicoton

quickwit/Cargo.lock

@@ -36,7 +36,6 @@ members = [
   "quickwit-serve",
   "quickwit-storage",
   "quickwit-telemetry",
-  "quickwit-telemetry",
 ]
 
 # The following list excludes `quickwit-metastore-utils` and `quickwit-lambda`

quickwit/Cargo.toml

@@ -0,0 +1,18 @@
+[package]
+name = "quickwit-auth"
+version.workspace = true
+edition.workspace = true
+homepage.workspace = true
+documentation.workspace = true
+repository.workspace = true
+authors.workspace = true
+license.workspace = true
+
+[dependencies]
+biscuit-auth = { workspace = true }
+http = { workspace = true }
+serde = { workspace = true }
+thiserror = { workspace = true }
+tonic = { workspace = true }
+tokio = { workspace = true }
+tracing = { workspace = true }

quickwit/quickwit-auth/Cargo.toml

@@ -0,0 +1,216 @@
+use std::sync::{Arc, OnceLock};
+
+use biscuit_auth::macros::authorizer;
+use biscuit_auth::{Authorizer, Biscuit, RootKeyProvider};
+use serde::{Deserialize, Serialize};
+
+pub type AuthorizationToken = Biscuit;
+
+tokio::task_local! {
+    pub static AUTHORIZATION_TOKEN: AuthorizationToken;
+}
+
+#[derive(thiserror::Error, Debug, Clone, Copy, Serialize, Deserialize, Eq, PartialEq)]
+pub enum AuthorizationError {
+    #[error("authorization token missing")]
+    AuthorizationTokenMissing,
+    #[error("invalid token")]
+    InvalidToken,
+    #[error("permission denied")]
+    PermissionDenied,
+}
+
+const AUTHORIZATION_VALUE_PREFIX: &str = "Bearer ";
+
+fn default_operation_authorizer<T: ?Sized>(
+    auth_token: &AuthorizationToken,
+) -> Result<Authorizer, AuthorizationError> {
+    let request_type = std::any::type_name::<T>();
+    let operation: &str = request_type.strip_suffix("Request").unwrap();
+    let mut authorizer: Authorizer = authorizer!(
+        r#"
+        operation({operation});
+
+        // We generate the actual user role, by doing an union of the rights granted via roles.
+        user_right($operation) <- role($role), right($role, $operation);
+        user_right($operation, $resource) <- role($role), right($role, $operation, $resource);
+        user_right($operation) <- role("root"), operation($operation);
+        user_right($operation, $resource) <- role("root"), operation($operation), resource($resource);
+
+        // Finally we check that we have access to index1 and index2.
+        check all operation($operation), right($operation);
+
+        allow if true;
+    "#
+    );
+    authorizer.set_time();
+    authorizer.add_token(auth_token)?;
+    Ok(authorizer)
+}
+
+pub trait Authorization {
+    fn attenuate(
+        &self,
+        auth_token: AuthorizationToken,
+    ) -> Result<AuthorizationToken, AuthorizationError>;
+    fn authorizer(
+        &self,
+        auth_token: &AuthorizationToken,
+    ) -> Result<Authorizer, AuthorizationError> {
+        default_operation_authorizer::<Self>(auth_token)
+    }
+}
+
+pub trait StreamAuthorization {
+    fn attenuate(
+        auth_token: AuthorizationToken,
+    ) -> std::result::Result<AuthorizationToken, AuthorizationError> {
+        Ok(auth_token)
+    }
+    fn authorizer(
+        auth_token: &AuthorizationToken,
+    ) -> std::result::Result<Authorizer, AuthorizationError> {
+        default_operation_authorizer::<Self>(&auth_token)
+    }
+}
+
+static ROOT_KEY_PROVIDER: OnceLock<Arc<dyn RootKeyProvider + Sync + Send>> = OnceLock::new();
+
+pub fn set_root_key_provider(key_provider: Arc<dyn RootKeyProvider + Sync + Send>) {
+    if ROOT_KEY_PROVIDER.set(key_provider).is_err() {
+        tracing::error!("root key provider was already initialized");
+    }
+}
+
+fn get_root_key_provider() -> Arc<dyn RootKeyProvider> {
+    ROOT_KEY_PROVIDER
+        .get()
+        .expect("root key provider should have been initialized beforehand")
+        .clone()
+}
+
+impl From<biscuit_auth::error::Token> for AuthorizationError {
+    fn from(_token_error: biscuit_auth::error::Token) -> AuthorizationError {
+        AuthorizationError::InvalidToken
+    }
+}
+
+pub fn get_auth_token(
+    req_metadata: &tonic::metadata::MetadataMap,
+) -> Result<AuthorizationToken, AuthorizationError> {
+    let authorization_header_value: &str = req_metadata
+        .get(http::header::AUTHORIZATION.as_str())
+        .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(biscuit)
+}
+
+pub fn set_auth_token(
+    auth_token: &AuthorizationToken,
+    req_metadata: &mut tonic::metadata::MetadataMap,
+) {
+    let authorization_header_value = format!("{AUTHORIZATION_VALUE_PREFIX}{auth_token}");
+    req_metadata.insert(
+        http::header::AUTHORIZATION.as_str(),
+        authorization_header_value.parse().unwrap(),
+    );
+}
+
+pub fn authorize<R: Authorization>(
+    req: &R,
+    auth_token: &Biscuit,
+) -> Result<(), AuthorizationError> {
+    let mut authorizer = req.authorizer(auth_token)?;
+    authorizer.add_token(&auth_token)?;
+    authorizer.authorize()?;
+    Ok(())
+}
+
+pub fn build_tonic_stream_request_with_auth_token<R>(
+    req: R,
+) -> Result<tonic::Request<R>, AuthorizationError> {
+    AUTHORIZATION_TOKEN
+        .try_with(|token| {
+            let mut request = tonic::Request::new(req);
+            set_auth_token(token, request.metadata_mut());
+            Ok(request)
+        })
+        .unwrap_or(Err(AuthorizationError::AuthorizationTokenMissing))
+}
+
+pub fn build_tonic_request_with_auth_token<R: Authorization>(
+    req: R,
+) -> Result<tonic::Request<R>, AuthorizationError> {
+    AUTHORIZATION_TOKEN
+        .try_with(|token| {
+            let mut request = tonic::Request::new(req);
+            set_auth_token(token, request.metadata_mut());
+            Ok(request)
+        })
+        .unwrap_or(Err(AuthorizationError::AuthorizationTokenMissing))
+}
+
+pub fn authorize_stream<R: StreamAuthorization>(
+    auth_token: &Biscuit,
+) -> Result<(), AuthorizationError> {
+    let mut authorizer = R::authorizer(auth_token)?;
+    authorizer.add_token(&auth_token)?;
+    authorizer.authorize()?;
+    Ok(())
+}
+
+impl From<AuthorizationError> for tonic::Status {
+    fn from(authorization_error: AuthorizationError) -> tonic::Status {
+        match authorization_error {
+            AuthorizationError::AuthorizationTokenMissing => {
+                tonic::Status::unauthenticated("Authorization token missing")
+            }
+            AuthorizationError::InvalidToken => {
+                tonic::Status::unauthenticated("Invalid authorization token")
+            }
+            AuthorizationError::PermissionDenied => {
+                tonic::Status::permission_denied("Permission denied")
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_auth_token() {
+        let mut req_metadata = tonic::metadata::MetadataMap::new();
+        let auth_token = "test_token".to_string();
+        set_auth_token(&auth_token, &mut req_metadata);
+        let auth_token_retrieved = get_auth_token(&req_metadata).unwrap();
+        assert_eq!(auth_token_retrieved, auth_token);
+    }
+
+    #[test]
+    fn test_auth_token_missing() {
+        let req_metadata = tonic::metadata::MetadataMap::new();
+        let missing_error = get_auth_token(&req_metadata).unwrap_err();
+        assert!(matches!(
+            missing_error,
+            AuthorizationError::AuthorizationTokenMissing
+        ));
+    }
+
+    #[test]
+    fn test_auth_token_invalid() {
+        let mut req_metadata = tonic::metadata::MetadataMap::new();
+        req_metadata.insert(
+            http::header::AUTHORIZATION.as_str(),
+            "some_token".parse().unwrap(),
+        );
+        let missing_error = get_auth_token(&req_metadata).unwrap_err();
+        assert!(matches!(missing_error, AuthorizationError::InvalidToken));
+    }
+}

quickwit/quickwit-auth/src/lib.rs

@@ -1,26 +1,46 @@
-use quickwit_auth::Authorization;
-use quickwit_auth::AuthorizationError;
-use quickwit_auth::AuthorizationToken;
-use quickwit_auth::StreamAuthorization;
+// The Quickwit Enterprise Edition (EE) license
+// Copyright (c) 2024-present Quickwit Inc.
+//
+// With regard to the Quickwit Software:
+//
+// This software and associated documentation files (the "Software") may only be
+// used in production, if you (and any entity that you represent) hold a valid
+// Quickwit Enterprise license corresponding to your usage.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+// SOFTWARE.
 
-use crate::GoodbyeRequest;
-use crate::HelloRequest;
-use crate::PingRequest;
+use quickwit_auth::{Authorization, AuthorizationError, AuthorizationToken, StreamAuthorization};
+
+use crate::{GoodbyeRequest, HelloRequest, PingRequest};
 
 impl Authorization for HelloRequest {
-    fn attenuate(&self, auth_token: quickwit_auth::AuthorizationToken) -> Result<quickwit_auth::AuthorizationToken, AuthorizationError> {
+    fn attenuate(
+        &self,
+        auth_token: quickwit_auth::AuthorizationToken,
+    ) -> Result<quickwit_auth::AuthorizationToken, AuthorizationError> {
         Ok(auth_token)
     }
 }
 
 impl Authorization for GoodbyeRequest {
-    fn attenuate(&self, auth_token: quickwit_auth::AuthorizationToken) -> Result<AuthorizationToken, AuthorizationError> {
+    fn attenuate(
+        &self,
+        auth_token: quickwit_auth::AuthorizationToken,
+    ) -> Result<AuthorizationToken, AuthorizationError> {
         Ok(auth_token)
     }
 }
 
 impl StreamAuthorization for PingRequest {
-    fn attenuate(auth_token: quickwit_auth::AuthorizationToken) -> Result<AuthorizationToken, AuthorizationError> {
+    fn attenuate(
+        auth_token: quickwit_auth::AuthorizationToken,
+    ) -> Result<AuthorizationToken, AuthorizationError> {
         Ok(auth_token)
     }
 }

quickwit/quickwit-codegen/example/src/authorization.rs

@@ -20,6 +20,7 @@
 use std::fmt;
 
 use quickwit_actors::AskError;
+use quickwit_auth::AuthorizationError;
 use quickwit_proto::error::GrpcServiceError;
 pub use quickwit_proto::error::{grpc_error_to_grpc_status, grpc_status_to_service_error};
 use quickwit_proto::{ServiceError, ServiceErrorCode};
@@ -38,6 +39,8 @@ pub enum HelloError {
     TooManyRequests,
     #[error("service unavailable: {0}")]
     Unavailable(String),
+    #[error("unauthorized: {0}")]
+    Unauthorized(#[from] AuthorizationError),
 }
 
 impl ServiceError for HelloError {
@@ -48,6 +51,7 @@ impl ServiceError for HelloError {
             Self::Timeout(_) => ServiceErrorCode::Timeout,
             Self::TooManyRequests => ServiceErrorCode::TooManyRequests,
             Self::Unavailable(_) => ServiceErrorCode::Unavailable,
+            Self::Unauthorized(_) => ServiceErrorCode::Unauthorized,
         }
     }
 }

quickwit/quickwit-codegen/example/src/codegen/hello.rs

@@ -19,9 +19,9 @@
 
 mod error;
 
+mod authorization;
 #[path = "codegen/hello.rs"]
 mod hello;
-mod authorization;
 
 use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::Arc;