0

I have an ApiGateway RestApi with a custom authoriser in AWS CDK v2. Now I want to create a WebSocket with an authoriser.

I started by following this guide Stack 3: Api Gateway Websocket API AWS CDK Stack Walk-thru, which has got me as far as creating the ApiGatewayV2 WebSocket. I'm struggling to figure out how to create a custom authoriser for that.

Some of the questions I have:

  • Can I use the same authoriser function with ApiGatewayV2 CfnAuthoriser?
  • Does the authorisation need to happen in some kind of websocket style?
  • How do I authorise with the WebSocket from the front-end application? Is it just an authentication header like in an HTTP request?

I'm having a hard time googling it, keep getting CDK v1 articles. If anyone has some time to point me in the right direction I'd really appreciate it.

Main stack

export class ThingCdkStack extends Stack {
    private authoriserLogicalId: string;

    constructor(scope: Construct, id: string, props: StackProps, private envs: Environment) {
        super(scope, id, props);

        const api = new RestApi(this, 'ThingApi');

        const role = new Role(this, 'ThingRole', {
            roleName: 'thing-role',
            assumedBy: new ServicePrincipal('apigateway.amazonaws.com'),
            inlinePolicies: {
                allowLambdaInvocation: PolicyDocument.fromJson({
                    Version: '2012-10-17',
                    Statement: [
                        {
                            Effect: 'Allow',
                            Action: ['lambda:InvokeFunction', 'lambda:InvokeAsync'],
                            Resource: `arn:aws:lambda:${envs.REGION}:${envs.ACCOUNT}:function:*`,
                        },
                    ],
                }),
            },
        });

        const authorizerHandler = new NodejsFunction(this, 'ThingCustomAuthorizer', {
            entry: 'lambda/handlers/auth/auth0-authoriser.ts',
            runtime: Runtime.NODEJS_18_X,
            environment: {
                AUTH0_ISSUER: 'https://my-auth.eu.auth0.com/',
                AUTH0_AUDIENCE: 'https://my-demo.com',
                REGION: envs.REGION,
                ACCOUNT: envs.ACCOUNT,
            }
        });

        const authorizer = new CfnAuthorizer(this, 'ThingAuthoriser', {
            restApiId: api.restApiId,
            type: 'TOKEN',
            name: 'thing-authoriser',
            identitySource: 'method.request.header.Authorization',
            authorizerUri: `arn:aws:apigateway:${envs.REGION}:lambda:path/2015-03-31/functions/${authorizerHandler.functionArn}/invocations`,
            authorizerCredentials: role.roleArn
        });

        this.authoriserLogicalId = authorizer.logicalId;

        const createThingHandler = new NodejsFunction(this, 'CreateThingLambda', {
            entry: 'lambda/handlers/thing/create-thing.ts',
            runtime: Runtime.NODEJS_18_X,
        });

        this.addAuthMethod('post', api.addResource('thing'), createThingHandler);

        this.addWebsocket(envs, authorizer);
    }

    private addAuthMethod(method: string, resource: Resource, handler: NodejsFunction, integrationOptions?: LambdaIntegrationOptions) {
        const route = resource.addMethod(
            method,
            new LambdaIntegration(handler, integrationOptions),
            {
                authorizationType: AuthorizationType.CUSTOM,
            }
        );
        const childResource = route.node.findChild('Resource');

        (childResource as CfnResource).addPropertyOverride('AuthorizationType', AuthorizationType.CUSTOM);
        (childResource as CfnResource).addPropertyOverride('AuthorizerId', {Ref: this.authoriserLogicalId});
    }

