Add apiKey authentication

.gitignore

@@ -1,3 +1,4 @@
+example/**/*.js
 lib/**/*.js
 test/**/*.js
 !jest.config.js

README.md

@@ -82,6 +82,23 @@ new EventBridgeWebSocket(this, 'sockets', {
 });
 ```
 
+##### Secure api with apikey
+
+```typescript
+new EventBridgeWebSocket(this, 'sockets', {
+  bus: 'your-event-bus-name',
+
+  // Listens for all UserCreated events
+  eventPattern: {
+    detailType: ['UserCreated'],
+  },
+  stage: 'dev',
+  authentication: true,
+});
+```
+
+This will create an aws secret in the secretsmanager with the api key used for authentication. The apikey must be added as query param to the api endpoint url `wss://<apiId>.execute-api.<region>.amazonaws.com/<stage>?apiKey=<valueFromSecret>`
+
 You can find more [here on the AWS documentation](https://docs.aws.amazon.com/cdk/api/latest/docs/@aws-cdk_aws-events.EventPattern.html)
 
 # Contributing

example/event-publisher.ts

@@ -0,0 +1,17 @@
+import { EventBridge } from 'aws-sdk';
+import { randomUUID } from 'crypto';
+
+var events = new EventBridge();
+export const handler = async () => {
+  var params = {
+    Entries: [
+      {
+        Detail: JSON.stringify({ timestamp: new Date().toLocaleDateString(), payload: randomUUID() }),
+        DetailType: 'ExampleEvent',
+        EventBusName: process.env.EVENT_BUS_NAME,
+        Source: 'example.event.source',
+      },
+    ],
+  };
+  await events.putEvents(params).promise();
+};

example/example-app.ts

@@ -0,0 +1,36 @@
+import { App, Duration, Stack } from 'aws-cdk-lib';
+import { EventBus, Rule, Schedule } from 'aws-cdk-lib/aws-events';
+import { AwsApi, LambdaFunction } from 'aws-cdk-lib/aws-events-targets';
+import { Runtime } from 'aws-cdk-lib/aws-lambda';
+import { NodejsFunction } from 'aws-cdk-lib/aws-lambda-nodejs';
+import { EventBridgeWebSocket } from '../lib';
+
+const app = new App();
+
+const stack = new Stack(app, 'example-stack');
+
+const eventbus = new EventBus(stack, 'EventBus');
+
+const eventPublisher = new NodejsFunction(stack, 'EventPublisher', {
+  entry: 'example/event-publisher.ts',
+  runtime: Runtime.NODEJS_16_X,
+  environment: {
+    EVENT_BUS_NAME: eventbus.eventBusName,
+  },
+});
+eventbus.grantPutEventsTo(eventPublisher);
+
+new Rule(stack, 'Rule', {
+  schedule: Schedule.rate(Duration.minutes(1)),
+  targets: [new LambdaFunction(eventPublisher)],
+});
+
+new EventBridgeWebSocket(stack, 'EventBridgeWebSocket', {
+  bus: eventbus.eventBusName,
+
+  eventPattern: {
+    source: ['example.event.source'],
+  },
+  stage: 'test',
+  authentication: true,
+});

lib/eventbridge-sockets/eventbridge-sockets-contruct.ts

@@ -4,17 +4,21 @@ import { Runtime } from 'aws-cdk-lib/aws-lambda';
 import { EventBus, EventPattern, Rule } from 'aws-cdk-lib/aws-events';
 import { Effect, PolicyStatement } from 'aws-cdk-lib/aws-iam';
 import { Construct } from 'constructs';
-import { WebSocketApi, WebSocketStage } from '@aws-cdk/aws-apigatewayv2-alpha';
+import { WebSocketApi, WebSocketAuthorizer, WebSocketAuthorizerType, WebSocketStage } from '@aws-cdk/aws-apigatewayv2-alpha';
 import { WebSocketLambdaIntegration } from '@aws-cdk/aws-apigatewayv2-integrations-alpha';
 import { LambdaFunction } from 'aws-cdk-lib/aws-events-targets';
 import { NodejsFunction } from 'aws-cdk-lib/aws-lambda-nodejs';
+import { Authorizer } from 'aws-cdk-lib/aws-apigateway';
+import { WebSocketLambdaAuthorizer } from '@aws-cdk/aws-apigatewayv2-authorizers-alpha';
+import { Secret } from 'aws-cdk-lib/aws-secretsmanager';
 
 const path = require('path');
 
 export interface EventBridgeWebSocketProps {
   readonly bus: string;
   readonly stage?: string;
   readonly eventPattern?: EventPattern;
+  readonly authentication?: boolean;
 }
 
 export class EventBridgeWebSocket extends Construct {
@@ -71,9 +75,13 @@ export class EventBridgeWebSocket extends Construct {
     table.grantReadWriteData(disconnectFunc);
     table.grantReadWriteData(eventBridgeBrokerFunc);
 
+    // authorizer
+    const authorizer: WebSocketLambdaAuthorizer | undefined = config.authentication ? this.addAuthorization() : undefined;
+
     // create routes for API Gateway
     api.addRoute('$connect', {
       integration: new WebSocketLambdaIntegration('ConnectIntegration', connectFunc),
+      authorizer,
     });
     api.addRoute('$disconnect', {
       integration: new WebSocketLambdaIntegration('DisconnectIntegration', disconnectFunc),
@@ -100,6 +108,24 @@ export class EventBridgeWebSocket extends Construct {
     });
   }
 
+  private addAuthorization() {
+    const apiKeySecret = new Secret(this, 'apiKeySecret', { removalPolicy: RemovalPolicy.DESTROY });
+    const authorizerLambda = new NodejsFunction(this, 'authorizer', {
+      entry: path.join(__dirname, '../lambda-fns/authorizer/authorizer.ts'),
+      runtime: Runtime.NODEJS_16_X,
+      memorySize: 256,
+      environment: {
+        API_KEY_SECRET: apiKeySecret.secretName,
+      },
+    });
+    apiKeySecret.grantRead(authorizerLambda);
+    const authorizer: WebSocketLambdaAuthorizer = new WebSocketLambdaAuthorizer('authorizer', authorizerLambda, {
+      identitySource: ['route.request.querystring.apiKey'],
+    });
+    Stack.of(this).exportValue(apiKeySecret.secretName, { name: 'ApiKeySecretName' });
+    return authorizer;
+  }
+
   private createFunction(name: string, tableName: string, options: any = {}) {
     return new NodejsFunction(this, name, {
       entry: path.join(__dirname, `../lambda-fns/${name}/${name}.ts`),

lib/lambda-fns/authorizer/authorizer.ts

@@ -0,0 +1,39 @@
+import { getSecretValue } from './secretsmanager';
+
+type ApigatewayWebsocketLambdaAuthorizerEvent = {
+  headers: Record<string, string>;
+  queryStringParameters: Record<string, string>;
+  methodArn: string;
+};
+export const handler = async ({ methodArn, queryStringParameters }: ApigatewayWebsocketLambdaAuthorizerEvent) => {
+  // Retrieve request parameters from the Lambda function input:
+
+  if (!queryStringParameters || queryStringParameters == null) {
+    console.log('No queryStringParameters found');
+    return 'Unauthorized';
+  }
+
+  const apiKey = await getSecretValue();
+
+  if (queryStringParameters['apiKey'] !== apiKey) {
+    console.log("API Key doesn't match");
+    return 'Unauthorized';
+  }
+  const apigatewayAuthorizerAllowPolicy = {
+    Version: '2012-10-17',
+    Statement: [
+      {
+        Action: 'execute-api:Invoke',
+        Effect: 'Allow',
+        Resource: methodArn,
+      },
+    ],
+  };
+  console.log('Access granted');
+  const resp = {
+    principalId: 'authenticated-user',
+    policyDocument: apigatewayAuthorizerAllowPolicy,
+  };
+  console.log(JSON.stringify(resp));
+  return resp;
+};

lib/lambda-fns/authorizer/secretsmanager.ts

@@ -0,0 +1,11 @@
+import { SecretsManager } from 'aws-sdk';
+
+const secretClient = new SecretsManager();
+
+export const getSecretValue = async (): Promise<string> => {
+  const secretString = await (await secretClient.getSecretValue({ SecretId: process.env.API_KEY_SECRET! }).promise()).SecretString;
+  if (!secretString) {
+    throw new Error(`Secret string not found in secret ${process.env.API_KEY_SECRET}`);
+  }
+  return secretString;
+};

package.json

@@ -11,6 +11,7 @@
     "build": "tsc",
     "watch": "tsc -w",
     "test": "jest",
+    "cdk-example": "cdk --app 'ts-node example/example-app.ts'",
     "bump:lib": "npm --no-git-tag-version --allow-same-version version $VERSION"
   },
   "devDependencies": {
@@ -27,6 +28,7 @@
   },
   "dependencies": {
     "@aws-cdk/aws-apigatewayv2-alpha": "^2.45.0-alpha.0",
+    "@aws-cdk/aws-apigatewayv2-authorizers-alpha": "^2.45.0-alpha.0",
     "@aws-cdk/aws-apigatewayv2-integrations-alpha": "^2.45.0-alpha.0",
     "aws-cdk-lib": "^2.45.0",
     "constructs": "^10.1.127",
@@ -45,4 +47,4 @@
     "type": "git",
     "url": "https://github.com/boyney123/cdk-eventbridge-socket.git"
   }
-}
+}

test/__snapshots__/cdk-eventbridge-socket.test.ts.snap

@@ -270,16 +270,13 @@ exports[`EventBridgeWebSocket snapshot test EventsRuleToSns default params 1`] =
             "AttributeType": "S",
           },
         ],
+        "BillingMode": "PAY_PER_REQUEST",
         "KeySchema": [
           {
             "AttributeName": "connectionId",
             "KeyType": "HASH",
           },
         ],
-        "ProvisionedThroughput": {
-          "ReadCapacityUnits": 5,
-          "WriteCapacityUnits": 5,
-        },
         "TableName": "eventBridgeSocketDeploy-connections-table",
       },
       "Type": "AWS::DynamoDB::Table",
