Skip to content

Commit

Permalink
blop
Browse files Browse the repository at this point in the history
  • Loading branch information
fulmicoton committed Nov 5, 2024
1 parent 94b575a commit 2ab75ef
Show file tree
Hide file tree
Showing 23 changed files with 304 additions and 45 deletions.
6 changes: 6 additions & 0 deletions config/quickwit.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -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}
36 changes: 36 additions & 0 deletions quickwit/Cargo.lock

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

1 change: 1 addition & 0 deletions quickwit/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"] }
Expand Down
4 changes: 3 additions & 1 deletion quickwit/quickwit-authorize/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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 }
Expand All @@ -23,4 +25,4 @@ pin-project = { workspace = true }
quickwit-common = { workspace = true }

[features]
enterprise = ["biscuit-auth"]
enterprise = ["dep:biscuit-auth", "dep:anyhow"]
Original file line number Diff line number Diff line change
@@ -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};

Expand All @@ -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 {
Expand Down
Original file line number Diff line number Diff line change
@@ -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)
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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() {
Expand All @@ -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>(
Expand Down Expand Up @@ -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> {
Expand All @@ -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(
Expand Down Expand Up @@ -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>,
{
Expand Down
6 changes: 2 additions & 4 deletions quickwit/quickwit-authorize/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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::*;
Expand Down
2 changes: 1 addition & 1 deletion quickwit/quickwit-cli/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"]
Expand Down
Loading

0 comments on commit 2ab75ef

Please sign in to comment.