From cf360b9c15594b0923c79adf3a401e29d84fad23 Mon Sep 17 00:00:00 2001
From: simeng-li <simeng@silverhand.io>
Date: Tue, 24 May 2022 16:42:28 +0800
Subject: [PATCH] feat(core): add admin role validation to the koaAuth (#920)

* feat(core): add admin role validation to the koaAuth

add admin role validation to the koaAuth

* fix(core): cr fix

cr fix
---
 packages/core/src/middleware/koa-auth.test.ts | 32 ++++++++++++++++++-
 packages/core/src/middleware/koa-auth.ts      | 14 +++++---
 packages/core/src/oidc/init.ts                | 13 ++++++++
 3 files changed, 54 insertions(+), 5 deletions(-)

diff --git a/packages/core/src/middleware/koa-auth.test.ts b/packages/core/src/middleware/koa-auth.test.ts
index be9896b9c..63257e6fe 100644
--- a/packages/core/src/middleware/koa-auth.test.ts
+++ b/packages/core/src/middleware/koa-auth.test.ts
@@ -9,7 +9,7 @@ import { createContextWithRouteParameters } from '@/utils/test-utils';
 import koaAuth, { WithAuthContext } from './koa-auth';
 
 jest.mock('jose', () => ({
-  jwtVerify: jest.fn(() => ({ payload: { sub: 'fooUser' } })),
+  jwtVerify: jest.fn(() => ({ payload: { sub: 'fooUser', roleNames: ['admin'] } })),
 }));
 
 describe('koaAuth middleware', () => {
@@ -80,4 +80,34 @@ describe('koaAuth middleware', () => {
 
     await expect(koaAuth()(ctx, next)).rejects.toMatchError(unauthorizedError);
   });
+
+  it('expect to throw if jwt roleNames is missing', async () => {
+    const mockJwtVerify = jwtVerify as jest.Mock;
+    mockJwtVerify.mockImplementationOnce(() => ({ payload: { sub: 'fooUser' } }));
+
+    ctx.request = {
+      ...ctx.request,
+      headers: {
+        authorization: 'Bearer access_token',
+      },
+    };
+
+    await expect(koaAuth()(ctx, next)).rejects.toMatchError(unauthorizedError);
+  });
+
+  it('expect to throw if jwt roleNames does not include admin', async () => {
+    const mockJwtVerify = jwtVerify as jest.Mock;
+    mockJwtVerify.mockImplementationOnce(() => ({
+      payload: { sub: 'fooUser', roleNames: ['foo'] },
+    }));
+
+    ctx.request = {
+      ...ctx.request,
+      headers: {
+        authorization: 'Bearer access_token',
+      },
+    };
+
+    await expect(koaAuth()(ctx, next)).rejects.toMatchError(unauthorizedError);
+  });
 });
diff --git a/packages/core/src/middleware/koa-auth.ts b/packages/core/src/middleware/koa-auth.ts
index cd2575394..72112a646 100644
--- a/packages/core/src/middleware/koa-auth.ts
+++ b/packages/core/src/middleware/koa-auth.ts
@@ -1,6 +1,6 @@
 import { IncomingHttpHeaders } from 'http';
 
-import { managementApiResource } from '@logto/schemas';
+import { managementApiResource, UserRole } from '@logto/schemas';
 import { jwtVerify } from 'jose';
 import { MiddlewareType, Request } from 'koa';
 import { IRouterParamContext } from 'koa-router';
@@ -32,7 +32,7 @@ const extractBearerTokenFromHeaders = ({ authorization }: IncomingHttpHeaders) =
   return authorization.slice(bearerTokenIdentifier.length + 1);
 };
 
-const getUserIdFromRequest = async (request: Request) => {
+const getUserInfoFromRequest = async (request: Request) => {
   const { isProduction, developmentUserId, oidc } = envSet.values;
 
   if (!isProduction && developmentUserId) {
@@ -41,13 +41,19 @@ const getUserIdFromRequest = async (request: Request) => {
 
   const { publicKey, issuer } = oidc;
   const {
-    payload: { sub },
+    payload: { sub, roleNames },
   } = await jwtVerify(extractBearerTokenFromHeaders(request.headers), publicKey, {
     issuer,
     audience: managementApiResource,
   });
+
   assertThat(sub, new RequestError({ code: 'auth.jwt_sub_missing', status: 401 }));
 
+  assertThat(
+    Array.isArray(roleNames) && roleNames.includes(UserRole.Admin),
+    new RequestError({ code: 'auth.unauthorized', status: 401 })
+  );
+
   return sub;
 };
 
@@ -58,7 +64,7 @@ export default function koaAuth<
 >(): MiddlewareType<StateT, WithAuthContext<ContextT>, ResponseBodyT> {
   return async (ctx, next) => {
     try {
-      const userId = await getUserIdFromRequest(ctx.request);
+      const userId = await getUserInfoFromRequest(ctx.request);
       ctx.auth = userId;
     } catch {
       throw new RequestError({ code: 'auth.unauthorized', status: 401 });
diff --git a/packages/core/src/oidc/init.ts b/packages/core/src/oidc/init.ts
index 8c9056a6a..99ea63c84 100644
--- a/packages/core/src/oidc/init.ts
+++ b/packages/core/src/oidc/init.ts
@@ -115,6 +115,19 @@ export default async function initOidc(app: Koa): Promise<Provider> {
         return refreshTokenTtl ?? defaultRefreshTokenTtl;
       },
     },
+    extraTokenClaims: async (ctx, token) => {
+      // AccessToken type is not exported by default, need to asset token is AccessToken
+      if (token.kind === 'AccessToken') {
+        const { accountId } = token;
+        const { roleNames } = await findUserById(accountId);
+
+        // Add User Roles to the AccessToken claims. Should be removed once we have RBAC implemented.
+        // User Roles should be hidden and  determined by the AccessToken scope only.
+        return {
+          roleNames,
+        };
+      }
+    },
   });
 
   addOidcEventListeners(oidc);