@@ -295,7 +292,7 @@ exports[`EventBridgeWebSocket snapshot test EventsRuleToSns default params 1`] =
           "S3Bucket": {
             "Fn::Sub": "cdk-hnb659fds-assets-\${AWS::AccountId}-\${AWS::Region}",
           },
-          "S3Key": "3e0a0f754714a39c50b8d4b23a69e2ffea1ff6f186f10c257f57b04f17740569.zip",
+          "S3Key": "a549b85065ed6603d1aac6495ead30df29d29dab7ce01adb9c4f9d7c376fe472.zip",
         },
         "Environment": {
           "Variables": {
@@ -325,7 +322,7 @@ exports[`EventBridgeWebSocket snapshot test EventsRuleToSns default params 1`] =
             "Arn",
           ],
         },
-        "Runtime": "nodejs12.x",
+        "Runtime": "nodejs16.x",
         "Timeout": 300,
       },
       "Type": "AWS::Lambda::Function",
@@ -455,7 +452,7 @@ exports[`EventBridgeWebSocket snapshot test EventsRuleToSns default params 1`] =
             "Arn",
           ],
         },
-        "Runtime": "nodejs12.x",
+        "Runtime": "nodejs16.x",
         "Timeout": 300,
       },
       "Type": "AWS::Lambda::Function",
