Fix for verifyToken failure bug in closed network mode

packages/cdk/lambda/utils/auth.ts

@@ -1,17 +1,29 @@
-import { CognitoJwtVerifier } from 'aws-jwt-verify';
+import { JwtRsaVerifier } from 'aws-jwt-verify';
 import { CognitoIdTokenPayload } from 'aws-jwt-verify/jwt-model';
 
 export const verifyToken = async (
   token: string
 ): Promise<CognitoIdTokenPayload | undefined> => {
   try {
-    const verifier = CognitoJwtVerifier.create({
-      userPoolId: process.env.USER_POOL_ID || '',
+    const region = process.env.AWS_REGION!;
+    const userPoolId = process.env.USER_POOL_ID!;
+    const clientId = process.env.USER_POOL_CLIENT_ID!;
+    const proxyEndpoint = process.env.USER_POOL_PROXY_ENDPOINT!;
+
+    const jwksUri =
+      proxyEndpoint.length > 0
+        ? `${proxyEndpoint}${userPoolId}/.well-known/jwks.json`
+        : `https://cognito-idp.${region}.amazonaws.com/${userPoolId}/.well-known/jwks.json`;
+
+    const verifier = JwtRsaVerifier.create({
+      issuer: `https://cognito-idp.${region}.amazonaws.com/${userPoolId}`,
+      audience: clientId,
+      jwksUri,
       tokenUse: 'id',
-      clientId: process.env.USER_POOL_CLIENT_ID || '',
     });
+
     const payload = await verifier.verify(token);
-    return payload;
+    return payload as CognitoIdTokenPayload;
   } catch {
     return undefined;
   }

packages/cdk/lib/construct/api.ts

@@ -72,6 +72,7 @@ export interface BackendApiProps {
   readonly vpc?: IVpc;
   readonly securityGroups?: ISecurityGroup[];
   readonly apiGatewayVpcEndpoint?: InterfaceVpcEndpoint;
+  readonly cognitoUserPoolProxyEndpoint?: string;
 }
 
 export class Api extends Construct {
@@ -205,6 +206,7 @@ export class Api extends Construct {
       environment: {
         USER_POOL_ID: userPool.userPoolId,
         USER_POOL_CLIENT_ID: userPoolClient.userPoolClientId,
+        USER_POOL_PROXY_ENDPOINT: props.cognitoUserPoolProxyEndpoint ?? '',
         MODEL_REGION: modelRegion,
         MODEL_IDS: JSON.stringify(modelIds),
         IMAGE_GENERATION_MODEL_IDS: JSON.stringify(imageGenerationModelIds),

packages/cdk/lib/construct/closedNetwork/cognito-private-proxy.ts

@@ -26,7 +26,7 @@ export class CognitoPrivateProxy extends Construct {
         restApiName: 'GenU Cognito UserPool Proxy API',
         defaultCorsPreflightOptions: {
           allowOrigins: agw.Cors.ALL_ORIGINS,
-          allowMethods: ['POST', 'OPTIONS'],
+          allowMethods: ['GET', 'POST', 'OPTIONS'],
           allowHeaders: [
             'amz-sdk-invocation-id',
             'amz-sdk-request',
@@ -64,6 +64,44 @@ export class CognitoPrivateProxy extends Construct {
       }
     );
 
+    const userPoolResource =
+      this.cognitoUserPoolProxyApi.root.addResource('{userPoolId}');
+    const wellKnownResource = userPoolResource.addResource('.well-known');
+    const jwksResource = wellKnownResource.addResource('jwks.json');
+
+    // Add GET method to the root resource to serve jwtks.json
+    jwksResource.addMethod(
+      'GET',
+      new agw.HttpIntegration(
+        `${cognitoUserPoolEndpoint}/{userPoolId}/.well-known/jwks.json`,
+        {
+          proxy: false,
+          httpMethod: 'GET',
+          options: {
+            requestParameters: {
+              'integration.request.path.userPoolId':
+                'method.request.path.userPoolId',
+            },
+            integrationResponses: [
+              {
+                statusCode: '200',
+              },
+            ],
+          },
+        }
+      ),
+      {
+        requestParameters: {
+          'method.request.path.userPoolId': true,
+        },
+        methodResponses: [
+          {
+            statusCode: '200',
+          },
+        ],
+      }
+    );
+
     // Add POST method to the root resource to proxy Cognito UserPool requests
     this.cognitoUserPoolProxyApi.root.addMethod(
       'POST',

packages/cdk/lib/generative-ai-use-cases-stack.ts

@@ -124,6 +124,7 @@ export class GenerativeAiUseCasesStack extends Stack {
       vpc: props.vpc,
       securityGroups,
       apiGatewayVpcEndpoint: props.apiGatewayVpcEndpoint,
+      cognitoUserPoolProxyEndpoint: props.cognitoUserPoolProxyEndpoint,
     });
 
     // WAF

packages/cdk/test/__snapshots__/generative-ai-use-cases.test.ts.snap

@@ -1808,8 +1808,15 @@ exports[`GenerativeAiUseCases matches the snapshot (closed network mode) 1`] = `
       "Type": "AWS::IAM::Role",
       "UpdateReplacePolicy": "Retain",
     },
-    "CognitoPrivateProxyCognitoUserPoolProxyApiDeployment11D70EC105771a0029216e27aa20a3859719b309": {
+    "CognitoPrivateProxyCognitoUserPoolProxyApiDeployment11D70EC1b5bb02dcd4261d6e44e76aa0863f6287": {
       "DependsOn": [
+        "CognitoPrivateProxyCognitoUserPoolProxyApiuserPoolIdwellknownjwksjsonGET21D71187",
+        "CognitoPrivateProxyCognitoUserPoolProxyApiuserPoolIdwellknownjwksjsonOPTIONSA2D13B9E",
+        "CognitoPrivateProxyCognitoUserPoolProxyApiuserPoolIdwellknownjwksjsonD9640238",
+        "CognitoPrivateProxyCognitoUserPoolProxyApiuserPoolIdwellknownOPTIONSB9A49912",
+        "CognitoPrivateProxyCognitoUserPoolProxyApiuserPoolIdwellknownCE432EBE",
+        "CognitoPrivateProxyCognitoUserPoolProxyApiuserPoolIdOPTIONSC878F5C3",
+        "CognitoPrivateProxyCognitoUserPoolProxyApiuserPoolId5D61C608",
         "CognitoPrivateProxyCognitoUserPoolProxyApiOPTIONSCE339C8F",
         "CognitoPrivateProxyCognitoUserPoolProxyApiPOST0CAE9A7C",
       ],
@@ -1830,7 +1837,7 @@ exports[`GenerativeAiUseCases matches the snapshot (closed network mode) 1`] = `
       ],
       "Properties": {
         "DeploymentId": {
-          "Ref": "CognitoPrivateProxyCognitoUserPoolProxyApiDeployment11D70EC105771a0029216e27aa20a3859719b309",
+          "Ref": "CognitoPrivateProxyCognitoUserPoolProxyApiDeployment11D70EC1b5bb02dcd4261d6e44e76aa0863f6287",
         },
         "RestApiId": {
           "Ref": "CognitoPrivateProxyCognitoUserPoolProxyApi818B53B1",
@@ -1849,7 +1856,7 @@ exports[`GenerativeAiUseCases matches the snapshot (closed network mode) 1`] = `
             {
               "ResponseParameters": {
                 "method.response.header.Access-Control-Allow-Headers": "'amz-sdk-invocation-id,amz-sdk-request,cache-control,content-type,x-amz-target,x-amz-user-agent'",
-                "method.response.header.Access-Control-Allow-Methods": "'POST,OPTIONS'",
+                "method.response.header.Access-Control-Allow-Methods": "'GET,POST,OPTIONS'",
                 "method.response.header.Access-Control-Allow-Origin": "'*'",
               },
               "StatusCode": "204",
@@ -1936,6 +1943,199 @@ exports[`GenerativeAiUseCases matches the snapshot (closed network mode) 1`] = `
       },
       "Type": "AWS::ApiGateway::Method",
     },
+    "CognitoPrivateProxyCognitoUserPoolProxyApiuserPoolId5D61C608": {
+      "Properties": {
+        "ParentId": {
+          "Fn::GetAtt": [
+            "CognitoPrivateProxyCognitoUserPoolProxyApi818B53B1",
+            "RootResourceId",
+          ],
+        },
+        "PathPart": "{userPoolId}",
+        "RestApiId": {
+          "Ref": "CognitoPrivateProxyCognitoUserPoolProxyApi818B53B1",
+        },
+      },
+      "Type": "AWS::ApiGateway::Resource",
+    },
+    "CognitoPrivateProxyCognitoUserPoolProxyApiuserPoolIdOPTIONSC878F5C3": {
+      "Properties": {
+        "ApiKeyRequired": false,
+        "AuthorizationType": "NONE",
+        "HttpMethod": "OPTIONS",
+        "Integration": {
+          "IntegrationResponses": [
+            {
+              "ResponseParameters": {
+                "method.response.header.Access-Control-Allow-Headers": "'amz-sdk-invocation-id,amz-sdk-request,cache-control,content-type,x-amz-target,x-amz-user-agent'",
+                "method.response.header.Access-Control-Allow-Methods": "'GET,POST,OPTIONS'",
+                "method.response.header.Access-Control-Allow-Origin": "'*'",
+              },
+              "StatusCode": "204",
+            },
+          ],
+          "RequestTemplates": {
+            "application/json": "{ statusCode: 200 }",
+          },
+          "Type": "MOCK",
+        },
+        "MethodResponses": [
+          {
+            "ResponseParameters": {
+              "method.response.header.Access-Control-Allow-Headers": true,
+              "method.response.header.Access-Control-Allow-Methods": true,
+              "method.response.header.Access-Control-Allow-Origin": true,
+            },
+            "StatusCode": "204",
+          },
+        ],
+        "ResourceId": {
+          "Ref": "CognitoPrivateProxyCognitoUserPoolProxyApiuserPoolId5D61C608",
+        },
+        "RestApiId": {
+          "Ref": "CognitoPrivateProxyCognitoUserPoolProxyApi818B53B1",
+        },
+      },
+      "Type": "AWS::ApiGateway::Method",
+    },
+    "CognitoPrivateProxyCognitoUserPoolProxyApiuserPoolIdwellknownCE432EBE": {
+      "Properties": {
+        "ParentId": {
+          "Ref": "CognitoPrivateProxyCognitoUserPoolProxyApiuserPoolId5D61C608",
+        },
+        "PathPart": ".well-known",
+        "RestApiId": {
+          "Ref": "CognitoPrivateProxyCognitoUserPoolProxyApi818B53B1",
+        },
+      },
+      "Type": "AWS::ApiGateway::Resource",
+    },
+    "CognitoPrivateProxyCognitoUserPoolProxyApiuserPoolIdwellknownOPTIONSB9A49912": {
+      "Properties": {
+        "ApiKeyRequired": false,
+        "AuthorizationType": "NONE",
+        "HttpMethod": "OPTIONS",
+        "Integration": {
+          "IntegrationResponses": [
+            {
+              "ResponseParameters": {
+                "method.response.header.Access-Control-Allow-Headers": "'amz-sdk-invocation-id,amz-sdk-request,cache-control,content-type,x-amz-target,x-amz-user-agent'",
+                "method.response.header.Access-Control-Allow-Methods": "'GET,POST,OPTIONS'",
+                "method.response.header.Access-Control-Allow-Origin": "'*'",
+              },
+              "StatusCode": "204",
+            },
+          ],
+          "RequestTemplates": {
+            "application/json": "{ statusCode: 200 }",
+          },
+          "Type": "MOCK",
+        },
+        "MethodResponses": [
+          {
+            "ResponseParameters": {
+              "method.response.header.Access-Control-Allow-Headers": true,
+              "method.response.header.Access-Control-Allow-Methods": true,
+              "method.response.header.Access-Control-Allow-Origin": true,
+            },
+            "StatusCode": "204",
+          },
+        ],
+        "ResourceId": {
+          "Ref": "CognitoPrivateProxyCognitoUserPoolProxyApiuserPoolIdwellknownCE432EBE",
+        },
+        "RestApiId": {
+          "Ref": "CognitoPrivateProxyCognitoUserPoolProxyApi818B53B1",
+        },
+      },
+      "Type": "AWS::ApiGateway::Method",
+    },
+    "CognitoPrivateProxyCognitoUserPoolProxyApiuserPoolIdwellknownjwksjsonD9640238": {
+      "Properties": {
+        "ParentId": {
+          "Ref": "CognitoPrivateProxyCognitoUserPoolProxyApiuserPoolIdwellknownCE432EBE",
+        },
+        "PathPart": "jwks.json",
+        "RestApiId": {
+          "Ref": "CognitoPrivateProxyCognitoUserPoolProxyApi818B53B1",
+        },
+      },
+      "Type": "AWS::ApiGateway::Resource",
+    },
+    "CognitoPrivateProxyCognitoUserPoolProxyApiuserPoolIdwellknownjwksjsonGET21D71187": {
+      "Properties": {
+        "AuthorizationType": "NONE",
+        "HttpMethod": "GET",
+        "Integration": {
+          "IntegrationHttpMethod": "GET",
+          "IntegrationResponses": [
+            {
+              "StatusCode": "200",
+            },
+          ],
+          "RequestParameters": {
+            "integration.request.path.userPoolId": "method.request.path.userPoolId",
+          },
+          "Type": "HTTP",
+          "Uri": "https://cognito-idp.us-east-1.amazonaws.com/{userPoolId}/.well-known/jwks.json",
+        },
+        "MethodResponses": [
+          {
+            "StatusCode": "200",
+          },
+        ],
+        "RequestParameters": {
+          "method.request.path.userPoolId": true,
+        },
+        "ResourceId": {
+          "Ref": "CognitoPrivateProxyCognitoUserPoolProxyApiuserPoolIdwellknownjwksjsonD9640238",
+        },
+        "RestApiId": {
+          "Ref": "CognitoPrivateProxyCognitoUserPoolProxyApi818B53B1",
+        },
+      },
+      "Type": "AWS::ApiGateway::Method",
+    },
+    "CognitoPrivateProxyCognitoUserPoolProxyApiuserPoolIdwellknownjwksjsonOPTIONSA2D13B9E": {
+      "Properties": {
+        "ApiKeyRequired": false,
+        "AuthorizationType": "NONE",
+        "HttpMethod": "OPTIONS",
+        "Integration": {
+          "IntegrationResponses": [
+            {
+              "ResponseParameters": {
+                "method.response.header.Access-Control-Allow-Headers": "'amz-sdk-invocation-id,amz-sdk-request,cache-control,content-type,x-amz-target,x-amz-user-agent'",
+                "method.response.header.Access-Control-Allow-Methods": "'GET,POST,OPTIONS'",
+                "method.response.header.Access-Control-Allow-Origin": "'*'",
+              },
+              "StatusCode": "204",
+            },
+          ],
+          "RequestTemplates": {
+            "application/json": "{ statusCode: 200 }",
+          },
+          "Type": "MOCK",
+        },
+        "MethodResponses": [
+          {
+            "ResponseParameters": {
+              "method.response.header.Access-Control-Allow-Headers": true,
+              "method.response.header.Access-Control-Allow-Methods": true,
+              "method.response.header.Access-Control-Allow-Origin": true,
+            },
+            "StatusCode": "204",
+          },
+        ],
+        "ResourceId": {
+          "Ref": "CognitoPrivateProxyCognitoUserPoolProxyApiuserPoolIdwellknownjwksjsonD9640238",
+        },
+        "RestApiId": {
+          "Ref": "CognitoPrivateProxyCognitoUserPoolProxyApi818B53B1",
+        },
+      },
+      "Type": "AWS::ApiGateway::Method",
+    },
     "CustomS3AutoDeleteObjectsCustomResourceProviderHandler9D90184F": {
       "DependsOn": [
         "CustomS3AutoDeleteObjectsCustomResourceProviderRole3B1BD092",
@@ -17224,6 +17424,26 @@ exports[`GenerativeAiUseCases matches the snapshot (closed network mode) 6`] = `
             "USER_POOL_ID": {
               "Ref": "AuthUserPool8115E87F",
             },
+            "USER_POOL_PROXY_ENDPOINT": {
+              "Fn::Join": [
+                "",
+                [
+                  "https://",
+                  {
+                    "Fn::ImportValue": "ClosedNetworkStack:ExportsOutputRefCognitoPrivateProxyCognitoUserPoolProxyApi818B53B11EFC652B",
+                  },
+                  ".execute-api.us-east-1.",
+                  {
+                    "Ref": "AWS::URLSuffix",
+                  },
+                  "/",
+                  {
+                    "Fn::ImportValue": "ClosedNetworkStack:ExportsOutputRefCognitoPrivateProxyCognitoUserPoolProxyApiDeploymentStageprod0C2BA9A5793BEF8B",
+                  },
+                  "/",
+                ],
+              ],
+            },
             "VIDEO_GENERATION_MODEL_IDS": {
               "Fn::Join": [
                 "",
@@ -37620,6 +37840,7 @@ exports[`GenerativeAiUseCases matches the snapshot 6`] = `
             "USER_POOL_ID": {
               "Ref": "AuthUserPool8115E87F",
             },
+            "USER_POOL_PROXY_ENDPOINT": "",
             "VIDEO_GENERATION_MODEL_IDS": {
               "Fn::Join": [
                 "",