feat: introduce "protocol" and "method" props for request param in checkSqlAuth (#9525)

Author: MikeNitsenko

packages/cubejs-api-gateway/src/sql-server.ts

@@ -31,6 +31,11 @@ export type SQLServerConstructorOptions = {
   gatewayPort?: number,
 };
 
+export type SqlAuthServiceAuthenticateRequest = {
+  protocol: string;
+  method: string;
+};
+
 export class SQLServer {
   protected sqlInterfaceInstance: SqlInterfaceInstance | null = null;
 
@@ -88,10 +93,14 @@ export class SQLServer {
       let { securityContext } = session;
 
       if (request.meta.changeUser && request.meta.changeUser !== session.user) {
+        const sqlAuthRequest: SqlAuthServiceAuthenticateRequest = {
+          protocol: request.meta.protocol,
+          method: 'password',
+        };
         const canSwitch = session.superuser || await canSwitchSqlUser(session.user, request.meta.changeUser);
         if (canSwitch) {
           userForContext = request.meta.changeUser;
-          const current = await checkSqlAuth(request, userForContext, null);
+          const current = await checkSqlAuth({ ...request, ...sqlAuthRequest }, userForContext, null);
           securityContext = current.securityContext;
         } else {
           throw new Error(

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

@@ -1,7 +1,7 @@
 use async_trait::async_trait;
 use cubesql::{
     di_service,
-    sql::{AuthContext, AuthenticateResponse, SqlAuthService},
+    sql::{AuthContext, AuthenticateResponse, SqlAuthService, SqlAuthServiceAuthenticateRequest},
     transport::LoadRequestMeta,
     CubeError,
 };
@@ -50,9 +50,28 @@ pub struct TransportRequest {
     pub meta: Option<LoadRequestMeta>,
 }
 
+#[derive(Debug, Serialize)]
+pub struct TransportAuthRequest {
+    pub id: String,
+    pub meta: Option<LoadRequestMeta>,
+    pub protocol: String,
+    pub method: String,
+}
+
+impl From<(TransportRequest, SqlAuthServiceAuthenticateRequest)> for TransportAuthRequest {
+    fn from((t, a): (TransportRequest, SqlAuthServiceAuthenticateRequest)) -> Self {
+        Self {
+            id: t.id,
+            meta: t.meta,
+            protocol: a.protocol,
+            method: a.method,
+        }
+    }
+}
+
 #[derive(Debug, Serialize)]
 struct CheckSQLAuthTransportRequest {
-    request: TransportRequest,
+    request: TransportAuthRequest,
     user: Option<String>,
     password: Option<String>,
 }
@@ -92,6 +111,7 @@ impl AuthContext for NativeSQLAuthContext {
 impl SqlAuthService for NodeBridgeAuthService {
     async fn authenticate(
         &self,
+        request: SqlAuthServiceAuthenticateRequest,
         user: Option<String>,
         password: Option<String>,
     ) -> Result<AuthenticateResponse, CubeError> {
@@ -100,9 +120,11 @@ impl SqlAuthService for NodeBridgeAuthService {
         let request_id = Uuid::new_v4().to_string();
 
         let extra = serde_json::to_string(&CheckSQLAuthTransportRequest {
-            request: TransportRequest {
+            request: TransportAuthRequest {
                 id: format!("{}-span-1", request_id),
                 meta: None,
+                protocol: request.protocol,
+                method: request.method,
             },
             user: user.clone(),
             password: password.clone(),

packages/cubejs-backend-native/test/sql.test.ts

@@ -201,6 +201,8 @@ describe('SQLInterface', () => {
           request: {
             id: expect.any(String),
             meta: null,
+            method: expect.any(String),
+            protocol: expect.any(String),
           },
           user: user || null,
           password:
@@ -258,6 +260,8 @@ describe('SQLInterface', () => {
         request: {
           id: expect.any(String),
           meta: null,
+          method: expect.any(String),
+          protocol: expect.any(String),
         },
         user: 'allowed_user',
         password: 'password_for_allowed_user',

packages/cubejs-testing/birdbox-fixtures/postgresql/single/sqlapi.js

@@ -13,6 +13,15 @@ module.exports = {
     return query;
   },
   checkSqlAuth: async (req, user, password) => {
+    if (!req) {
+      throw new Error('Request is not defined');
+    }
+
+    const missing = ['protocol', 'method'].filter(key => !(key in req));
+    if (missing.length) {
+      throw new Error(`Request object is missing required field(s): ${missing.join(', ')}`);
+    }
+
     if (user === 'admin') {
       if (password && password !== 'admin_password') {
         throw new Error(`Password doesn't match for ${user}`);

rust/cubesql/cubesql/src/compile/router.rs

@@ -12,6 +12,7 @@ use crate::{
         DatabaseVariable, DatabaseVariablesToUpdate,
     },
     sql::{
+        auth_service::SqlAuthServiceAuthenticateRequest,
         dataframe,
         statement::{
             ApproximateCountDistinctVisitor, CastReplacer, DateTokenNormalizeReplacer,
@@ -447,12 +448,16 @@ impl QueryRouter {
                 })?
             {
                 self.state.set_user(Some(to_user.clone()));
+                let sql_auth_request = SqlAuthServiceAuthenticateRequest {
+                    protocol: "postgres".to_string(),
+                    method: "password".to_string(),
+                };
                 let authenticate_response = self
                     .session_manager
                     .server
                     .auth
                     // TODO do we want to send actual password here?
-                    .authenticate(Some(to_user.clone()), None)
+                    .authenticate(sql_auth_request, Some(to_user.clone()), None)
                     .await
                     .map_err(|e| {
                         CompilationError::internal(format!("Error calling authenticate: {}", e))
@@ -562,11 +567,15 @@ impl QueryRouter {
 
     async fn reauthenticate_if_needed(&self) -> CompilationResult<()> {
         if self.state.is_auth_context_expired() {
+            let sql_auth_request = SqlAuthServiceAuthenticateRequest {
+                protocol: "postgres".to_string(),
+                method: "password".to_string(),
+            };
             let authenticate_response = self
                 .session_manager
                 .server
                 .auth
-                .authenticate(self.state.user(), None)
+                .authenticate(sql_auth_request, self.state.user(), None)
                 .await
                 .map_err(|e| {
                     CompilationError::fatal(format!(

rust/cubesql/cubesql/src/compile/test/mod.rs

@@ -8,9 +8,10 @@ use crate::{
     },
     config::{ConfigObj, ConfigObjImpl},
     sql::{
-        compiler_cache::CompilerCacheImpl, dataframe::batches_to_dataframe,
-        pg_auth_service::PostgresAuthServiceDefaultImpl, AuthContextRef, AuthenticateResponse,
-        HttpAuthContext, ServerManager, Session, SessionManager, SqlAuthService,
+        auth_service::SqlAuthServiceAuthenticateRequest, compiler_cache::CompilerCacheImpl,
+        dataframe::batches_to_dataframe, pg_auth_service::PostgresAuthServiceDefaultImpl,
+        AuthContextRef, AuthenticateResponse, HttpAuthContext, ServerManager, Session,
+        SessionManager, SqlAuthService,
     },
     transport::{
         CubeMeta, CubeMetaDimension, CubeMetaJoin, CubeMetaMeasure, CubeMetaSegment,
@@ -747,6 +748,7 @@ pub fn get_test_auth() -> Arc<dyn SqlAuthService> {
     impl SqlAuthService for TestSqlAuth {
         async fn authenticate(
             &self,
+            _request: SqlAuthServiceAuthenticateRequest,
             _user: Option<String>,
             password: Option<String>,
         ) -> Result<AuthenticateResponse, CubeError> {

rust/cubesql/cubesql/src/sql/auth_service.rs

@@ -2,6 +2,7 @@ use std::{any::Any, env, fmt::Debug, sync::Arc};
 
 use crate::CubeError;
 use async_trait::async_trait;
+use serde::{Deserialize, Serialize};
 use serde_json::Value;
 
 // We cannot use generic here. It's why there is this trait
@@ -43,10 +44,17 @@ pub struct AuthenticateResponse {
     pub skip_password_check: bool,
 }
 
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct SqlAuthServiceAuthenticateRequest {
+    pub protocol: String,
+    pub method: String,
+}
+
 #[async_trait]
 pub trait SqlAuthService: Send + Sync + Debug {
     async fn authenticate(
         &self,
+        request: SqlAuthServiceAuthenticateRequest,
         user: Option<String>,
         password: Option<String>,
     ) -> Result<AuthenticateResponse, CubeError>;
@@ -61,6 +69,7 @@ crate::di_service!(SqlAuthDefaultImpl, [SqlAuthService]);
 impl SqlAuthService for SqlAuthDefaultImpl {
     async fn authenticate(
         &self,
+        _request: SqlAuthServiceAuthenticateRequest,
         _user: Option<String>,
         password: Option<String>,
     ) -> Result<AuthenticateResponse, CubeError> {

rust/cubesql/cubesql/src/sql/mod.rs

@@ -13,7 +13,7 @@ pub(crate) mod types;
 // Public API
 pub use auth_service::{
     AuthContext, AuthContextRef, AuthenticateResponse, HttpAuthContext, SqlAuthDefaultImpl,
-    SqlAuthService,
+    SqlAuthService, SqlAuthServiceAuthenticateRequest,
 };
 pub use database_variables::postgres::session_vars::CUBESQL_PENALIZE_POST_PROCESSING_VAR;
 pub use postgres::*;

rust/cubesql/cubesql/src/sql/postgres/pg_auth_service.rs

@@ -3,6 +3,7 @@ use std::{collections::HashMap, fmt::Debug, sync::Arc};
 use async_trait::async_trait;
 
 use crate::{
+    sql::auth_service::SqlAuthServiceAuthenticateRequest,
     sql::{AuthContextRef, SqlAuthService},
     CubeError,
 };
@@ -74,8 +75,16 @@ impl PostgresAuthService for PostgresAuthServiceDefaultImpl {
         }
 
         let user = parameters.get("user").unwrap().clone();
+        let sql_auth_request = SqlAuthServiceAuthenticateRequest {
+            protocol: "postgres".to_string(),
+            method: "password".to_string(),
+        };
         let authenticate_response = service
-            .authenticate(Some(user.clone()), Some(password_message.password.clone()))
+            .authenticate(
+                sql_auth_request,
+                Some(user.clone()),
+                Some(password_message.password.clone()),
+            )
             .await;
 
         let auth_fail = || {