Project Doorbell: How to use Amazon API Gateway WebSocket APIs with Rust

Project Doorbell: How to use Amazon API Gateway WebSocket APIs with Rust

This blog post is all about Amazon API Gateway and in the specific working with WebSocket.

As a part of the series Project Doorbell you can check out the other parts:

Part 1 a light introduction to Amazon Rekognition, the AI service that will identify people.

Part 2 is about the setup taken to run the project and some essential parts to connect AWS S3 with AWS Step Function by Amazon EventBridge.

WebSocket

A WebSocket is a persistent connection between a client and a server. It is a two-way communications channel over HTTP through a single TCP/IP socket connection. WebSocket is distinct from HTTP. Both protocols are located at layer 7 in the OSI model and depend on TCP at layer 4.

Websocket_connection.png

Before WebSocket, we had no way to send data from the server to the client because HTTP is one-way communication. So a workaround was what we know as Long-polling, a technique where the client regularly asks the server for new data.

WebSockets do not use the HTTP(s):// scheme but WS(s):// schema but for the rest of the URI is the same with a host, port, path and any query parameters.

Amazon API Gateway with WebSocket

AWS announced the WebSocket support back in Dec 2018

websockets-arch.png

With WebSocket API, we have the concept of a route. A route describes how to handle a specific type of request and includes a routeKey parameter and a value to identify the route.

There are three routeKey values that API Gateway allows you to use for a route:

  • $default: when the route does not match any route keys usually used to implement a generic error handling.
  • $connect: when a client first connects.
  • $disconnect: when a client disconnects.

I had the opportunity to play around with WebSocket APIs in Amazon API Gateway, and you can find some code in GitHub:

  • Basic usage, here
  • Trying to use it at scale, here

To secure the WebSocket API, there are the following mechanisms for authentication and authorization:

  • AWS IAM roles: for controlling who can create and manage your APIs, as well as who can invoke them.
  • IAM tags: can be used together with IAM policies to control access.
  • Lambda authorizers.

In this project, I do not need to work at a particular scale with Amazon API Gateway with WebSocket because you ring a doorbell not often at the door, and for now, I do not control access.

Rust with WebSocket

If you have followed me with Serverless Rust, you have seen that there are two main creates from AWS to use Rust in a serverless context.

  1. lambda-runtime: is a library that provides a Lambda runtime for applications written in Rust.
  2. lambda-http: is a library that makes it easy to write API Gateway proxy event focused Lambda functions in Rust.

And like me, you would use the lambda-http create to discover that the WebSocket request is not supported out of the box, and you will have this error:

