API Gateway WebSocket APIs with the Serverless Framework
Tags: apigateway, aws, serverlessframework, websocket
Published: 2019-07-01 22:35:43

---
title: API Gateway WebSocket APIs with the Serverless Framework
description: A walkthrough of a simple WebSocket application.
published: true
tags: api-gateway,aws,serverless-framework,websocket
canonical_url: https://blog.neverendingqs.com/2019/07/01/serverless-websocket-example.html
---

I recently built a web application in JavaScript that leveraged WebSockets to display live data from a server. I used [AWS API Gateway WebSocket APIs](https://docs.aws.amazon.com/apigateway/latest/developerguide/apigateway-websocket-api.html) in the back-end and the [WebSocket API](https://developer.mozilla.org/en-US/docs/Web/API/WebSockets_API) in the front-end. I also used the [Serverless Framework](https://serverless.com/framework/docs/providers/aws/events/websocket/) ("Serverless" for short) to deploy the back-end.

In this post, I will do a quick walkthrough of [a simple WebSocket application](https://github.com/neverendingqs/serverless-websocket-example) I wrote. Although I will be sharing snippets of Serverless configuration, the examples should be informative enough to apply it to the deployment method you are using.

## Back-End

The backend consists of an API Gateway WebSockets API, some Lambda functions, and a persistent store.

### Routing WebSocket Messages

The first thing I had to do was create a WebSocket API and configure the three predefined routes:

```yaml
# serverless.yml
functions:
  # Create a Lambda function that responds to messages for the predefined routes
  websocket:
    handler: src/websocket.handler
    events:
      - websocket:
          # Handles new connection requests
          route: $connect

      - websocket:
          # Route messages here if '$request.body.action' is 'routeA'.
          # You can adjust which property to use for routing by adjusting
          # 'websocketsApiRouteSelectionExpression'.
          route: routeA

      - websocket:
          # Handles all unrouted messages
          route: $default

      - websocket:
          # Handles disconnect messages
          route: $disconnect
```

Serverless will configure API Gateway to call the Lambda function defined at `src/websocket.handler` whenever it receives a message. The function should look something like this:

```js
// src/websocket.js
exports.handler = async function(event, context) {
  const { requestContext: { routeKey }} = event;
  switch(routeKey) {
    case '$connect':
      ...
      break;

    case '$disconnect':
      ...
      break;

    case 'routeA':
      ...
      break;

    case '$default':
    default:
      ...
  }

  // Return a 200 status to tell API Gateway the message was processed
  // successfully.
  // Otherwise, API Gateway will return a 500 to the client.
  return { statusCode: 200 };
}
```

I chose to handle messages from any route in the same Lambda function, but you can instead split that out into more than one function. For more details on how to do that in Serverless, see their [WebSocket documentation](https://serverless.com/framework/docs/providers/aws/events/websocket/).

### WebSocket Sessions

When a WebSocket client requests a new connection, API Gateway assigns a connection ID to that session and invokes your `$connect` Lambda function with the ID in the `event` payload (`event.requestContext.connectionId`). You can use that right away to start sending messages to the client in the same function invocation.  If you would like to send messages at a later time, you need to store the connection ID somewhere.

In the code example, I used DynamoDB to store the connection ID:

```js
// src/websocket.js
case '$connect':
  await dynamodb.put({
    TableName: connectionTable,
    Item: {
      connectionId,
      // Expire the connection an hour later. This is optional, but recommended.
      // You will have to decide how often to time out and/or refresh the ttl.
      ttl: parseInt((Date.now() / 1000) + 3600)
    }
  }).promise();
  break;
```

This allows the [broadcast Lambda function](https://github.com/neverendingqs/serverless-websocket-example/blob/master/src/broadcast.js) to get a list of all connected WebSocket clients when broadcasting a message.

I also made sure to clean up sessions as clients close their WebSocket connections:

```js
// src/websocket.js
case '$disconnect':
  await dynamodb.delete({
    TableName: connectionTable,
    Key: { connectionId }
  }).promise();
  break;
```

### Using the ApiGatewayManagementApi Class

To send a message back to the client, you need to use the [ApiGatewayManagementApi class](https://docs.aws.amazon.com/AWSJavaScriptSDK/latest/AWS/ApiGatewayManagementApi.html). Unlike other AWS services where you provide the resource location when calling an instance method (e.g. the S3 bucket name for `putObject()` or the DynamoDB table name for `putItem()`), you need to provide the URL to your WebSocket API on class creation:

```js
const apig = new AWS.ApiGatewayManagementApi({
  endpoint: // your WebSocket API endpoint here
});
```

You can then use `postToConnection()` to send messages to the client:

```js
await apig.postToConnection({
  ConnectionId: connectionId,
  Data: body
}).promise();
```

The endpoint follows a specific format and you can use CloudFormation's join and reference functions to format it:

```yaml
# serverless.yml
# <logical ID of the WebSocket API>.execute-api.<AWS region>.amazonaws.com/<stage>
APIG_ENDPOINT:
  Fn::Join:
    - ''
    - - Ref: WebsocketsApi
      - .execute-api.
      - Ref: AWS::Region
      - .amazonaws.com/
      - ${self:custom.stage}
```

If you do not provide an endpoint when creating a new `ApiGatewayManagementApi` object, you will get an error like the following when you call `postToConnection()`:

```
{
    "errorMessage": "Inaccessible host: `execute-api.us-east-1.amazonaws.com'. This service may not be available in the `us-east-1' region.",
    "errorType": "Error",
    "stackTrace": [
        "UnknownEndpoint: Inaccessible host: `execute-api.us-east-1.amazonaws.com'. This service may not be available in the `us-east-1' region.",
        "    at Request.ENOTFOUND_ERROR (/serverless-websocket-example/node_modules/aws-sdk/lib/event_listeners.js:494:46)",
        "    at Request.callListeners (/serverless-websocket-example/node_modules/aws-sdk/lib/sequential_executor.js:106:20)",
        "    at Request.emit (/serverless-websocket-example/node_modules/aws-sdk/lib/sequential_executor.js:78:10)",
        "    at Request.emit (/serverless-websocket-example/node_modules/aws-sdk/lib/request.js:683:14)",
        "    at ClientRequest.error (/serverless-websocket-example/node_modules/aws-sdk/lib/event_listeners.js:333:22)",
        "    at ClientRequest.<anonymous> (/serverless-websocket-example/node_modules/aws-sdk/lib/http/node.js:96:19)",
        "    at ClientRequest.emit (events.js:189:13)",
        "    at ClientRequest.EventEmitter.emit (domain.js:441:20)",
        "    at TLSSocket.socketErrorListener (_http_client.js:392:9)",
        "    at TLSSocket.emit (events.js:189:13)",
        "    at TLSSocket.EventEmitter.emit (domain.js:441:20)",
        "    at emitErrorNT (internal/streams/destroy.js:82:8)",
        "    at emitErrorAndCloseNT (internal/streams/destroy.js:50:3)",
        "    at process._tickCallback (internal/process/next_tick.js:63:19)"
    ]
}
```

Also, if you are not using the `nodejs10.x` runtime, you might get the following error:

```
{
    "errorMessage": "AWS.ApiGatewayManagementApi is not a constructor",
    "errorType": "TypeError",
    "stackTrace": [
        "Module._compile (module.js:652:30)",
        "Object.Module._extensions..js (module.js:663:10)",
        "Module.load (module.js:565:32)",
        "tryModuleLoad (module.js:505:12)",
        "Function.Module._load (module.js:497:3)",
        "Module.require (module.js:596:17)",
        "require (internal/module.js:11:18)"
    ]
}
```

The AWS JavaScript SDK introduced the `ApiGatewayManagementApi` class in `v2.379.0`, but the `nodejs8.10` runtime only has `v2.290.0` of the SDK. If you are using `nodejs8.10`, you either have to switch to using `nodejs10.x` (`nodejs10.x` is currently at `v2.437.0`) or bundle your own version of `aws-sdk`. See [AWS's Lambda runtime documentation](https://docs.aws.amazon.com/lambda/latest/dg/lambda-runtimes.html) for more details, including information around other language runtimes.

### Creating an Authorizer Lambda Function

An authorizer Lambda function is optional (but recommended). It gets called before the `$connect` Lambda function gets called to make a decision around authorization. For example, you can check for a token in the Authorization header and reject the request if the token is invalid.

In the sample code, I created a very simple authorizer Lambda function that does an uninteresting check:

```js
// src/authorizer.js
function createAuthorizedResponse(resource) {
  return {
    principalId: 'me',
    policyDocument: {
      Version: '2012-10-17',
      Statement: [{
        Action: 'execute-api:Invoke',
        Effect: 'Allow',
        Resource: resource
      }]
    }
  };
}

exports.handler = async function(event, context) {
  // For debug purposes only.
  // You should not log any sensitive information in production.
  console.log("EVENT: \n" + JSON.stringify(event, null, 2));

  const { headers, methodArn } = event;

  // This is for demo purposes only.
  // This check is probably not valuable in production.
  if(headers['X-Forwarded-Proto'] === 'https') {
    return createAuthorizedResponse(methodArn);
  } else {
    throw new Error('Unauthorized');
  }
}
```

For REST APIs, you can configure an authorizer Lambda function for any route. For WebSocket APIs, only the `$connect` route can have an authorizer function associated with it. If you try to configure other routes with an authorizer function, you get an error like the following:

```
An error occurred: routeAWebsocketsRoute - Currently, authorization is restricted to the $connect route only (Service: AmazonApiGatewayV2; Status Code: 400; Error Code: BadRequestException; Request ID: a04ea057-954f-436f-962e-3946639f3d15).
```

This implies authorization is only checked on initial connection, and that the connection ID is the only authorization check for messages sent to other routes. You should protect it like you would a session ID.

If you have an existing session management system for your application that also makes authorization decisions, it must associate the current session with the connection ID assigned by API Gateway. Discussing how to do this is out-of-scope of this post, but you may want to take a look at how you can [pass context from the authorizer Lambda function](https://docs.aws.amazon.com/apigateway/latest/developerguide/api-gateway-lambda-authorizer-output.html) to the Lambda function handling `$connect` messages.

If the authorizer Lambda function rejects the connection request, the WebSocket client will receive a 401.

## Front-End

The [front-end code](https://github.com/neverendingqs/serverless-websocket-example/tree/master/client) is centered around the `WebSocket` class:

```js
// The endpoint's protocol is either 'ws://' or 'wss://'.
// The preferred protocol is 'wss://'.
const websocket = new WebSocket(endpoint);
```

There are four events / event handlers of interest:
* `close` / `onclose`
* `error` / `onerror`
* `message` / `onmessage`
* `open` / `onopen`

To listen to any of these events, create a corresponding function for it:

```js
websocket.onmessage = ({ data }) => {
  // handle the message
};
```

`onmessage` is the most important one, as that is where you handle new messages from the server, while `onopen`, `onerror`, and `onclose` are useful for managing the application's state. `onclose` is also useful so that you know when API Gateway closes the connection, as it does so upon inactivity.

Each event has different properties available, and details are available on the [HTML Living Standard](https://html.spec.whatwg.org/multipage/indices.html#events-2) page.

As a side note, `WebSocket` sends a `$disconnect` message to the server on page reload or close, so you do not have to close the connection yourself in those scenarios.

## Conclusion

In this post, I did a quick walkthrough of what I learned when building a simple `WebSocket` application. Feel free to use any of the [working sample code](https://github.com/neverendingqs/serverless-websocket-example) and please let me know of any questions or comments below!