AWS API Gateway Websocket with Cloudformation and Golang (with authentication)

Tue, 24 Nov 2020

You already know "how to make a web chat" or a real-time system with websocket, but in reality there are many other use cases. Consider the following example or in general an asynchronous communication between services.

Your client may call an API endpoint which sends a message into a queue, a VM will poll the message, process it and save the result into a database; all of this without the user needing to wait for a response. This is very convenient and makes the system less susceptible to latency, but there is a little caveat; the client usually needs to get a response or at least an acknowledgement of the process status. And how do we get something back in a situation like that? Well, we have some options here:

• Store the result somewhere and let the client poll for it at fixed interval

• Push the update directly to the client when the process has terminated or has some change in its state

The former is clearly the simplest one, but it is very inefficient and you risk to unnecessarily overload your infrastructures and perhaps getting a delay. The latter is clearly a bit more complex as you need to keep track of the connected clients and add additional resources to support the websocket communication, but on the other hand is a very efficient and scalable solution.

I guess many other posts have done this topic, however almost all of them use Serverless framework, thus I will be using SAM/Cloudformation and golang as a backend language. As always if you are looking for a quick solution you can skip the whole article and visit this repository where you can get all the code, otherwise you can follow along.

The API

At the time I wrote this article AWS SAM does not fully support API Gateway websocket, therefore it has to be written in raw Cloudformation:

WebSocket:
  Type: AWS::ApiGatewayV2::Api
  Properties:
    ProtocolType: WEBSOCKET
    RouteSelectionExpression: "$request.body.action"
    Name: !Sub ${Appname}-${Environment}

Notice the RouteSelectionExpression which means that if you want to trigger a custom route you need to add a property in the body called action, an example payload that will call your pushtowebsocket route could be the following:

{"message": "hello Matteo", "action": "pushtowebsocket"}

Your API needs a stage and a deployment as well:

Make sure to add this resource, if you do not include it you will still be able to make everything work, however, the next time you will try to deploy the template, all the changes you made to API Gateway will not be reflected

WebSocketStage:
  Type: AWS::ApiGatewayV2::Stage
  Properties:
    StageName: !Ref Environment
    AutoDeploy: true
    ApiId: !Ref WebSocket

Deployment:
  Type: AWS::ApiGatewayV2::Deployment
  DependsOn:
    - ConnectRoute
    - PushToWebSocket
  Properties:
    ApiId: !Ref WebSocket

Your deployments needs at least a route in order to work, therefore it needs to be created after. The DependsOn should be specified on possibly all the routes you have.

Authentication

Before adding the authentication there are some important remarks:

• Authentication works only for the $connect route.

• Authorizer must have a specified credential, otherwise API Gateway will reject the request with 401 without any specific error.

• Context object must contains only primitive types (e.g. no arrays or nested objects).

• Only custom REQUEST type authorizers work with websocket.

An API Gateway websocket has two default routes/events:

• $connect: executed after the upgrade request has succeed. You can use it to save the connection Id and authenticate the client.

• $disconnect: is executed after the connection is closed, however API Gateway cannot guarantee its delivery, which means that it might not be always triggered. Usually to delete the connection Id.

Both routes will trigger an event, It is up to the user to attach them to a function or any other service. Let's add the `$connect` route:

1. Route declaration:

   ConnectRoute:
     Type: AWS::ApiGatewayV2::Route
     Properties:
       ApiId: !Ref WebSocket
       RouteKey: "$connect"
       AuthorizationType: CUSTOM
       AuthorizerId: !Ref Authorizer
       Target: !Join
         - '/'
         - - 'integrations'
           - !Ref ConnectRouteIntegration
   

2. Route integration:

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

3. The function itself:

   onConnect:
     Type: AWS::Serverless::Function
     Properties:
       CodeUri: ./functions/onConnect
       Handler: main
       Policies:
         - AWSLambdaBasicExecutionRole
   

4. Function's permission:

   OnConnectPermission:
     Type: AWS::Lambda::Permission
     Properties:
       Action: lambda:InvokeFunction
       FunctionName: !Ref onConnect
       Principal: apigateway.amazonaws.com

For the authorizer, instead, we need three components:

1. Authorizer declaration:

   Authorizer:
     Type: AWS::ApiGatewayV2::Authorizer
     Properties:
       Name: LambdaAuthorizer
       ApiId: !Ref WebSocket
       AuthorizerType: REQUEST
       AuthorizerUri:
         Fn::Sub: arn:aws:apigateway:${AWS::Region}:lambda:path/2015-03-31/functions/${RequestAuthorizer.Arn}/invocations
       IdentitySource:
         - route.request.querystring.Auth
   

2. Function:

   RequestAuthorizer:
     Type: AWS::Serverless::Function
     Properties:
       Policies:
         - AWSLambdaBasicExecutionRole
       CodeUri: ./functions/authorizer
       Handler: main
   

3. Permission:

   RequestAuthorizerPermission:
     Type: AWS::Lambda::Permission
     Properties:
       Action: lambda:InvokeFunction
       Principal: apigateway.amazonaws.com
       FunctionName:
         Ref: RequestAuthorizer

One question that might immediately arise is: how can I pass the token in the Authorization header? Well the answer is: you can't. The JavaScript API does not allow you to pass headers in general. Therefore one solution would be to pass it as query string parameter. Thus we need to specify

route.request.querystring.Auth

as identity source. When you try to connect to the websocket you must pass the Auth parameter even if you are only testing the function, otherwise API Gateway will reject your call no matter what:

wscat -c 'wss://<api-id>.execute-api.<region>.amazonaws.com/<stage>?Auth=myFakeToken'

A custom authorizer must return a policy response with this format:

