From 739e27b80277d3238ee912f18c54971ea5ebaaee Mon Sep 17 00:00:00 2001
From: Max Konin <maxim21214@gmail.com>
Date: Wed, 17 May 2023 14:56:05 +0700
Subject: [PATCH] Use Cognito JWT to extract auth data

---
 packages/server/src/index.ts                   | 2 ++
 packages/server/src/utils/useLambdaIdentity.ts | 8 +++++---
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/packages/server/src/index.ts b/packages/server/src/index.ts
index 855d154..1f3daf8 100644
--- a/packages/server/src/index.ts
+++ b/packages/server/src/index.ts
@@ -93,6 +93,7 @@ export async function createServer({ defaultQuery, lambdaHandler, port, schema,
                             try { prismaAppSyncHeader = JSON.parse(request?.headers?.['x-prisma-appsync']) }
                             catch { prismaAppSyncHeader = {} }
 
+                            const authorizationHeader = request?.headers?.authorization
                             const authorization = prismaAppSyncHeader?.authorization || Authorizations.AWS_IAM || null
                             const signature = prismaAppSyncHeader?.signature || {}
 
@@ -102,6 +103,7 @@ export async function createServer({ defaultQuery, lambdaHandler, port, schema,
                                     username: 'johndoe',
                                     sub: 'xxxxxx',
                                     resolverContext: {},
+                                    jwt: authorizationHeader,
                                 },
                                 ...signature,
                             })
diff --git a/packages/server/src/utils/useLambdaIdentity.ts b/packages/server/src/utils/useLambdaIdentity.ts
index d71e7c3..856ca87 100644
--- a/packages/server/src/utils/useLambdaIdentity.ts
+++ b/packages/server/src/utils/useLambdaIdentity.ts
@@ -26,11 +26,12 @@ export default function useLambdaIdentity(identity: Authorization, opts?: mockOp
         return mock
     }
     else if (identity === Authorizations.AMAZON_COGNITO_USER_POOLS) {
+        const decodedJWTToken = opts?.jwt ? JSON.parse(Buffer.from(opts?.jwt?.split('.')[1], 'base64').toString()) : {}
         const mock: AMAZON_COGNITO_USER_POOLS = {
-            sub: opts?.sub || 'undefined',
+            sub: decodedJWTToken?.sub || 'undefined',
             issuer: 'string',
-            username: opts?.username || 'undefined',
-            claims: {},
+            username: decodedJWTToken?.['cognito:username'] || 'undefined',
+            claims: Object.fromEntries(Object.entries(decodedJWTToken).filter(([key]) => key.includes('custom:'))),
             sourceIp: [opts?.sourceIp || 'undefined'],
             defaultAuthStrategy: 'string',
             groups: ['admin', 'member'],
@@ -70,4 +71,5 @@ interface mockOptions {
     username: string
     sourceIp: string
     resolverContext: any
+    jwt: string
 }