authorization as a layer

Author: fulmicoton

quickwit/Cargo.lock

@@ -9,13 +9,18 @@ authors.workspace = true
 license.workspace = true
 
 [dependencies]
+tower = { workspace = true}
 biscuit-auth = { workspace = true, optional=true }
+futures = { workspace = true }
 http = { workspace = true }
 serde = { workspace = true }
 thiserror = { workspace = true }
 tonic = { workspace = true }
 tokio = { workspace = true }
 tracing = { workspace = true }
+pin-project = { workspace = true }
+
+quickwit-common = { workspace = true }
 
 [features]
 enterprise = ["biscuit-auth"]

quickwit/quickwit-auth/Cargo.toml

@@ -0,0 +1,51 @@
+use std::fmt;
+use std::task::{Context, Poll};
+
+use futures::future::Either;
+use quickwit_common::tower::RpcName;
+use tower::{Layer, Service};
+
+use crate::AuthorizationError;
+
+pub struct AuthorizationLayer;
+
+impl<S: Clone> Layer<S> for AuthorizationLayer {
+    type Service = AuthorizationService<S>;
+
+    fn layer(&self, service: S) -> Self::Service {
+        AuthorizationService { service }
+    }
+}
+
+#[derive(Clone)]
+pub struct AuthorizationService<S> {
+    service: S,
+}
+
+impl<S, Request> Service<Request> for AuthorizationService<S>
+where
+    S: Service<Request>,
+    S::Future: Send + 'static,
+    S::Response: Send + 'static,
+    S::Error: From<AuthorizationError> + Send + 'static,
+    Request: fmt::Debug + Send + RpcName + crate::Authorization + 'static,
+{
+    type Response = S::Response;
+    type Error = S::Error;
+    type Future =
+        futures::future::Either<futures::future::Ready<Result<S::Response, S::Error>>, 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) -> Self::Future {
+        if let Err(authorization_err) = crate::authorize_request(&request) {
+            let err = S::Error::from(authorization_err);
+            let result: Result<S::Response, S::Error> =  Err(err);
+            return Either::Left(futures::future::ready(result));
+        }
+        let service_fut = self.service.call(request);
+        Either::Right(service_fut)
+    }
+}

quickwit/quickwit-auth/src/authorization_layer.rs

@@ -85,3 +85,8 @@ pub fn execute_with_authorization<F, O>(_: AuthorizationToken, f: F) -> impl Fut
 where F: Future<Output = O> {
     f
 }
+
+
+pub fn authorize_request<R: Authorization>(_req: &R) -> Result<(), AuthorizationError> {
+    Ok(())
+}

quickwit/quickwit-auth/src/community.rs

@@ -215,6 +215,14 @@ pub fn authorize_stream<R: StreamAuthorization>(
     Ok(())
 }
 
+pub fn authorize_request<R: Authorization>(req: &R) -> Result<(), AuthorizationError> {
+    AUTHORIZATION_TOKEN
+        .try_with(|auth_token| {
+            authorize(req, auth_token)
+        })
+        .unwrap_or(Err(AuthorizationError::AuthorizationTokenMissing))
+}
+
 pub fn execute_with_authorization<F, O>(
     token: AuthorizationToken,
     f: F,

quickwit/quickwit-auth/src/enterprise.rs

@@ -17,7 +17,7 @@
 // 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 serde::{Deserialize, Serialize};
+mod authorization_layer;
 
 #[cfg(not(feature = "enterprise"))]
 #[path = "community.rs"]
@@ -28,6 +28,7 @@ mod implementation;
 mod implementation;
 
 pub use implementation::*;
+use serde::{Deserialize, Serialize};
 
 #[derive(thiserror::Error, Debug, Clone, Copy, Serialize, Deserialize, Eq, PartialEq)]
 pub enum AuthorizationError {

quickwit/quickwit-auth/src/lib.rs

@@ -1271,24 +1271,12 @@ fn generate_grpc_server_adapter_methods(context: &CodegenContext) -> TokenStream
             quote! { tonic::Response::new }
         };
 
-        let authorize_block = if syn_method.client_streaming {
-            let stream_item = &syn_method.request_type;
-            quote! {
-                quickwit_auth::authorize_stream::<#stream_item>(&auth_token)?;
-            }
-        } else {
-            quote! {
-                quickwit_auth::authorize(&req, &auth_token)?;
-            }
-        };
         let method = quote! {
             #associated_type
 
             async fn #method_name(&self, request: tonic::Request<#request_type>) -> Result<tonic::Response<#response_type>, tonic::Status> {
                 let auth_token = quickwit_auth::get_auth_token(request.metadata())?;
-                let req = #method_arg;
-                #authorize_block;
-                quickwit_auth::execute_with_authorization(auth_token, self.inner.0.#method_name(req)).await
+                quickwit_auth::execute_with_authorization(auth_token, self.inner.0.#method_name(#method_arg)).await
                     .map(#into_response_type)
                     .map_err(crate::error::grpc_error_to_grpc_status)
             }