Error: Error("data did not match any variant of untagged enum LambdaRequest"

I have opened an issue to create some visibility to the topic:

WebSockets would require some design thoughts to get a friendly interface, and I doubt they will support it as part of lambda_http, but AWS teams could surprise us. Until then, the way to work is to create your input struct.

An example of a WebSocket request is the following:

{
   "headers":"Object("{
      "Host":"String(""xxxx.execute-api.xxxx.amazonaws.com"")",
      "Sec-WebSocket-Extensions":"String(""permessage-deflate; client_max_window_bits"")",
      "Sec-WebSocket-Key":"String(""xxxx"")",
      "Sec-WebSocket-Version":"String(""13"")",
      "X-Amzn-Trace-Id":"String(""Root=1-61e2cb5d-xxxx"")",
      "X-Forwarded-For":"String(""xxx.xxx.xxx.xxx"")",
      "X-Forwarded-Port":"String(""443"")",
      "X-Forwarded-Proto":"String(""https"")"
   }")",
   "isBase64Encoded":"Bool(false)",
   "multiValueHeaders":"Object("{
      "Host":"Array("[
         "String(""xxxx.execute-api.xxxx.amazonaws.com"")"
      ]")",
      "Sec-WebSocket-Extensions":"Array("[
         "String(""permessage-deflate; client_max_window_bits"")"
      ]")",
      "Sec-WebSocket-Key":"Array("[
         "String(""xxxx"")"
      ]")",
      "Sec-WebSocket-Version":"Array("[
         "String(""13"")"
      ]")",
      "X-Amzn-Trace-Id":"Array("[
         "String(""Root=1-61e2cb5d-xxxx"")"
      ]")",
      "X-Forwarded-For":"Array("[
         "String(""xxx.xxx.xxx.xxx"")"
      ]")",
      "X-Forwarded-Port":"Array("[
         "String(""443"")"
      ]")",
      "X-Forwarded-Proto":"Array("[
         "String(""https"")"
      ]")"
   }")",
   "requestContext":"Object("{
      "apiId":"String(""xxxxxx"")",
      "connectedAt":Number(1642253149666),
      "connectionId":"String(""L_S2qf_VliACH_g="")",
      "domainName":"String(""xxxx.execute-api.xxxx.amazonaws.com"")",
      "eventType":"String(""CONNECT"")",
      "extendedRequestId":"String(""L_S2qHEVFiAFk9w="")",
      "identity":"Object("{
         "sourceIp":"String(""xxx.xxx.xxx.xxx"")"
      }")",
      "messageDirection":"String(""IN"")",
      "requestId":"String(""L_S2qHEVFiAFk9w="")",
      "requestTime":"String(""15/Jan/2022:13:25:49 +0000"")",
      "requestTimeEpoch":Number(1642253149666),
      "routeKey":"String(""$connect"")",
      "stage":"String(""test"")"
   }")"
}

I am interested in the connectionId, and so I built a custom structure:

use serde::{Deserialize, Serialize};

#[derive(Deserialize, Debug, Clone)]
#[serde(rename_all = "camelCase")]
pub struct Event {
    pub request_context: RequestContext,
}

#[derive(Deserialize, Debug, Clone)]
#[serde(rename_all = "camelCase")]
pub struct RequestContext {
    pub event_type: EventType,
    pub connection_id: String,
}

#[derive(Serialize, Deserialize, Debug, Clone)]
#[serde(rename_all = "UPPERCASE")]
pub enum EventType {
    Connect,
    Disconnect,
}

Having this place will allow me to accept the request and start coding.

use rust_doorbell::aws::client::AWSClient;
use rust_doorbell::dtos::request::Event;
use lambda_runtime::{handler_fn, Context, Error};
use tracing::{info};
use rust_doorbell::{utils::*};

#[tokio::main]
async fn main() -> Result<(), Error> {
  // Initialize service
  setup_tracing();
  // Initialize AWS client
  let aws_client = get_aws_client().await;

  lambda_runtime::run(handler_fn(|event: Event, ctx: Context| {
    execute(&aws_client, event, ctx)
  }))
  .await?;
  Ok(())
}

pub async fn execute(aws_client: &AWSClient, event: Event, _ctx: Context) -> Result<(), Error> {
  info!("EVENT {:?}", event);
  info!("event_type {:?}", event.request_context.event_type);
  info!("connection_id {:?}", event.request_context.connection_id);

  Ok(())
}

Provision Infrastructure As Code

AWS CloudFormation is an infrastructure as code (IaC) service that allows you to provision and manage AWS resources. As per schematic:

websockets-arch.png

I need to provision an Amazon API Gateway with these routes.

  • $connect: when a client first connects.
  • $disconnect: when a client disconnects.

Each route will integrate with a Lambda function. An integration request involves the following:

  • Choosing a route key to integrate to the backend.
  • Specifying the backend endpoint to invoke, such as a Lambda function.
  • Configuring how to transform the route request data (optional).

Each Lambda function should handle the connectionId with DynamoDB, and so they will require two different permission patterns for the same table.

AWSTemplateFormatVersion: 2010-09-09
Transform: AWS::Serverless-2016-10-31
Description: An Amazon API Gateway WebSocket API and an AWS Lambda function.

# Global values that are applied to all applicable resources in this template
Globals:
  Function:
    MemorySize: 1024
    Architectures: ["arm64"]
    Handler: bootstrap
    Runtime: provided.al2
    Timeout: 29
    Environment:
      Variables:
        RUST_BACKTRACE: 1
        RUST_LOG: info

Resources:
##########################################################################
#   DYNAMODB                                                            #
##########################################################################
  WebSocketTable:
    Type: AWS::DynamoDB::Table
    Properties:
      AttributeDefinitions:
        - AttributeName: connectionId
          AttributeType: S
      KeySchema:
        - AttributeName: connectionId
          KeyType: HASH
      BillingMode: PAY_PER_REQUEST

##########################################################################
#   Lambda Function                                                      #
##########################################################################
  OnConnectLambdaFunction:
    Type: AWS::Serverless::Function
    Properties:
      CodeUri: ../build/on-connect/
      Policies:
        - AWSLambdaBasicExecutionRole
        - Version: "2012-10-17"
          Statement:
            - Effect: Allow
              Action: 
                - dynamodb:PutItem
              Resource: !GetAtt WebSocketTable.Arn
      Environment:
        Variables:
          TABLE_NAME: !Ref WebSocketTable

  OnConnectFunctionResourcePermission:
    Type: AWS::Lambda::Permission
    Properties:
      Action: lambda:InvokeFunction
      Principal: apigateway.amazonaws.com
      FunctionName: !Ref OnConnectLambdaFunction
      SourceArn: !Sub arn:aws:execute-api:${AWS::Region}:${AWS::AccountId}:${WebSocketApi}/*

  OnDisconnectLambdaFunction:
    Type: AWS::Serverless::Function
    Properties:
      CodeUri: ../build/on-disconnect/
      Policies:
        - AWSLambdaBasicExecutionRole
        - Version: "2012-10-17"
          Statement:
            - Effect: Allow
              Action: 
                - dynamodb:DeleteItem
              Resource: !GetAtt WebSocketTable.Arn
      Environment:
        Variables:
          TABLE_NAME: !Ref WebSocketTable

  OnDisconnectFunctionResourcePermission:
    Type: AWS::Lambda::Permission
    Properties:
      Action: lambda:InvokeFunction
      Principal: apigateway.amazonaws.com
      FunctionName: !Ref OnDisconnectLambdaFunction
      SourceArn: !Sub arn:aws:execute-api:${AWS::Region}:${AWS::AccountId}:${WebSocketApi}/*

##########################################################################
#   API Gateway WebSocket API                                            #
##########################################################################
  WebSocketApi:
    Type: AWS::ApiGatewayV2::Api
    Properties:
      Name: !Ref AWS::StackName
      Description: An Amazon API Gateway WebSocket API and an AWS Lambda function.
      ProtocolType: WEBSOCKET
      RouteSelectionExpression: "$request.body.action"

  OnConnectIntegration:
    Type: AWS::ApiGatewayV2::Integration
    Properties:
      ApiId: !Ref WebSocketApi
      Description: OnConnect Integration
      IntegrationType: AWS_PROXY
      IntegrationUri: 
        Fn::Sub:
            arn:aws:apigateway:${AWS::Region}:lambda:path/2015-03-31/functions/${OnConnectLambdaFunction.Arn}/invocations

  OnConnectRoute:
    Type: AWS::ApiGatewayV2::Route
    Properties:
      ApiId: !Ref WebSocketApi
      RouteKey: $connect
      AuthorizationType: NONE
      OperationName: OnConnectRoute
      Target: !Join
        - /
        - - integrations
          - !Ref OnConnectIntegration

  OnDisconnectIntegration:
    Type: AWS::ApiGatewayV2::Integration
    Properties:
      ApiId: !Ref WebSocketApi
      Description: OnDisconnect Integration
      IntegrationType: AWS_PROXY
      IntegrationUri: 
        Fn::Sub:
            arn:aws:apigateway:${AWS::Region}:lambda:path/2015-03-31/functions/${OnDisconnectLambdaFunction.Arn}/invocations

  OnDisconnectRoute:
    Type: AWS::ApiGatewayV2::Route
    Properties:
      ApiId: !Ref WebSocketApi
      RouteKey: $disconnect
      AuthorizationType: NONE
      OperationName: OnDisconnectRoute
      Target: !Join
        - /
        - - integrations
          - !Ref OnDisconnectIntegration

  Deployment:
    Type: AWS::ApiGatewayV2::Deployment
    DependsOn:
      - OnConnectRoute
      - OnDisconnectRoute
    Properties:
      ApiId: !Ref WebSocketApi

  Stage:
    Type: AWS::ApiGatewayV2::Stage
    Properties:
      StageName: prod
      Description: Prod Stage
      DeploymentId: !Ref Deployment
      ApiId: !Ref WebSocketApi

Outputs:
  OnConnectLambdaFunctionArn:
    Description: "OnConnect function ARN"
    Value: !GetAtt OnConnectLambdaFunction.Arn
  OnDisconnectLambdaFunctionArn:
    Description: "OnDisconnect function ARN"
    Value: !GetAtt OnDisconnectLambdaFunction.Arn
  WebSocketURL:
    Description: "The WSS Protocol URL to connect to"
    Value: !Join [ '', [ wss://, !Ref WebSocketApi, .execute-api.,!Ref AWS::Region,.amazonaws.com/,!Ref Stage] ]

Conclusion

WebSockets are handy if you need two-way communications. Using AWS Lambda and Amazon DynamoDB, we built a real-time application. As usual, this is made possible using managed services. Because I do not need to scale this solution, the WebSockets fit this IoT project. I should have it all in place for a dry run on the next episode of this series. After that, I will show you the steps to simulate the connections between the doorbell and your AWS backend.