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

@@ -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;

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

@@ -1167,14 +1167,29 @@ fn generate_grpc_client_adapter_methods(context: &CodegenContext) -> TokenStream
         } else {
             quote! { |response| response.into_inner() }
         };
-        let method = quote! {
-            async fn #method_name(&self, request: #request_type) -> #result_type<#response_type> {
-                self.inner
+        let method = if syn_method.client_streaming {
+            quote! {
+                async fn #method_name(&self, request: #request_type) -> #result_type<#response_type> {
+                    let tonic_request = quickwit_auth::build_tonic_stream_request_with_auth_token(request)?;
+                    self.inner
                     .clone()
-                    .#method_name(request)
+                    .#method_name(tonic_request)
                     .await
                     .map(#into_response_type)
                     .map_err(|status| crate::error::grpc_status_to_service_error(status, #rpc_name))
+                }
+            }
+        } else {
+            quote! {
+                async fn #method_name(&self, request: #request_type) -> #result_type<#response_type> {
+                    let tonic_request = quickwit_auth::build_tonic_request_with_auth_token(request)?;
+                    self.inner
+                        .clone()
+                        .#method_name(tonic_request)
+                        .await
+                        .map(#into_response_type)
+                        .map_err(|status| crate::error::grpc_status_to_service_error(status, #rpc_name))
+                }
             }
         };
         stream.extend(method);
@@ -1286,8 +1301,6 @@ fn generate_grpc_server_adapter_methods(context: &CodegenContext) -> TokenStream
     stream
 }
 
-
-
 /// A [`ServiceGenerator`] wrapper that appends a suffix to the name of the wrapped service. It is
 /// used to add a `Grpc` suffix to the service, client, and server generated by tonic.
 struct WithSuffixServiceGenerator {

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

@@ -1,25 +1,47 @@
-use quickwit_auth::Authorization;
-use quickwit_auth::AuthorizationError;
-use quickwit_auth::AuthorizationToken;
+// 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::FetchRequest;
-use crate::IngestRequest;
-use crate::TailRequest;
+use quickwit_auth::{Authorization, AuthorizationError, AuthorizationToken};
+
+use crate::{FetchRequest, IngestRequest, TailRequest};
 
 impl Authorization for TailRequest {
-    fn attenuate(&self, auth_token: AuthorizationToken) -> Result<AuthorizationToken, AuthorizationError> {
+    fn attenuate(
+        &self,
+        auth_token: AuthorizationToken,
+    ) -> Result<AuthorizationToken, AuthorizationError> {
         Ok(auth_token)
     }
 }
 
 impl Authorization for IngestRequest {
-    fn attenuate(&self, auth_token: AuthorizationToken) -> Result<AuthorizationToken, AuthorizationError> {
+    fn attenuate(
+        &self,
+        auth_token: AuthorizationToken,
+    ) -> Result<AuthorizationToken, AuthorizationError> {
         Ok(auth_token)
     }
 }
 
 impl Authorization for FetchRequest {
-    fn attenuate(&self, auth_token: AuthorizationToken) -> Result<AuthorizationToken, AuthorizationError> {
+    fn attenuate(
+        &self,
+        auth_token: AuthorizationToken,
+    ) -> Result<AuthorizationToken, AuthorizationError> {
         Ok(auth_token)
     }
 }

quickwit/quickwit-codegen/src/codegen.rs

@@ -21,6 +21,7 @@ use std::io;
 
 use mrecordlog::error::*;
 use quickwit_actors::AskError;
+use quickwit_auth::AuthorizationError;
 use quickwit_common::rate_limited_error;
 use quickwit_common::tower::BufferError;
 pub(crate) use quickwit_proto::error::{grpc_error_to_grpc_status, grpc_status_to_service_error};
@@ -48,6 +49,8 @@ pub enum IngestServiceError {
     RateLimited(RateLimitingCause),
     #[error("ingest service is unavailable ({0})")]
     Unavailable(String),
+    #[error("unauthorized: {0}")]
+    Unauthorized(#[from] AuthorizationError),
 }
 
 impl From<AskError<IngestServiceError>> for IngestServiceError {
@@ -93,6 +96,9 @@ impl From<IngestV2Error> for IngestServiceError {
             IngestV2Error::TooManyRequests(rate_limiting_cause) => {
                 IngestServiceError::RateLimited(rate_limiting_cause)
             }
+            IngestV2Error::Unauthorized(authorization_error) => {
+                IngestServiceError::Unauthorized(authorization_error)
+            }
         }
     }
 }
@@ -134,6 +140,9 @@ impl From<IngestFailure> for IngestServiceError {
             IngestFailureReason::CircuitBreaker => {
                 IngestServiceError::RateLimited(RateLimitingCause::CircuitBreaker)
             }
+            IngestFailureReason::Unauthorized => {
+                IngestServiceError::Unauthorized(AuthorizationError::PermissionDenied)
+            }
         }
     }
 }
@@ -161,6 +170,7 @@ impl ServiceError for IngestServiceError {
             }
             Self::RateLimited(_) => ServiceErrorCode::TooManyRequests,
             Self::Unavailable(_) => ServiceErrorCode::Unavailable,
+            Self::Unauthorized(_) => ServiceErrorCode::Unauthorized,
         }
     }
 }
@@ -204,6 +214,9 @@ impl From<IngestServiceError> for tonic::Status {
             IngestServiceError::IoError { .. } => tonic::Code::Internal,
             IngestServiceError::RateLimited(_) => tonic::Code::ResourceExhausted,
             IngestServiceError::Unavailable(_) => tonic::Code::Unavailable,
+            IngestServiceError::Unauthorized(authorized_error) => {
+                return (*authorized_error).into();
+            }
         };
         let message = error.to_string();
         tonic::Status::new(code, message)

quickwit/quickwit-ingest/src/authorize.rs

@@ -44,6 +44,7 @@ pub(crate) struct IngestResultMetrics {
     pub load_shedding: IntCounter,
     pub shard_not_found: IntCounter,
     pub unavailable: IntCounter,
+    pub unauthorized: IntCounter,
 }
 
 impl Default for IngestResultMetrics {
@@ -72,6 +73,7 @@ impl Default for IngestResultMetrics {
             load_shedding: ingest_result_total_vec.with_label_values(["load_shedding"]),
             unavailable: ingest_result_total_vec.with_label_values(["unavailable"]),
             shard_not_found: ingest_result_total_vec.with_label_values(["shard_not_found"]),
+            unauthorized: ingest_result_total_vec.with_label_values(["unauthorized"]),
         }
     }
 }

quickwit/quickwit-ingest/src/codegen/ingest_service.rs

@@ -542,6 +542,7 @@ fn update_ingest_metrics(ingest_result: &IngestV2Result<IngestResponseV2>, num_s
                         ingest_results_metrics.router_load_shedding.inc()
                     }
                     IngestFailureReason::LoadShedding => ingest_results_metrics.load_shedding.inc(),
+                    IngestFailureReason::Unauthorized => ingest_results_metrics.unauthorized.inc(),
                 }
             }
         }
@@ -588,6 +589,9 @@ fn update_ingest_metrics(ingest_result: &IngestV2Result<IngestResponseV2>, num_s
             IngestV2Error::Internal(_) => {
                 ingest_results_metrics.internal.inc_by(num_subrequests);
             }
+            IngestV2Error::Unauthorized(_) => {
+                ingest_results_metrics.unauthorized.inc_by(num_subrequests);
+            }
         },
     }
 }

quickwit/quickwit-ingest/src/error.rs

@@ -224,6 +224,12 @@ impl IngestWorkbench {
                     self.record_too_many_requests(subrequest_id, rate_limiting_cause);
                 }
             }
+            IngestV2Error::Unauthorized(_) => {
+                for subrequest_id in persist_summary.subrequest_ids {
+                    let failure = SubworkbenchFailure::Persist(PersistFailureReason::Unauthorized);
+                    self.record_failure(subrequest_id, failure);
+                }
+            }
         }
     }