    private addWebsocket(environment: Environment, authorizer: CfnAuthorizer) {
        const connectionsTable = new Table(this, 'ConnectionsTable', {
            partitionKey: {name: 'connectionId', type: AttributeType.STRING},
            readCapacity: 2,
            writeCapacity: 1,
            timeToLiveAttribute: "ttl"
        });

        const commonHandlerProps: NodejsFunctionProps = {
            bundling: {minify: true, sourceMap: true, target: 'es2019'},
            runtime: Runtime.NODEJS_18_X,
            logRetention: RetentionDays.THREE_DAYS
        };

        const connectHandler = new NodejsFunction(this, 'ConnectHandler', {
            ...commonHandlerProps,
            entry: 'lambda/websocket/connect.ts',
            environment: {
                CONNECTIONS_TBL: connectionsTable.tableName
            }
        });

        const defaultHandler = new NodejsFunction(this, 'defaultHandler', {
            ...commonHandlerProps,
            entry: 'lambda/websocket/default.ts',
            environment: {
                CONNECTIONS_TBL: connectionsTable.tableName
            }
        });

        const disconnectHandler = new NodejsFunction(this, 'DisconnectHandler', {
            ...commonHandlerProps,
            entry: 'lambda/websocket/disconnect.ts',
            environment: {
                CONNECTIONS_TBL: connectionsTable.tableName
            }
        });

        const websocketApi = new WebsocketApi(this, "CompletionWebsocketApi", {
            apiName: "completions-api",
            apiDescription: "Web Socket API for Completions",
            stageName: environment.STAGE,
            connectHandler,
            disconnectHandler,
            defaultHandler,
            connectionsTable
        });

        const CONNECTION_URL = `https://${websocketApi.api.ref}.execute-api.${Aws.REGION}.amazonaws.com/${environment.STAGE}`;

        const completionHandler = new NodejsFunction(this, 'CompletionHandler', {
            ...commonHandlerProps,
            entry: 'lambda/websocket/completions.ts',
            environment: {
                CONNECTION_TBL: connectionsTable.tableName,
                CONNECTION_URL: CONNECTION_URL
            },
        });

        websocketApi.addLambdaIntegration(completionHandler, 'completions', 'CompletionsRoute')

        const managementApiPolicyStatement = new PolicyStatement({
            effect: Effect.ALLOW,
            actions: ["execute-api:ManageConnections"],
            resources: [`arn:aws:execute-api:${Aws.REGION}:${Aws.ACCOUNT_ID}:${websocketApi.api.ref}/*`]
        })
        defaultHandler.addToRolePolicy(managementApiPolicyStatement);
        completionHandler.addToRolePolicy(managementApiPolicyStatement);

        new CfnOutput(this, 'WebsocketConnectionUrl', {value: CONNECTION_URL});

        const websocketApiUrl = `${websocketApi.api.attrApiEndpoint}/${environment.STAGE}`
        new CfnOutput(this, "websocketUrl", {
            value: websocketApiUrl
        });
    }
}

WebsocketApi construct

import { ITable } from "aws-cdk-lib/aws-dynamodb";
import { IFunction } from "aws-cdk-lib/aws-lambda";
import { CfnApi, CfnIntegration, CfnRoute, CfnStage, CfnDeployment } from "aws-cdk-lib/aws-apigatewayv2";
import { ServicePrincipal } from "aws-cdk-lib/aws-iam";
import { Aws, Stack } from 'aws-cdk-lib';
import { Construct } from 'constructs';

export interface WebsocketApiProps {
    readonly apiName: string;
    readonly apiDescription: string;
    readonly stageName: string;
    readonly connectHandler: IFunction;
    readonly disconnectHandler: IFunction;
    readonly connectionsTable: ITable;
    readonly defaultHandler?: IFunction;
}

export class WebsocketApi extends Construct {
    readonly props: WebsocketApiProps;
    readonly api: CfnApi;
    readonly deployment: CfnDeployment;

    constructor(parent: Stack, name: string, props: WebsocketApiProps) {
        super(parent, name);
        this.props = props;

        this.api = new CfnApi(this, 'CompletionsWebSocketApi', {
            name: props.apiName,
            description: props.apiDescription,
            protocolType: "WEBSOCKET",
            routeSelectionExpression: "$request.body.action",
        });
        this.deployment = new CfnDeployment(this, "WebsocketDeployment", {
            apiId: this.api.ref,
        });

        new CfnStage(this, "WebsocketStage", {
            stageName: props.stageName,
            apiId: this.api.ref,
            deploymentId: this.deployment.ref,
        });

        props.connectionsTable.grantWriteData(props.connectHandler);
        props.connectionsTable.grantWriteData(props.disconnectHandler);

        this.addLambdaIntegration(props.connectHandler, "$connect", "ConnectionRoute");
        this.addLambdaIntegration(props.disconnectHandler, "$disconnect", "DisconnectRoute");

        if(props.defaultHandler) {
            props.connectionsTable.grantWriteData(props.defaultHandler);
            this.addLambdaIntegration(props.defaultHandler, "$default", "DefaultRoute");
        }
    }

