Breaking: [AEA-0000] - An attempt at making a single lambda work on both real and mock auth (#241)

## Summary

- ❗ Breaking Change

### Details

This is my attempt at making the lambda able to receive both real and
mock auth tokens.

There's an environment variable that we can set to enable mock auth, and
if it's present AND the user is using what appears to be a mock token,
then we process it as such. Otherwise, we fall back on proper
authentication.

---------

Co-authored-by: jonathanwelch1-nhs <[email protected]>

Author: wildjames

packages/auth_demo/src/App.tsx

@@ -10,7 +10,6 @@ import { authConfig } from './configureAmplify'
 Amplify.configure(authConfig, { ssr: true })
 
 const trackerUserInfoEndpoint = "/api/tracker-user-info"
-const mockTrackerUserInfoEndpoint = "/api/mock-tracker-user-info"
 
 function App() {
   const [user, setUser] = useState(null)
@@ -66,14 +65,13 @@ function App() {
     }
   }
 
-  const fetchTrackerUserInfo = async (isMock: boolean) => {
+  const fetchTrackerUserInfo = async () => {
     setLoading(true)
     setTrackerUserInfoData(null)
     setError(null)
 
-    let endpoint = isMock ? mockTrackerUserInfoEndpoint : trackerUserInfoEndpoint
     try {
-      const response = await axios.get(endpoint, {
+      const response = await axios.get(trackerUserInfoEndpoint, {
         headers: {
           Authorization: `Bearer ${idToken}`,
           'NHSD-Session-URID': '555254242106'
@@ -108,20 +106,12 @@ function App() {
 
       <div style={{ marginTop: '20px' }}>
         <button
-          onClick={() => fetchTrackerUserInfo(false)}
+          onClick={fetchTrackerUserInfo}
           disabled={!isSignedIn}
         >
           Fetch Tracker User Info
         </button>
       </div>
-      <div style={{ marginTop: '20px' }}>
-        <button
-          onClick={() => fetchTrackerUserInfo(true)}
-          disabled={!isSignedIn}
-        >
-          Fetch Mock Tracker User Info
-        </button>
-      </div>
 
       {loading && <p>Loading...</p>}
       {trackerUserInfoData && (

packages/cdk/resources/RestApiGatewayMethods.ts

@@ -18,7 +18,6 @@ export interface RestApiGatewayMethodsProps {
   readonly tokenLambda: NodejsFunction
   readonly mockTokenLambda: NodejsFunction
   readonly trackerUserInfoLambda: LambdaFunction
-  readonly mockTrackerUserInfoLambda: LambdaFunction
   readonly useMockOidc: boolean
   readonly authorizer: CognitoUserPoolsAuthorizer
 
@@ -61,17 +60,6 @@ export class RestApiGatewayMethods extends Construct{
       authorizer: props.authorizer
     })
 
-    // mock tracker-user-info endpoint
-    if (props.useMockOidc) {
-      const mockTrackerUserInfoLambdaResource = props.restApiGateway.root.addResource("mock-tracker-user-info")
-      mockTrackerUserInfoLambdaResource.addMethod("GET", new LambdaIntegration(props.mockTrackerUserInfoLambda.lambda, {
-        credentialsRole: props.restAPiGatewayRole
-      }), {
-        authorizationType: AuthorizationType.COGNITO,
-        authorizer: props.authorizer
-      })
-    }
-
     /* Dummy Method/Resource to test cognito auth */
 
     const mockNoAuth = new MockIntegration({

packages/cdk/resources/api/apiFunctions.ts

@@ -1,6 +1,4 @@
-
 import {Construct} from "constructs"
-
 import {LambdaFunction} from "../LambdaFunction"
 import {ITableV2} from "aws-cdk-lib/aws-dynamodb"
 import {
@@ -39,7 +37,6 @@ export interface ApiFunctionsProps {
   readonly mockPoolIdentityProviderName: string
   readonly logRetentionInDays: number
   readonly deploymentRole: IRole
-
 }
 
 /**
@@ -48,7 +45,6 @@ export interface ApiFunctionsProps {
 export class ApiFunctions extends Construct {
   public readonly apiFunctionsPolicies: Array<IManagedPolicy>
   public readonly trackerUserInfoLambda: LambdaFunction
-  public readonly mockTrackerUserInfoLambda: LambdaFunction
   public readonly primaryJwtPrivateKey: Secret
 
   public constructor(scope: Construct, id: string, props: ApiFunctionsProps) {
@@ -107,7 +103,7 @@ export class ApiFunctions extends Construct {
       ]
     })
 
-    // Secret used by lambdas that holds the JWT private key
+    // Real JWT Secret
     const primaryJwtPrivateKey = new Secret(this, "PrimaryJwtPrivateKey", {
       secretName: `${props.stackName}-primaryJwtPrivateKeyTrkUsrNfo`,
       secretStringValue: SecretValue.unsafePlainText("ChangeMe"),
@@ -136,9 +132,9 @@ export class ApiFunctions extends Construct {
       ]
     })
 
-    // Set up things for mock authorization
-    let getMockJWTPrivateKeySecret: ManagedPolicy | undefined = undefined
+    // Setup mock keys if needed
     let mockJwtPrivateKey: Secret | undefined = undefined
+    let getMockJWTPrivateKeySecret: ManagedPolicy | undefined = undefined
 
     if (props.useMockOidc) {
       if (
@@ -151,7 +147,6 @@ export class ApiFunctions extends Construct {
         throw new Error("Attempt to use mock oidc but variables are not defined")
       }
 
-      // Secret used by mock user info lambda that holds the JWT private key
       mockJwtPrivateKey = new Secret(this, "MockJwtPrivateKey", {
         secretName: `${props.stackName}-mockJwtPrivateKeyTrkUsrNfo`,
         secretStringValue: SecretValue.unsafePlainText("ChangeMe"),
@@ -173,123 +168,63 @@ export class ApiFunctions extends Construct {
       })
     }
 
-    // *-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-* //
-    //           Lambda setups           //
-    // *-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-* //
+    // Combine policies
+    const additionalPolicies = [
+      props.tokenMappingTableWritePolicy,
+      props.tokenMappingTableReadPolicy,
+      props.useTokensMappingKmsKeyPolicy,
+      useJwtKmsKeyPolicy,
+      getJWTPrivateKeySecret
+    ]
 
-    // CPT USER INFORMATION LAMBDA //
+    if (props.useMockOidc && getMockJWTPrivateKeySecret) {
+      additionalPolicies.push(getMockJWTPrivateKeySecret)
+    }
 
-    // Set up things for mock authorization
-    let mockTrackerUserInfoLambda: LambdaFunction
-    if (props.useMockOidc) {
-      // These checks are technically unnecessary, but they are here to prevent the linter from complaining
-      if (
-        props.mockOidcjwksEndpoint === undefined ||
-        props.mockOidcUserInfoEndpoint === undefined ||
-        props.mockOidcTokenEndpoint === undefined ||
-        props.mockOidcClientId === undefined ||
-        props.mockOidcIssuer === undefined ||
-        getMockJWTPrivateKeySecret === undefined ||
-        mockJwtPrivateKey === undefined
-      ) {
-        throw new Error("Attempt to use mock oidc but variables are not defined")
-      }
+    // Environment variables
+    // We pass in both sets of endpoints and keys. The Lambda code determines at runtime which to use.
+    const lambdaEnv: { [key: string]: string } = {
+      TokenMappingTableName: props.tokenMappingTable.tableName,
+
+      // Real endpoints/credentials
+      REAL_IDP_TOKEN_PATH: props.primaryOidcTokenEndpoint,
+      REAL_USER_INFO_ENDPOINT: props.primaryOidcUserInfoEndpoint,
+      REAL_OIDCJWKS_ENDPOINT: props.primaryOidcjwksEndpoint,
+      REAL_JWT_PRIVATE_KEY_ARN: primaryJwtPrivateKey.secretArn,
+      REAL_USER_POOL_IDP: props.primaryPoolIdentityProviderName,
+      REAL_USE_SIGNED_JWT: "true",
+      REAL_OIDC_CLIENT_ID: props.primaryOidcClientId,
+      REAL_OIDC_ISSUER: props.primaryOidcIssuer,
+
+      // Indicate if mock mode is available
+      MOCK_MODE_ENABLED: props.useMockOidc ? "true" : "false"
+    }
 
-      mockTrackerUserInfoLambda = new LambdaFunction(this, "MockTrackerUserInfo", {
-        serviceName: props.serviceName,
-        stackName: props.stackName,
-        lambdaName: `${props.stackName}-mockTrkUsrNfo`,
-        additionalPolicies: [
-          props.tokenMappingTableWritePolicy,
-          props.tokenMappingTableReadPolicy,
-          props.useTokensMappingKmsKeyPolicy,
-          useJwtKmsKeyPolicy,
-          getMockJWTPrivateKeySecret
-        ],
-        logRetentionInDays: props.logRetentionInDays,
-        packageBasePath: "packages/trackerUserInfoLambda",
-        entryPoint: "src/handler.ts",
-        lambdaEnvironmentVariables: {
-          idpTokenPath: props.mockOidcTokenEndpoint,
-          TokenMappingTableName: props.tokenMappingTable.tableName,
-          UserPoolIdentityProvider: props.mockPoolIdentityProviderName,
-          oidcjwksEndpoint: props.mockOidcjwksEndpoint,
-          jwtPrivateKeyArn: mockJwtPrivateKey.secretArn,
-          userInfoEndpoint: props.mockOidcUserInfoEndpoint,
-          useSignedJWT: "false",
-          oidcClientId: props.mockOidcClientId,
-          oidcIssuer: props.mockOidcIssuer
-        }
-      })
+    // If mock OIDC is enabled, add mock environment variables
+    if (props.useMockOidc && mockJwtPrivateKey) {
+      lambdaEnv["MOCK_IDP_TOKEN_PATH"] = props.mockOidcTokenEndpoint!
+      lambdaEnv["MOCK_USER_INFO_ENDPOINT"] = props.mockOidcUserInfoEndpoint!
+      lambdaEnv["MOCK_OIDCJWKS_ENDPOINT"] = props.mockOidcjwksEndpoint!
+      lambdaEnv["MOCK_JWT_PRIVATE_KEY_ARN"] = mockJwtPrivateKey.secretArn
+      lambdaEnv["MOCK_USER_POOL_IDP"] = props.mockPoolIdentityProviderName
+      lambdaEnv["MOCK_USE_SIGNED_JWT"] = "false"
+      lambdaEnv["MOCK_OIDC_CLIENT_ID"] = props.mockOidcClientId!
+      lambdaEnv["MOCK_OIDC_ISSUER"] = props.mockOidcIssuer!
     }
 
-    // Proper lambda function for user info
+    // Single Lambda for both real and mock scenarios
     const trackerUserInfoLambda = new LambdaFunction(this, "TrackerUserInfo", {
       serviceName: props.serviceName,
       stackName: props.stackName,
-      lambdaName: `${props.stackName}-TrkUsrNfo`,
-      additionalPolicies: [
-        props.tokenMappingTableWritePolicy,
-        props.tokenMappingTableReadPolicy,
-        props.useTokensMappingKmsKeyPolicy,
-        useJwtKmsKeyPolicy,
-        getJWTPrivateKeySecret
-      ],
+      lambdaName: `${props.stackName}-TrkUsrNfoUnified`,
+      additionalPolicies: additionalPolicies,
       logRetentionInDays: props.logRetentionInDays,
       packageBasePath: "packages/trackerUserInfoLambda",
       entryPoint: "src/handler.ts",
-      lambdaEnvironmentVariables: {
-        idpTokenPath: props.primaryOidcTokenEndpoint,
-        TokenMappingTableName: props.tokenMappingTable.tableName,
-        UserPoolIdentityProvider: props.primaryPoolIdentityProviderName,
-        oidcjwksEndpoint: props.primaryOidcjwksEndpoint,
-        jwtPrivateKeyArn: primaryJwtPrivateKey.secretArn,
-        userInfoEndpoint: props.primaryOidcUserInfoEndpoint,
-        useSignedJWT: "true",
-        oidcClientId: props.primaryOidcClientId,
-        oidcIssuer: props.primaryOidcIssuer
-      }
+      lambdaEnvironmentVariables: lambdaEnv
     })
 
-    // PRESCRIPTION SEARCH LAMBDA //
-
-    // const prescriptionSearchLambda = new LambdaFunction(this, "PrescriptionSearch", {
-    //   runtime: Runtime.NODEJS_20_X,
-    //   serviceName: props.serviceName,
-    //   stackName: props.stackName,
-    //   lambdaName: `${props.stackName}-prescSearch`,
-    //   additionalPolicies: [
-    //     props.tokenMappingTableWritePolicy,
-    //     props.tokenMappingTableReadPolicy,
-    //     props.useTokensMappingKmsKeyPolicy,
-    //     props.sharedSecrets.useJwtKmsKeyPolicy,
-    //     props.sharedSecrets.getPrimaryJwtPrivateKeyPolicy
-    //   ],
-    //   logRetentionInDays: props.logRetentionInDays,
-    //   packageBasePath: "packages/prescriptionSearchLambda",
-    //   entryPoint: "src/handler.ts",
-    //   lambdaEnvironmentVariables: {
-    //     idpTokenPath: props.primaryOidcTokenEndpoint,
-    //     TokenMappingTableName: props.tokenMappingTable.tableName,
-    //     UserPoolIdentityProvider: props.primaryPoolIdentityProviderName,
-    //     oidcjwksEndpoint: props.primaryOidcjwksEndpoint,
-    //     jwtPrivateKeyArn: props.sharedSecrets.primaryJwtPrivateKey.secretArn,
-    //     userInfoEndpoint: props.primaryOidcUserInfoEndpoint,
-    //     useSignedJWT: "true",
-    //     oidcClientId: props.primaryOidcClientId,
-    //     oidcIssuer: props.primaryOidcIssuer,
-    //     apigeeApiKey: props.apigeeApiKey,
-    //     jwtKid: props.jwtKid
-    //   }
-    // })
-
-    // // Suppress the AwsSolutions-L1 rule for the prescription search Lambda function
-    // NagSuppressions.addResourceSuppressions(prescriptionSearchLambda.lambda, [
-    //   {
-    //     id: "AwsSolutions-L1",
-    //     reason: "The Lambda function uses the latest runtime version supported at the time of implementation."
-    //   }
-    // ])
+    // This one Lambda can handle both mock and real requests at runtime
 
     // *-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-* //
     //            Permissions            //
@@ -300,9 +235,6 @@ export class ApiFunctions extends Construct {
       trackerUserInfoLambda.executeLambdaManagedPolicy
       // prescriptionSearchLambda.executeLambdaManagedPolicy
     ]
-    if (props.useMockOidc) {
-      apiFunctionsPolicies.push(mockTrackerUserInfoLambda!.executeLambdaManagedPolicy)
-    }
 
     // *-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-* //
     //              Outputs              //
@@ -314,9 +246,5 @@ export class ApiFunctions extends Construct {
 
     // CPT user info lambda outputs
     this.trackerUserInfoLambda = trackerUserInfoLambda
-    this.mockTrackerUserInfoLambda = mockTrackerUserInfoLambda!
-
-    // // Prescription search lambda outputs
-    // this.prescriptionSearchLambda = prescriptionSearchLambda.lambda
   }
 }

packages/cdk/stacks/StatelessResourcesStack.ts

@@ -187,7 +187,6 @@ export class StatelessResourcesStack extends Stack {
       tokenLambda: cognitoFunctions.tokenLambda,
       mockTokenLambda: cognitoFunctions.mockTokenLambda,
       trackerUserInfoLambda: apiFunctions.trackerUserInfoLambda,
-      mockTrackerUserInfoLambda: apiFunctions.mockTrackerUserInfoLambda,
       useMockOidc: useMockOidc,
       authorizer: apiGateway.authorizer
     })

packages/cpt-ui/components/EpsFooter.tsx

@@ -12,7 +12,7 @@ export default function EpsFooter() {
 
     useEffect(() => {
         console.log("Viewing site version of commit ID:", COMMIT_ID)
-    }, [COMMIT_ID])
+    }, [])
 
     return (
         <Footer id="eps_footer" className="eps_footer">

packages/trackerUserInfoLambda/jest.config.ts

@@ -4,7 +4,8 @@ import type {JestConfigWithTsJest} from "ts-jest"
 const jestConfig: JestConfigWithTsJest = {
   ...defaultConfig,
   "rootDir": "./",
-  setupFiles: ["<rootDir>/.jest/setEnvVars.js"]
+  setupFiles: ["<rootDir>/.jest/setEnvVars.js"],
+  moduleNameMapper: {"@/(.*)$": ["<rootDir>/src/$1"]}
 }
 
 export default jestConfig