API Gateway WebSocket APIs with the Serverless Framework

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 in the back-end and the WebSocket API in the front-end. I also used the Serverless Framework (“Serverless” for short) to deploy the back-end.

In this post, I will do a quick walkthrough of a simple WebSocket application 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:

# https://github.com/neverendingqs/serverless-websocket-example/blob/master/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:

// https://github.com/neverendingqs/serverless-websocket-example/blob/master/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.

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:

// https://github.com/neverendingqs/serverless-websocket-example/blob/master/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 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:

// https://github.com/neverendingqs/serverless-websocket-example/blob/master/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. 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:

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

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

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:

# https://github.com/neverendingqs/serverless-websocket-example/blob/master/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 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:

// https://github.com/neverendingqs/serverless-websocket-example/blob/master/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 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 is centered around the WebSocket class:

// 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:

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 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 and please let me know of any questions or comments below!

Previous

Express.js on Netlify

Netlify started off as a service for hosting static websites. Not only was it an amazing service from the start,...
Next

Waste GUIDs Even Faster with Waste Some GUIDs!

What can be better than being able to indulge in your deepest and most base desire of [wasting GUIDs one...

Software Developer | Lifelong Learner. neverendingqs.com