feat(cubejs-api-gateway): Support returning new security context from…

… check_auth (#8585) This is useful with Python configs, where it's input checkAuth context is a copy of original context --------- Co-authored-by: Igor Lukanin <[email protected]>
cube-js · Aug 28, 2024 · 704a96c · 704a96c
1 parent 11b6dec
commit 704a96c
Show file tree

Hide file tree

Showing 9 changed files with 169 additions and 38 deletions.
diff --git a/docs/pages/product/auth.mdx b/docs/pages/product/auth.mdx
@@ -230,7 +230,7 @@ Cube caches JWKS by default when
 
 Cube also allows you to provide your own JWT verification logic by setting a
 [`checkAuth()`][ref-config-check-auth] function in the `cube.js` configuration
-file. This function is expected to verify a JWT and assigns its' claims to the
+file. This function is expected to verify a JWT and return its claims as the
 security context.
 
 <WarningBox>
@@ -249,7 +249,7 @@ module.exports = {
   checkAuth: async (req, auth) => {
     try {
       const userInfo = await getUserFromLDAP(req.get("X-LDAP-User-ID"));
-      req.securityContext = userInfo;
+      return { security_context: userInfo };
     } catch {
       throw new Error("Could not authenticate user from LDAP");
     }

diff --git a/docs/pages/reference/configuration/config.mdx b/docs/pages/reference/configuration/config.mdx
@@ -1079,23 +1079,14 @@ environment documentation][ref-exec-environment-globals].
 
 ### `check_auth`
 
-Used in both REST and WebSockets API.
+Used in the [REST API][ref-rest-api]. Default implementation parses the [JSON
+Web Token][link-jwt] in the `Authorization` header, verifies it, and sets its
+payload to the `securityContext`. [Read more][ref-sec-ctx] about JWT generation.
 
 Called on each request.
 
-Default implementation parses [JSON Web Token (JWT)][link-jwt] in `Authorization`
-header and sets payload to `securityContext` if it's verified. More
-information on how to generate these tokens is [here][ref-sec-ctx].
-
-You can set `securityContext = userContextObj` inside the middleware if you
-want to customize [`SECURITY_CONTEXT`][ref-schema-cube-ref-ctx-sec-ctx].
-
-<ReferenceBox>
-
-Currently, assigning to security context doesn't work in Python.
-Please [track this issue](https://github.com/cube-js/cube/issues/8133).
-
-</ReferenceBox>
+You can return an object with the `security_context` field if you want to
+customize [`SECURITY_CONTEXT`][ref-schema-cube-ref-ctx-sec-ctx].
 
 You can use empty `check_auth` function to disable built-in security or
 raise an exception to fail the authentication check.
@@ -1107,19 +1098,25 @@ from cube import config
 
 @config('check_auth')
 def check_auth(ctx: dict, token: str) -> None:
-  context = ctx['securityContext']
-
   if token == 'my_secret_token':
-    return
+    return {
+      'security_context': {
+        'user_id': 42
+      }
+    }
 
   raise Exception('Access denied')
 ```
 
 ```javascript
 module.exports = {
-  checkAuth: ({ securityContext }, token) => {
+  checkAuth: (ctx, token) => {
     if (token === 'my_secret_token') {
-      return;
+      return {
+        security_context: {
+          user_id: 42
+        }
+      }
     }
 
     throw new Error('Access denied');

diff --git a/packages/cubejs-api-gateway/src/gateway.ts b/packages/cubejs-api-gateway/src/gateway.ts
@@ -113,6 +113,9 @@ function systemAsyncHandler(handler: (req: Request & { context: ExtendedRequestC
   };
 }
 
+// Prepared CheckAuthFn, default or from config: always async, returns nothing
+type PreparedCheckAuthFn = (ctx: any, authorization?: string) => Promise<void>;
+
 class ApiGateway {
   protected readonly refreshScheduler: any;
 
@@ -134,9 +137,9 @@ class ApiGateway {
 
   protected readonly dataSourceStorage: any;
 
-  public readonly checkAuthFn: CheckAuthFn;
+  public readonly checkAuthFn: PreparedCheckAuthFn;
 
-  public readonly checkAuthSystemFn: CheckAuthFn;
+  public readonly checkAuthSystemFn: PreparedCheckAuthFn;
 
   protected readonly contextToApiScopesFn: ContextToApiScopesFn;
 
@@ -2148,14 +2151,19 @@ class ApiGateway {
     };
   }
 
-  protected wrapCheckAuth(fn: CheckAuthFn): CheckAuthFn {
+  protected wrapCheckAuth(fn: CheckAuthFn): PreparedCheckAuthFn {
     // We dont need to span all logs with deprecation message
     let warningShowed = false;
     // securityContext should be object
     let showWarningAboutNotObject = false;
 
     return async (req, auth) => {
-      await fn(req, auth);
+      const result = await fn(req, auth);
+
+      // checkAuth from config can return new security context, e.g from Python config
+      if (result?.security_context) {
+        req.securityContext = result?.security_context;
+      }
 
       // We renamed authInfo to securityContext, but users can continue to use both ways
       if (req.securityContext && !req.authInfo) {
@@ -2187,7 +2195,7 @@ class ApiGateway {
     };
   }
 
-  protected createDefaultCheckAuth(options?: JWTOptions, internalOptions?: CheckAuthInternalOptions): CheckAuthFn {
+  protected createDefaultCheckAuth(options?: JWTOptions, internalOptions?: CheckAuthInternalOptions): PreparedCheckAuthFn {
     type VerifyTokenFn = (auth: string, secret: string) => Promise<object | string> | object | string;
 
     const verifyToken = (auth, secret) => jwt.verify(auth, secret, {
@@ -2270,7 +2278,7 @@ class ApiGateway {
     };
   }
 
-  protected createCheckAuthFn(options: ApiGatewayOptions): CheckAuthFn {
+  protected createCheckAuthFn(options: ApiGatewayOptions): PreparedCheckAuthFn {
     const mainCheckAuthFn = options.checkAuth
       ? this.wrapCheckAuth(options.checkAuth)
       : this.createDefaultCheckAuth(options.jwt);
@@ -2289,7 +2297,7 @@ class ApiGateway {
     return (ctx, authorization) => mainCheckAuthFn(ctx, authorization);
   }
 
-  protected createCheckAuthSystemFn(): CheckAuthFn {
+  protected createCheckAuthSystemFn(): PreparedCheckAuthFn {
     const systemCheckAuthFn = this.createDefaultCheckAuth(
       {
         key: this.playgroundAuthSecret,
@@ -2370,7 +2378,7 @@ class ApiGateway {
     return undefined;
   }
 
-  protected async checkAuthWrapper(checkAuthFn: CheckAuthFn, req: Request, res: ExpressResponse, next) {
+  protected async checkAuthWrapper(checkAuthFn: PreparedCheckAuthFn, req: Request, res: ExpressResponse, next) {
     const token = this.extractAuthorizationHeaderWithSchema(req);
 
     try {

diff --git a/packages/cubejs-api-gateway/src/types/auth.ts b/packages/cubejs-api-gateway/src/types/auth.ts
@@ -39,15 +39,19 @@ interface JWTOptions {
   claimsNamespace?: string,
 }
 
+type CheckAuthResponse = {
+  'security_context'?: unknown,
+};
+
 /**
  * Function that should provides basic auth mechanic. Used as a part
  * of a main configuration object of the server-core to provide base
- * auth logic.
+ * auth logic. Can return new security context.
  * @todo ctx can be passed from SubscriptionServer that will cause
  * incapability with Express.Request
  */
 type CheckAuthFn =
-  (ctx: any, authorization?: string) => Promise<void> | void;
+  (ctx: any, authorization?: string) => Promise<void | CheckAuthResponse> | CheckAuthResponse | void;
 
 /**
  * Result of the SQL auth workflow.

diff --git a/packages/cubejs-api-gateway/test/auth.test.ts b/packages/cubejs-api-gateway/test/auth.test.ts
@@ -128,7 +128,7 @@ describe('test authorization', () => {
       .get('/test-auth-fake')
       .set('Authorization', `Authorization: ${badToken}`)
       .expect(403);
-      
+
     expect(loggerMock.mock.calls.length).toEqual(1);
     expect(handlerMock.mock.calls.length).toEqual(2);
 
@@ -247,6 +247,58 @@ describe('test authorization', () => {
     });
   });
 
+  test('custom checkAuth with async flow and return', async () => {
+    const loggerMock = jest.fn(() => {
+      //
+    });
+
+    const expectSecurityContext = (securityContext) => {
+      expect(securityContext.uid).toEqual(5);
+      expect(securityContext.iat).toBeDefined();
+      expect(securityContext.exp).toBeDefined();
+    };
+
+    const handlerMock = jest.fn((req, res) => {
+      expectSecurityContext(req.context.securityContext);
+      expectSecurityContext(req.context.authInfo);
+
+      res.status(200).end();
+    });
+
+    const { app } = createApiGateway(handlerMock, loggerMock, {
+      checkAuth: async (req: Request, auth?: string) => {
+        if (auth) {
+          await pausePromise(500);
+
+          const securityContext = jwt.verify(auth, 'secret');
+
+          req.securityContext = {
+            uid: 'should not be visible',
+          };
+
+          return {
+            security_context: securityContext,
+          };
+        }
+
+        return {};
+      }
+    });
+
+    const token = generateAuthToken({ uid: 5, });
+
+    await request(app)
+      .get('/test-auth-fake')
+      .set('Authorization', `Authorization: ${token}`)
+      .expect(200);
+
+    expect(handlerMock.mock.calls.length).toEqual(1);
+
+    expectSecurityContext(handlerMock.mock.calls[0][0].context.securityContext);
+    // authInfo was deprecated, but should exist as computability
+    expectSecurityContext(handlerMock.mock.calls[0][0].context.authInfo);
+  });
+
   test('custom checkAuth with deprecated authInfo', async () => {
     const loggerMock = jest.fn(() => {
       //

diff --git a/packages/cubejs-api-gateway/test/index.test.ts b/packages/cubejs-api-gateway/test/index.test.ts
@@ -255,6 +255,54 @@ describe('API Gateway', () => {
     expect(queryRewrite.mock.calls.length).toEqual(1);
   });
 
+  test('query transform with checkAuth with return', async () => {
+    const queryRewrite = jest.fn(async (query: Query, context) => {
+      expect(context.securityContext).toEqual({
+        exp: 2475857705,
+        iat: 1611857705,
+        uid: 5
+      });
+
+      expect(context.authInfo).toEqual({
+        exp: 2475857705,
+        iat: 1611857705,
+        uid: 5
+      });
+
+      return query;
+    });
+
+    const { app } = await createApiGateway(
+      new AdapterApiMock(),
+      new DataSourceStorageMock(),
+      {
+        checkAuth: (req: Request, authorization) => {
+          if (authorization) {
+            return {
+              security_context: jwt.verify(authorization, API_SECRET),
+            };
+          }
+
+          return {};
+        },
+        queryRewrite
+      }
+    );
+
+    const res = await request(app)
+      .get(
+        '/cubejs-api/v1/load?query={"measures":["Foo.bar"],"filters":[{"dimension":"Foo.id","operator":"equals","values":[null]}]}'
+      )
+      // console.log(generateAuthToken({ uid: 5, }));
+      .set('Authorization', 'Authorization: eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJ1aWQiOjUsImlhdCI6MTYxMTg1NzcwNSwiZXhwIjoyNDc1ODU3NzA1fQ.tTieqdIcxDLG8fHv8YWwfvg_rPVe1XpZKUvrCdzVn3g')
+      .expect(200);
+
+    console.log(res.body);
+    expect(res.body && res.body.data).toStrictEqual([{ 'Foo.bar': 42 }]);
+
+    expect(queryRewrite.mock.calls.length).toEqual(1);
+  });
+
   test('null filter values', async () => {
     const { app } = await createApiGateway();
 

diff --git a/packages/cubejs-backend-native/js/index.ts b/packages/cubejs-backend-native/js/index.ts
@@ -351,7 +351,7 @@ export const buildSqlAndParams = (cubeEvaluator: any): String => {
 export interface PyConfiguration {
   repositoryFactory?: (ctx: unknown) => Promise<unknown>,
   logger?: (msg: string, params: Record<string, any>) => void,
-  checkAuth?: (req: unknown, authorization: string) => Promise<void>
+  checkAuth?: (req: unknown, authorization: string) => Promise<{ 'security_context'?: unknown }>
   queryRewrite?: (query: unknown, ctx: unknown) => Promise<unknown>
   contextToApiScopes?: () => Promise<string[]>
 }
@@ -377,10 +377,16 @@ export const pythonLoadConfig = async (content: string, options: { fileName: str
 
   if (config.checkAuth) {
     const nativeCheckAuth = config.checkAuth;
-    config.checkAuth = async (req: ExpressRequest, authorization: string) => nativeCheckAuth(
-      simplifyExpressRequest(req),
-      authorization,
-    );
+    config.checkAuth = async (req: ExpressRequest, authorization: string) => {
+      const nativeResult = await nativeCheckAuth(
+        simplifyExpressRequest(req),
+        authorization,
+      );
+      const securityContext = nativeResult?.security_context;
+      return {
+        ...(securityContext ? { security_context: securityContext } : {})
+      };
+    };
   }
 
   if (config.extendContext) {

diff --git a/packages/cubejs-backend-native/test/config.py b/packages/cubejs-backend-native/test/config.py
@@ -15,6 +15,14 @@ def query_rewrite(query, ctx):
 @config
 async def check_auth(req, authorization):
     print('[python] check_auth req=', req, ' authorization=', authorization)
+    return {
+      'security_context': {
+        'sub': '1234567890',
+        'iat': 1516239022,
+        'user_id': 42
+      },
+      'ignoredField': 'should not be visible'
+    }
 
 @config
 async def repository_factory(ctx):

diff --git a/packages/cubejs-backend-native/test/python.test.ts b/packages/cubejs-backend-native/test/python.test.ts
@@ -52,10 +52,18 @@ suite('Python Config', () => {
       throw new Error('checkAuth was not defined in config.py');
     }
 
-    await config.checkAuth(
+    const result = await config.checkAuth(
       { requestId: 'test' },
       'MY_SECRET_TOKEN'
     );
+
+    expect(result).toEqual({
+      security_context: {
+        sub: '1234567890',
+        iat: 1516239022,
+        user_id: 42
+      },
+    });
   });
 
   test('context_to_api_scopes', async () => {