    addLambdaIntegration(handler: IFunction, routeKey: string, operationName: string, apiKeyRequired?: boolean, authorizationType?: string) {
        const integration = new CfnIntegration(this, `${operationName}Integration`, {
            apiId: this.api.ref,
            integrationType: "AWS_PROXY",
            integrationUri: `arn:aws:apigateway:${Aws.REGION}:lambda:path/2015-03-31/functions/${handler.functionArn}/invocations`
        });

        handler.grantInvoke(new ServicePrincipal('apigateway.amazonaws.com', {
            conditions: {
                "ArnLike": {
                    "aws:SourceArn": `arn:aws:execute-api:${Aws.REGION}:${Aws.ACCOUNT_ID}:${this.api.ref}/*/*`
                }
            }
        }));

        this.deployment.addDependency(new CfnRoute(this, `${operationName}Route`, {
            apiId: this.api.ref,
            routeKey,
            apiKeyRequired,
            authorizationType: authorizationType || "NONE",
            operationName,
            target: `integrations/${integration.ref}`
        }));
    }
}

A full working stack of what I have so far can be found here: https://github.com/OrderAndCh4oS/aws-cdk-v2-apigatewayv2-websocket-demo

Update
I've found some useful documentation that I'm working now through: https://docs.aws.amazon.com/apigateway/latest/developerguide/apigateway-websocket-api-lambda-auth.html

Have also found a sample project using Cognito: https://github.com/aws-samples/websocket-api-cognito-auth-sample

OrderAndChaos
  • 3,547
  • 2
  • 30
  • 57
  • This https://dev.to/rahulmlokurte/how-to-validate-requests-to-the-aws-api-gateway-using-cdk-291c ,can do validation.You thought of adding JWT? – Richard Rublev Apr 12 '23 at 08:52
  • 1
    @RichardRublev Yeah, it's the authorisation that I'm after. The custom authoriser handles JWT auth. I have that working with the RestApi, but the WebSockets use CloudFormation templates, and I'm a little lost. – OrderAndChaos Apr 12 '23 at 09:25
  • There is an example in https://github.com/aws-samples/websocket-api-cognito-auth-sample/blob/main/cdk/lib/construct/websocket.ts – quartaela Apr 14 '23 at 09:37

1 Answers1

0

Have made some progress, it's still rough around the edges but is working at least.

I created an ApiGatewayV2 CfnAuthoriser and hooked that up with the deployment.addDependency() on the $connect route, and set the authorizationType to CUSTOM.

The handler is just a regular custom authoriser lambda, but I had to use a querystring parameter to pass the token in.

