feat: commit

Author: ovr

packages/cubejs-backend-native/Cargo.lock

@@ -29,6 +29,7 @@ findshlibs = "0.10.2"
 futures = "0.3.30"
 http-body-util = "0.1"
 axum = { version = "0.7.5", features = ["default", "ws"] }
+tower = "0.5.2"
 libc = "0.2"
 log = "0.4.21"
 log-reroute = "0.1"

packages/cubejs-backend-native/Cargo.toml

@@ -145,7 +145,7 @@ impl GatewayAuthService for NodeBridgeAuthService {
         req: GatewayCheckAuthRequest,
         token: String,
     ) -> Result<GatewayAuthenticateResponse, CubeError> {
-        trace!("[sql auth] Request ->");
+        trace!("[auth] Request ->");
 
         let request_id = Uuid::new_v4().to_string();
 
@@ -163,7 +163,7 @@ impl GatewayAuthService for NodeBridgeAuthService {
             Some(extra),
         )
         .await?;
-        trace!("[sql auth] Request <- {:?}", response);
+        trace!("[auth] Request <- {:?}", response);
 
         Ok(GatewayAuthenticateResponse {
             context: Arc::new(NativeAuthContext {

packages/cubejs-backend-native/src/auth.rs

@@ -0,0 +1,42 @@
+use std::sync::Arc;
+use axum::extract::State;
+use axum::http::{header, HeaderMap, HeaderValue, Response, StatusCode};
+use axum::response::IntoResponse;
+use crate::gateway::{ApiGatewayState, GatewayAuthContextRef, GatewayAuthService, GatewayCheckAuthRequest};
+use crate::gateway::http_error::HttpError;
+
+#[derive(Debug, Clone)]
+pub struct AuthExtension {
+    auth_context: GatewayAuthContextRef,
+}
+
+pub async fn gateway_auth_middleware(
+    State(state): State<ApiGatewayState>,
+    mut req: axum::extract::Request,
+    next: axum::middleware::Next,
+) -> Result<impl IntoResponse, HttpError> {
+    let Some(token_header_value) = req.headers().get("authorization") else {
+        return Err(HttpError::unauthorized("No authorization header".to_string()));
+    };
+
+    let auth = state
+        .injector_ref()
+        .get_service_typed::<dyn GatewayAuthService>()
+        .await;
+
+    let auth_fut = auth.authenticate(
+        GatewayCheckAuthRequest {
+            protocol: "http".to_string(),
+        },
+        token_header_value.to_str()?.to_string()
+    );
+
+    let auth_response = auth_fut.await.map_err(|err| {;
+        HttpError::unauthorized("Authentication error".to_string())
+    })?;
+
+    req.extensions_mut().insert(AuthExtension { auth_context: auth_response.context });
+
+    let response = next.run(req).await;
+    Ok(response)
+}

packages/cubejs-backend-native/src/gateway/auth_middleware.rs

@@ -19,7 +19,7 @@ pub struct GatewayAuthenticateResponse {
 
 #[derive(Debug, Serialize)]
 pub struct GatewayCheckAuthRequest {
-    protocol: String,
+    pub(crate) protocol: String,
 }
 
 #[async_trait]

packages/cubejs-backend-native/src/gateway/auth_service.rs

@@ -1,16 +1,18 @@
-use crate::gateway::ApiGatewayState;
+use crate::gateway::{ApiGatewayRouterBuilder, ApiGatewayState};
 use axum::extract::State;
 use axum::http::StatusCode;
-use axum::Json;
+use axum::{Extension, Json};
 use serde::Serialize;
+use crate::gateway::auth_middleware::AuthExtension;
 
 #[derive(Serialize)]
 pub struct HandlerResponse {
     message: String,
 }
 
 pub async fn stream_handler_v2(
-    State(_state): State<ApiGatewayState>,
+    State(_gateway_state): State<ApiGatewayState>,
+    Extension(_auth): Extension<AuthExtension>,
 ) -> (StatusCode, Json<HandlerResponse>) {
     (
         StatusCode::NOT_IMPLEMENTED,

packages/cubejs-backend-native/src/gateway/handlers/stream.rs

@@ -0,0 +1,48 @@
+use axum::http::header::ToStrError;
+use axum::response::{IntoResponse, Response};
+
+pub enum HttpErrorCode {
+    NotFound,
+    Unauthorized,
+    InternalServerError,
+}
+
+pub struct HttpError {
+    code: HttpErrorCode,
+    message: String,
+}
+
+impl HttpError {
+    pub fn unauthorized(message: String) -> HttpError {
+        Self {
+            code: HttpErrorCode::Unauthorized,
+            message,
+        }
+    }
+
+    pub fn status_code(&self) -> axum::http::StatusCode {
+        match self.code {
+            HttpErrorCode::NotFound => axum::http::StatusCode::NOT_FOUND,
+            HttpErrorCode::InternalServerError => axum::http::StatusCode::INTERNAL_SERVER_ERROR,
+            HttpErrorCode::Unauthorized => axum::http::StatusCode::UNAUTHORIZED,
+        }
+    }
+}
+
+impl IntoResponse for HttpError {
+    fn into_response(self) -> Response {
+        let status_code = self.status_code();
+
+        (status_code, status_code.canonical_reason().unwrap_or_default())
+            .into_response()
+    }
+}
+
+impl From<ToStrError> for HttpError {
+    fn from(value: ToStrError) -> Self {
+        HttpError {
+            code: HttpErrorCode::InternalServerError,
+            message: value.to_string(),
+        }
+    }
+}

packages/cubejs-backend-native/src/gateway/http_error.rs

@@ -3,6 +3,8 @@ pub mod handlers;
 pub mod router;
 pub mod server;
 pub mod state;
+pub mod auth_middleware;
+pub mod http_error;
 
 pub use auth_service::{
     GatewayAuthContext, GatewayAuthContextRef, GatewayAuthService, GatewayAuthenticateResponse,
@@ -11,3 +13,4 @@ pub use auth_service::{
 pub use router::ApiGatewayRouterBuilder;
 pub use server::{ApiGatewayServer, ApiGatewayServerImpl};
 pub use state::ApiGatewayState;
+pub use auth_middleware::{gateway_auth_middleware};

packages/cubejs-backend-native/src/gateway/mod.rs

@@ -2,6 +2,7 @@ use crate::gateway::handlers::stream_handler_v2;
 use crate::gateway::ApiGatewayState;
 use axum::routing::{get, MethodRouter};
 use axum::Router;
+use tower::ServiceBuilder;
 
 #[derive(Debug, Clone)]
 pub struct ApiGatewayRouterBuilder {

packages/cubejs-backend-native/src/gateway/router.rs

@@ -6,6 +6,7 @@ use cubesql::CubeError;
 use std::sync::Arc;
 use tokio::net::TcpListener;
 use tokio::sync::{watch, Mutex};
+use tower::ServiceBuilder;
 
 pub trait ApiGatewayServer: ProcessingLoop {}
 
@@ -79,6 +80,12 @@ impl ProcessingLoop for ApiGatewayServerImpl {
             log::trace!("Shutdown signal received");
         };
 
+        let router = router.layer(
+            ServiceBuilder::new()
+                .layer(crate::gateway::auth_middleware::gateway_auth_middleware)
+                .into_inner(),
+        );
+
         axum::serve(listener, router)
             .with_graceful_shutdown(shutdown_signal())
             .await