Skip to content

Commit

Permalink
added check
Browse files Browse the repository at this point in the history
  • Loading branch information
fulmicoton committed Nov 1, 2024
1 parent 5a924f6 commit 1945ea4
Show file tree
Hide file tree
Showing 38 changed files with 475 additions and 190 deletions.
1 change: 1 addition & 0 deletions quickwit/Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 0 additions & 1 deletion quickwit/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,6 @@ members = [
"quickwit-serve",
"quickwit-storage",
"quickwit-telemetry",
"quickwit-telemetry",
]

# The following list excludes `quickwit-metastore-utils` and `quickwit-lambda`
Expand Down
40 changes: 30 additions & 10 deletions quickwit/quickwit-codegen/example/src/authorization.rs
Original file line number Diff line number Diff line change
@@ -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)
}
}
11 changes: 8 additions & 3 deletions quickwit/quickwit-codegen/example/src/codegen/hello.rs

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

4 changes: 4 additions & 0 deletions quickwit/quickwit-codegen/example/src/error.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand All @@ -38,6 +39,8 @@ pub enum HelloError {
TooManyRequests,
#[error("service unavailable: {0}")]
Unavailable(String),
#[error("unauthorized: {0}")]
Unauthorized(#[from] AuthorizationError),
}

impl ServiceError for HelloError {
Expand All @@ -48,6 +51,7 @@ impl ServiceError for HelloError {
Self::Timeout(_) => ServiceErrorCode::Timeout,
Self::TooManyRequests => ServiceErrorCode::TooManyRequests,
Self::Unavailable(_) => ServiceErrorCode::Unavailable,
Self::Unauthorized(_) => ServiceErrorCode::Unauthorized,
}
}
}
Expand Down
2 changes: 1 addition & 1 deletion quickwit/quickwit-codegen/example/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
25 changes: 19 additions & 6 deletions quickwit/quickwit-codegen/src/codegen.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -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 {
Expand Down
40 changes: 31 additions & 9 deletions quickwit/quickwit-ingest/src/authorize.rs
Original file line number Diff line number Diff line change
@@ -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)
}
}
9 changes: 6 additions & 3 deletions quickwit/quickwit-ingest/src/codegen/ingest_service.rs

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

13 changes: 13 additions & 0 deletions quickwit/quickwit-ingest/src/error.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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)
}
}
}
}
Expand Down Expand Up @@ -134,6 +140,9 @@ impl From<IngestFailure> for IngestServiceError {
IngestFailureReason::CircuitBreaker => {
IngestServiceError::RateLimited(RateLimitingCause::CircuitBreaker)
}
IngestFailureReason::Unauthorized => {
IngestServiceError::Unauthorized(AuthorizationError::PermissionDenied)
}
}
}
}
Expand Down Expand Up @@ -161,6 +170,7 @@ impl ServiceError for IngestServiceError {
}
Self::RateLimited(_) => ServiceErrorCode::TooManyRequests,
Self::Unavailable(_) => ServiceErrorCode::Unavailable,
Self::Unauthorized(_) => ServiceErrorCode::Unauthorized,
}
}
}
Expand Down Expand Up @@ -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)
Expand Down
2 changes: 2 additions & 0 deletions quickwit/quickwit-ingest/src/ingest_v2/metrics.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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"]),
}
}
}
Expand Down
4 changes: 4 additions & 0 deletions quickwit/quickwit-ingest/src/ingest_v2/router.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
}
}
}
Expand Down Expand Up @@ -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);
}
},
}
}
Expand Down
6 changes: 6 additions & 0 deletions quickwit/quickwit-ingest/src/ingest_v2/workbench.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
}
}
}

Expand Down
Loading

0 comments on commit 1945ea4

Please sign in to comment.