An authorisation token is passed in on the query string when calling the WebSocket const socket = new WebSocket(`wss://xxxxxxxxxx.execute-api.eu-west-1.amazonaws.com/dev?auth=${token}); `

Main Stack

import * as cdk from 'aws-cdk-lib';
import {Aws, CfnOutput} from 'aws-cdk-lib';
import {Construct} from 'constructs';
import {AttributeType, Table} from "aws-cdk-lib/aws-dynamodb";
import {WebsocketApi} from "./websocket-api";
import {RetentionDays} from "aws-cdk-lib/aws-logs";
import {NodejsFunction, NodejsFunctionProps} from "aws-cdk-lib/aws-lambda-nodejs";
import {Effect, PolicyStatement} from "aws-cdk-lib/aws-iam";
import {Environment} from "../bin/environment";
import {Runtime} from "aws-cdk-lib/aws-lambda";

export class AwsCdkV2WebsocketStack extends cdk.Stack {
    constructor(scope: Construct, id: string, props: cdk.StackProps, private envs: Environment) {
        super(scope, id, props);
        this.addWebsocket(envs);
    }

    private addWebsocket(envs: Environment) {
        const connectionsTable = new Table(this, 'ConnectionsTableWebsocketDemo', {
            partitionKey: {name: 'connectionId', type: AttributeType.STRING},
            readCapacity: 2,
            writeCapacity: 1,
            timeToLiveAttribute: "ttl"
        });

        const commonHandlerProps: NodejsFunctionProps = {
            bundling: {minify: true, sourceMap: true, target: 'es2019'},
            runtime: Runtime.NODEJS_18_X,
            logRetention: RetentionDays.THREE_DAYS
        };

        const connectHandler = new NodejsFunction(this, 'ConnectHandlerWebsocketDemo', {
            ...commonHandlerProps,
            entry: 'lambda/websocket/connect.ts',
            environment: {
                CONNECTIONS_TBL: connectionsTable.tableName
            }
        });

        const authorizationHandler = new NodejsFunction(this, 'AuthorisationHandlerWebsocketDemo', {
            ...commonHandlerProps,
            entry: 'lambda/handlers/authorisation.ts',
            environment: {
                // Todo: use env
                ISSUER: 'https://app-auth.eu.auth0.com/',
                AUDIENCE: 'https://app-demo.com',
            }
        });

        const defaultHandler = new NodejsFunction(this, 'DefaultHandlerWebsocketDemo', {
            ...commonHandlerProps,
            entry: 'lambda/websocket/default.ts',
            environment: {
                CONNECTIONS_TBL: connectionsTable.tableName
            }
        });

        const disconnectHandler = new NodejsFunction(this, 'DisconnectHandlerWebsocketDemo', {
            ...commonHandlerProps,
            entry: 'lambda/websocket/disconnect.ts',
            environment: {
                CONNECTIONS_TBL: connectionsTable.tableName
            }
        });

        const websocketApi = new WebsocketApi(
            this,
            "MessageWebsocketApiWebsocketDemo",
            {
                apiName: "messages-api",
                apiDescription: "Web Socket API for Completions",
                stageName: envs.STAGE,
                connectHandler,
                disconnectHandler,
                defaultHandler,
                connectionsTable,
                authorizationHandler
            },
            envs
        );

        const CONNECTION_URL = `https://${websocketApi.api.ref}.execute-api.${Aws.REGION}.amazonaws.com/${envs.STAGE}`;

        const messageHandler = new NodejsFunction(this, 'CompletionHandlerWebsocketDemo', {
            ...commonHandlerProps,
            entry: 'lambda/websocket/message.ts',
            environment: {
                CONNECTION_TBL: connectionsTable.tableName,
                CONNECTION_URL
            },
        });

        websocketApi.addLambdaIntegration(messageHandler, 'message', 'CompletionsRoute')

        const managementApiPolicyStatement = new PolicyStatement({
            effect: Effect.ALLOW,
            actions: ["execute-api:ManageConnections"],
            resources: [`arn:aws:execute-api:${Aws.REGION}:${Aws.ACCOUNT_ID}:${websocketApi.api.ref}/*`]
        })
        defaultHandler.addToRolePolicy(managementApiPolicyStatement);
        messageHandler.addToRolePolicy(managementApiPolicyStatement);

        new CfnOutput(this, 'WebsocketConnectionUrl', {value: CONNECTION_URL});

        const websocketApiUrl = `${websocketApi.api.attrApiEndpoint}/${envs.STAGE}`
        new CfnOutput(this, "WebsocketUrl", {
            value: websocketApiUrl
        });
    }
}

WebSocket API Stack

import { ITable } from "aws-cdk-lib/aws-dynamodb";
import { IFunction } from "aws-cdk-lib/aws-lambda";
import {CfnApi, CfnIntegration, CfnRoute, CfnStage, CfnDeployment, CfnAuthorizer} from "aws-cdk-lib/aws-apigatewayv2";
import {PolicyDocument, Role, ServicePrincipal} from "aws-cdk-lib/aws-iam";
import { Aws, Stack } from 'aws-cdk-lib';
import { Construct } from 'constructs';
import {Environment} from "../bin/environment";

export interface WebsocketApiProps {
    readonly apiName: string;
    readonly apiDescription: string;
    readonly stageName: string;
    readonly connectHandler: IFunction;
    readonly disconnectHandler: IFunction;
    readonly connectionsTable: ITable;
    readonly authorizationHandler: IFunction;
    readonly defaultHandler?: IFunction;
}

export class WebsocketApi extends Construct {
    readonly props: WebsocketApiProps;
    readonly api: CfnApi;
    readonly deployment: CfnDeployment;

    constructor(parent: Stack, name: string, props: WebsocketApiProps, envs: Environment) {
        super(parent, name);
        this.props = props;

        this.api = new CfnApi(this, 'CompletionsWebSocketApi', {
            name: props.apiName,
            description: props.apiDescription,
            protocolType: "WEBSOCKET",
            routeSelectionExpression: "$request.body.action"
        });

        this.deployment = new CfnDeployment(this, "WebsocketDeployment", {
            apiId: this.api.ref,
        });

        new CfnStage(this, "WebsocketStage", {
            stageName: props.stageName,
            apiId: this.api.ref,
            deploymentId: this.deployment.ref,
        });

        const role = new Role(this, 'AuthorisedRole', {
            roleName: 'authorised-role',
            assumedBy: new ServicePrincipal('apigateway.amazonaws.com'),
            inlinePolicies: {
                allowLambdaInvocation: PolicyDocument.fromJson({
                    Version: '2012-10-17',
                    Statement: [
                        {
                            Effect: 'Allow',
                            Action: ['lambda:InvokeFunction', 'lambda:InvokeAsync'],
                            Resource: `arn:aws:lambda:${envs.REGION}:${envs.ACCOUNT}:function:*`,
                        },
                    ],
                }),
            },
        });

        const authorizer = new CfnAuthorizer(this, 'WorkspaceAuthoriser', {
            name: 'workspace-authoriser',
            apiId: this.api.ref,
            authorizerType: 'REQUEST',
            identitySource: ['route.request.querystring.auth'],
            authorizerUri: `arn:aws:apigateway:${envs.REGION}:lambda:path/2015-03-31/functions/${this.props.authorizationHandler.functionArn}/invocations`,
            authorizerCredentialsArn: role.roleArn,
        });
        this.addLambdaIntegration(props.connectHandler, "$connect", "ConnectionRoute", false,"CUSTOM",  authorizer);

        props.connectionsTable.grantWriteData(props.connectHandler);
        props.connectionsTable.grantWriteData(props.disconnectHandler);

        this.addLambdaIntegration(props.disconnectHandler, "$disconnect", "DisconnectRoute");

        if(props.defaultHandler) {
            props.connectionsTable.grantWriteData(props.defaultHandler);
            this.addLambdaIntegration(props.defaultHandler, "$default", "DefaultRoute");
        }
    }

    addLambdaIntegration(handler: IFunction, routeKey: string, operationName: string, apiKeyRequired?: boolean, authorizationType?: string, authorizer?: CfnAuthorizer) {
        const integration = new CfnIntegration(this, `${operationName}Integration`, {
            apiId: this.api.ref,
            integrationType: "AWS_PROXY",
            integrationUri: `arn:aws:apigateway:${Aws.REGION}:lambda:path/2015-03-31/functions/${handler.functionArn}/invocations`,
        });

        handler.grantInvoke(new ServicePrincipal('apigateway.amazonaws.com', {
            conditions: {
                "ArnLike": {
                    "aws:SourceArn": `arn:aws:execute-api:${Aws.REGION}:${Aws.ACCOUNT_ID}:${this.api.ref}/*/*`
                }
            }
        }));

        this.deployment.addDependency(new CfnRoute(this, `${operationName}RouteWebsocketDemo`, {
            apiId: this.api.ref,
            routeKey,
            apiKeyRequired,
            authorizationType: authorizationType || "NONE",
            operationName,
            target: `integrations/${integration.ref}`,
            authorizerId: authorizer?.attrAuthorizerId,
        }));
    }
}

Authorisation Handler

if (!process.env.AUDIENCE) throw new Error('Missing AUDIENCE');
if (!process.env.ISSUER) throw new Error('Missing ISSUER');
if (!process.env.AWS_REGION) throw new Error('Missing AWS_REGION');

const JWKS_URI = `${process.env.ISSUER}.well-known/jwks.json`

export const handler = async (event: any, context: any, callback: any) => {
    let data;
    try {
        data = await authenticate(event);
    } catch (err) {
        console.log('UNAUTHORISED', err);
        return context.fail('Unauthorized');
    }

    console.log('AUTHORISED', data);
    return data;
};

const getPolicyDocument = (effect: any, resource: any) => {
    return {
        Version: '2012-10-17',
        Statement: [{
            Action: 'execute-api:Invoke',
            Effect: effect,
            Resource: resource,
        }]
    };
}

const getToken = (event: any) => {
    if (!event.type || event.type !== 'REQUEST') {
        throw new Error('Expected "event.type" parameter to have value "REQUEST"');
    }

    const tokenString = event.queryStringParameters?.auth;
    if (!tokenString) {
        throw new Error('Expected "event.queryStringParameters.auth" parameter to be set');
    }

    return tokenString;
}

const jwtOptions = {
    audience: process.env.AUDIENCE,
    issuer: process.env.ISSUER
};

const client = jwksClient({
    cache: true,
    rateLimit: true,
    jwksRequestsPerMinute: 10,
    jwksUri: JWKS_URI
});

const authenticate = (event: any) => {
    console.log(event);
    const token = getToken(event);

    const decoded = jwt.decode(token, {complete: true});
    if (!decoded || !decoded.header || !decoded.header.kid) {
        throw new Error('invalid token');
    }

    const getSigningKey = util.promisify(client.getSigningKey);
    return getSigningKey(decoded.header.kid)
        .then((key: any) => {
            const signingKey = key?.publicKey || key?.rsaPublicKey;
            return jwt.verify(token, signingKey, jwtOptions);
        })
        .then((decoded: any) => ({
            principalId: decoded.sub,
            policyDocument: getPolicyDocument('Allow', '*'),
            context: {scope: decoded.scope}
        }));
}

I made a scratch demo project for this which may or may not be useful to someone: https://github.com/OrderAndCh4oS/aws-cdk-v2-apigatewayv2-websocket-demo

Will leave the question open in case someone is able to improve on this, it certainly needs work still.

OrderAndChaos
  • 3,547
  • 2
  • 30
  • 57