@@ -561,7 +558,7 @@ exports[`EventBridgeWebSocket snapshot test EventsRuleToSns default params 1`] =
             "Arn",
           ],
         },
-        "Runtime": "nodejs12.x",
+        "Runtime": "nodejs16.x",
         "Timeout": 300,
       },
       "Type": "AWS::Lambda::Function",

test/lambdas/authorizer.test.ts

@@ -0,0 +1,53 @@
+import { handler } from '../../lib/lambda-fns/authorizer/authorizer';
+import * as secretsmanager from '../../lib/lambda-fns/authorizer/secretsmanager';
+
+describe('authorizer', () => {
+  it('return allow policy for method arn if apiKey match', async () => {
+    jest.spyOn(secretsmanager, 'getSecretValue').mockResolvedValue('some-api-key');
+    const event = {
+      headers: {},
+      queryStringParameters: {
+        apiKey: 'some-api-key',
+      },
+      methodArn: 'arn:aws:execute-api:us-east-1:123456789012:11111/prod/$connect',
+    };
+    const response = await handler(event);
+    expect(response).toEqual({
+      principalId: 'authenticated-user',
+      policyDocument: {
+        Version: '2012-10-17',
+        Statement: [
+          {
+            Action: 'execute-api:Invoke',
+            Effect: 'Allow',
+            Resource: 'arn:aws:execute-api:us-east-1:123456789012:11111/prod/$connect',
+          },
+        ],
+      },
+    });
+  });
+
+  it("should return 'Unauthorized' if apiKey doesn't match", async () => {
+    jest.spyOn(secretsmanager, 'getSecretValue').mockResolvedValue('some-api-key');
+    const event = {
+      headers: {},
+      queryStringParameters: {
+        apiKey: 'wrong-api-key',
+      },
+      methodArn: 'arn:aws:execute-api:us-east-1:123456789012:11111/prod/$connect',
+    };
+    const response = await handler(event);
+    expect(response).toEqual('Unauthorized');
+  });
+
+  it("should return 'Unauthorized' if apiKey is missing", async () => {
+    jest.spyOn(secretsmanager, 'getSecretValue').mockResolvedValue('some-api-key');
+    const event = {
+      headers: {},
+      queryStringParameters: {},
+      methodArn: 'arn:aws:execute-api:us-east-1:123456789012:11111/prod/$connect',
+    };
+    const response = await handler(event);
+    expect(response).toEqual('Unauthorized');
+  });
+});