Testing Cognito secured AWS API Gateways with AWS Custom Resources

Recently I experimented with Cognito secured AWS API Gateways and came across an annoying problem - how can I test my application end to end without having to log in to Cognito with a Browser? Removing the Cognito Authorizer for testing purposes would be one option, but there are cases where applications make use of the tokens for access control or other mechanisms and changing the application behaviour with regards to tokens just for tests doesn't seem be right. Fortunately, CDK supports custom AWS resources that can be used to create dummy users with all the attributes necessary!

Stage 1: Basic Setup

First, we define a Lambda that backs our API. For the purposes of this, the Lambda will just return the AWS Region in which it runs.

producer/main.go

package main

import (
    "context"
    "encoding/json"
    "github.com/aws/aws-lambda-go/events"
    "github.com/aws/aws-lambda-go/lambda"
    "github.com/aws/aws-sdk-go-v2/config"
)

type Handler struct {
    Region string
}

type Response struct {
    Region string `json:"region"`
}

func (h Handler) handleRequest(_ context.Context, _ events.APIGatewayProxyRequest) (*events.APIGatewayProxyResponse, error) {
    response := Response{
        Region: h.Region,
    }
    bytes, err := json.Marshal(response)
    if err != nil {
        return nil, err
    }

    return &events.APIGatewayProxyResponse{
        StatusCode: 200,
        Body:       string(bytes),
    }, nil
}

func main() {
    cfg, err := config.LoadDefaultConfig(context.TODO())
    if err != nil {
        panic(err)
    }

    handler := Handler{
        Region: cfg.Region,
    }
    lambda.Start(handler.handleRequest)
}

The respective CDK constructs then declare the Lambda, an API with a method and an Authorizer that makes use of an imported User Pool.

userPool := awscognito.UserPool_FromUserPoolId(stack, jsii.String("UserPool"), jsii.String(props.UserPoolId))

// ...

producer := awscdklambdagoalpha.NewGoFunction(stack, jsii.String("Producer"), &awscdklambdagoalpha.GoFunctionProps{
    Entry:        jsii.String("./producer"),
    FunctionName: jsii.String("Producer"),
    Runtime:      awslambda.Runtime_PROVIDED_AL2(),
    LogGroup: awslogs.NewLogGroup(stack, jsii.String("ProducerLogGroup"), &awslogs.LogGroupProps{
        LogGroupName:  jsii.String("/aws/lambda/Producer"),
        RemovalPolicy: awscdk.RemovalPolicy_DESTROY,
    }),
})

restApi := awsapigateway.NewRestApi(stack, jsii.String("RestApi"), &awsapigateway.RestApiProps{
    RestApiName: jsii.String("API"),
})
resource := restApi.Root().AddResource(jsii.String("hello"), nil)
resource.AddMethod(jsii.String("GET"), awsapigateway.NewLambdaIntegration(producer, nil), &awsapigateway.MethodOptions{
    Authorizer: awsapigateway.NewCognitoUserPoolsAuthorizer(stack, jsii.String("CognitoAuthorizer"), &awsapigateway.CognitoUserPoolsAuthorizerProps{
        AuthorizerName: jsii.String("CognitoAuthorizer"),
        CognitoUserPools: &[]awscognito.IUserPool{
            userPool,
        },
    }),
})

Now we can curl our API and as expected get a 401 with {"message":"Unauthorized"}. If we "somehow" would get a valid ID token, we could call the API and get the proper result, but depending on the Cognito setup that requires us to perform a few requests with a Browser etc. and we also need a user that depending on our Cognito setup might be federated and not as easy to prepare for tests. We could create test users via AWS CLI that allow us to use the existing login method but without further actions we could still not login headless and instead would still need a Browser to login and we would need to manually make sure the user gets also torn down after our tests.

Stage 2: Dummy User & Headless Login

So, to avoid this extra work we use a different approach: Define a new User Pool Client that allows username and password login and create a dummy user with credentials that allows us to login headless:

userPoolClient := awscognito.NewUserPoolClient(stack, jsii.String("UserPoolClient"), &awscognito.UserPoolClientProps{
    AuthFlows: &awscognito.AuthFlow{
        UserPassword: jsii.Bool(true),
    },
    UserPoolClientName: jsii.String("UserPoolClient"),
    UserPool:           userPool,
})