{
	"principalId": "my-username",
	"policyDocument": {
		"Version": "2012-10-17",
		"Statement": [
			{
				"Action": "execute-api:Invoke",
				"Effect": "Allow",
				"Resource": "arn:aws:execute-api:us-east-1:123456789012:qsxrty/test/GET/mydemoresource"
		]
	},
	"context": {
		"org": "my-org",
		"role": "admin",
		"createdAt": "2019-01-03T12:15:42"
	}
}

Essentially allow the caller to invoke the endpoint. In addition you can return a context object in which you can pass custom data to a downstream service and you can simply extract it in your lambda function:

func function(event events.APIGatewayWebsocketProxyRequest) (events.APIGatewayProxyResponse, error) {
   ctx := event.RequestContext.Authorizer
   ...
}

The lambda authorizer signature should be the following:

func function(event events.APIGatewayCustomAuthorizerRequestTypeRequest) (events.APIGatewayCustomAuthorizerResponse, error)

The event contains more information than a normal TOKEN type authorizer:

{
    "type": "REQUEST",
    "methodArn": "arn:aws:execute-api:us-east-1:123456789012:s4x3opwd6i/test/GET/request",
    "resource": "/request",
    "path": "/request",
    "httpMethod": "GET",
    "headers": {
      //...
    },
    "queryStringParameters": {
        //...
    },
    "pathParameters": {},
    "stageVariables": {
        //...
    },
    "requestContext": {
      //...
    }
}

The methodArn is the ARN of the incoming method request and is populated by API Gateway in accordance with the Lambda authorizer configuration; you can use this parameter as Resource of your IAM policy to return in the response. As I said before, you can use queryStringParameters to get your token from the parameter Auth:

token := event.QueryStringParameters["Auth"]

as you are using a custom authorizer you need to validate the token yourself.

Custom route and pushing messages

So let's add a custom route, you can use this route as you want, you may call it internally in your system to send updates to the client, but for sake of completeness I will put it as an echo route. The declaration is mostly the same, the only things that change is the RouteKey in this case can be an arbitrary name and with no authorization type:

PushToWebSocket:
  Type: AWS::ApiGatewayV2::Route
  Properties:
    ApiId: !Ref WebSocket
    RouteKey: pushtowebsocket
    AuthorizationType: NONE
    OperationName: PushToWebSocket
    Target: !Join
      - '/'
      - - 'integrations'
        - !Ref PushToWebSocketIntegration

Now let's attach our function integration:

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

And of course you need to attach the function itself and the appropriate permissions:

PushToWebSocketFunction:
  Type: AWS::Serverless::Function
  Properties:
    Policies:
      - AWSLambdaBasicExecutionRole
      - Statement:
          - Effect: Allow
            Action:
              - execute-api:ManageConnections
            Resource:
              - !Sub arn:aws:execute-api:${AWS::Region}:${AWS::AccountId}:${WebSocket}/*
    CodeUri: ./functions/pushToWebSocket
    Handler: main
    Environment:
      Variables:
        API_GATEWAY_ID: !Ref WebSocket

PushToWebSocketPermission:
  Type: AWS::Lambda::Permission
  Properties:
    Action: lambda:InvokeFunction
    FunctionName: !Ref PushToWebSocketFunction
    Principal: apigateway.amazonaws.com

Notice the

execute-api:ManageConnections 

this permission is needed to use the PostToConnection method. This API uses the https endpoint not the ws, and given a connectionId will push the message to the connected client:

input := &apigatewaymanagementapi.PostToConnectionInput{
   ConnectionId: aws.String(connectionId),
   Data:         []byte("Hello there!"),
}
if _, err := api.PostToConnection(input); err != nil {
   return events.APIGatewayProxyResponse{
      StatusCode: 500,
      Body:       "Something went wrong",
   }
}

Every lambda invocation from an API Gateway websocket will receive an event object which will carry the request context where we can find the connectionId

func function(request events.APIGatewayWebsocketProxyRequest) {
   connectionId := request.RequestContext.ConnectionID
   // rest of the code
}

func main() {
   lambda.Start(function)
}

You can read here more info about the request context. You could use the same context to initialize the apigatewaymanagementapi, as inside you can find the domain name and resource id, however I prefer to pass the API Gateway id as an environmental variable. In this way the function can be invoked by another service and you still can access the API Id:

var (
   initializedSession = session.Must(session.NewSessionWithOptions(session.Options{
      SharedConfigState: session.SharedConfigEnable,
   }))
   apiGatewayId = os.Getenv("API_GATEWAY_ID")
   region       = os.Getenv("AWS_REGION")
   environment  = os.Getenv("GO_ENV")
)

func GetApiGatewayEndpoint(apiGatewayId string) string {
   return fmt.Sprintf("%v.execute-api.%v.amazonaws.com/%v", apiGatewayId, region, environment)
}

func NewApiGatewayManagementApi() *apigatewaymanagementapi.ApiGatewayManagementApi {
   return apigatewaymanagementapi.New(initializedSession,
      aws.NewConfig().WithEndpoint(GetApiGatewayEndpoint(apiGatewayId)))
}

Now that we have all the pieces of our LEGO project, let's go back to our initial architecture and add those bricks in there.

What we can usually do in this case is to pass the connectionId downstream as part of our payload. For example you can extract it in the first Lambda and pass it through SNS, SQS message body and then as a final step retrieve it and use it to push the response do the client. This is a good approach if you have an highly event driven asynchronous architecture where the client that started it needs the acknowledgement.

As rule of thumb always consider the wire time: if the time of a normal HTTP request-response is shorter, than maybe you should avoid this strategy and adopt a synchronous communication.

As always you can clone my repository and dive into the code. I hope you found this post interesting, if you have any questions or corrections feel free to send me an email. I will leave you one more article to read from the aws blog.

Cheers

Thanks for reading!