dummyUser := customresources.NewAwsCustomResource(stack, jsii.String("DummyUserCustomResource"), &customresources.AwsCustomResourceProps{
    OnCreate: &customresources.AwsSdkCall{
        Service: jsii.String("CognitoIdentityServiceProvider"),
        Action:  jsii.String("adminCreateUser"),
        Parameters: map[string]interface{}{
            "UserPoolId": jsii.String(props.UserPoolId),
            "Username":   jsii.String("dummy"),
        },
        PhysicalResourceId: customresources.PhysicalResourceId_Of(jsii.String("DummyUserCustomResource")),
    },
    OnDelete: &customresources.AwsSdkCall{
        Service: jsii.String("CognitoIdentityServiceProvider"),
        Action:  jsii.String("adminDeleteUser"),
        Parameters: map[string]interface{}{
            "UserPoolId": jsii.String(props.UserPoolId),
            "Username":   jsii.String("dummy"),
        },
    },
    Policy: customresources.AwsCustomResourcePolicy_FromSdkCalls(&customresources.SdkCallsPolicyOptions{
        Resources: customresources.AwsCustomResourcePolicy_ANY_RESOURCE(),
    }),
})

dummyUserCredentials := customresources.NewAwsCustomResource(stack, jsii.String("DummyUserCredentialsCustomResource"), &customresources.AwsCustomResourceProps{
    OnCreate: &customresources.AwsSdkCall{
        Service: jsii.String("CognitoIdentityServiceProvider"),
        Action:  jsii.String("adminSetUserPassword"),
        Parameters: map[string]interface{}{
            "UserPoolId": jsii.String(props.UserPoolId),
            "Username":   jsii.String("dummy"),
            "Password":   jsii.String("Password1!"),
            "Permanent":  jsii.Bool(true),
        },
        PhysicalResourceId: customresources.PhysicalResourceId_Of(jsii.String("DummyUserCredentialsCustomResource")),
    },
    Policy: customresources.AwsCustomResourcePolicy_FromSdkCalls(&customresources.SdkCallsPolicyOptions{
        Resources: customresources.AwsCustomResourcePolicy_ANY_RESOURCE(),
    }),
})
dummyUserCredentials.Node().AddDependency(dummyUser)

The additional benefit of the AWS Custom Resources is that we already declared the behaviour for cdk destroy: The deletion of the user.

Stage 3: Successful Request

Now that we have our own User Pool Client and User we can get a token ourselves: aws cognito-idp initiate-auth --auth-flow USER_PASSWORD_AUTH --client-id $CLIENT_ID --auth-parameters USERNAME=dummy,PASSWORD=Password1! and perform a request: and finally get the correct result! Putting the same logic into code is similarly easy:

consumer/main.go

package main

import (
    "context"
    "fmt"
    "github.com/aws/aws-sdk-go-v2/config"
    "github.com/aws/aws-sdk-go-v2/service/cognitoidentityprovider"
    "github.com/aws/aws-sdk-go-v2/service/cognitoidentityprovider/types"
    "github.com/aws/jsii-runtime-go"
    "io"
    "net/http"
    "os"
)

func main() {
    clientId := os.Getenv("CLIENT_ID")
    apiUrl := os.Getenv("API_URL")
    idToken := login(context.TODO(), clientId, "dummy", "Password1!")
    callApi(apiUrl + "/hello", idToken)
}

func login(ctx context.Context, clientId string, username string, password string) string {
    cfg, err := config.LoadDefaultConfig(ctx)
    if err != nil {
        panic(err)
    }

    cognitoClient := cognitoidentityprovider.NewFromConfig(cfg)

    initiateAuthOutput, err := cognitoClient.InitiateAuth(ctx, &cognitoidentityprovider.InitiateAuthInput{
        AuthFlow: types.AuthFlowTypeUserPasswordAuth,
        ClientId: jsii.String(clientId),
        AuthParameters: map[string]string{
            "USERNAME": username,
            "PASSWORD": password,
        },
    })
    if err != nil {
        panic(err)
    }

    return *initiateAuthOutput.AuthenticationResult.IdToken
}

func callApi(url string, idToken string) {
    request, err := http.NewRequest("GET", url, nil)
    if err != nil {
        panic(err)
    }

    request.Header.Add("Authorization", idToken)

    httpClient := &http.Client{}
    response, err := httpClient.Do(request)
    if err != nil {
        panic(err)
    }

    body, err := io.ReadAll(response.Body)
    if err != nil {
        panic(err)
    }

    fmt.Printf("Response: %s\n", body)
}