Success With AWS Web Socket Endpoint Configuration

by Stephen Owens

I had a lot of trouble getting the API Gateway socket management calls to work using the AWS SDK V2 version of the Golang SDK. The problem was that I was not invoking the right endpoint. So that I remember in the future how to do this and to help anyone having similar problems here’s the magic formula I finally learned.

Endpoint Value

First you need to know the endpoint to use to make calls against. You can either look it up in the AWS console or get it dynamically from the request. I chose to get it from the request, but note that it is only set properly in the initial connect call to the socket, not during any other socket interaction. Here’s how I get it from the initial http connect call.

package main

import (
  "context"

  "github.com/aws/aws-lambda-go/events"
  "github.com/aws/aws-lambda-go/lambda"
)

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

func handler(ctx context.Context, request events.APIGatewayWebsocketProxyRequest) (awsproxy.Response, error) {

  hostName := request.RequestContext.DomainName
  InitHost(ctx, hostName)

  return events.APIGatewayProxyResponse{
    StatusCode:      200,
    IsBase64Encoded: false,
    Body:            "success",
    Headers: map[string]string{
      "Content-Type": "text/plain; charset=UTF-8",
    },
  }, nil
}

Now that hostName contains the correct endpoint you would need to store it somewhere. I chose to put it in a record in DynamoDB where it can be found whenever needed by other Lambdas. See the InitHost method shown below.

Invoking the Endpoint

The stored endpoint is used to initialize an api client interface for the API Gateway web socket command interface. Here’s the code I use to do that along with a sample method that uses the resulting API client to make a command request.

package notification

import (
  "bytes"
  "context"
  "encoding/json"
  "net/url"
  "time"

  "github.com/aws/aws-sdk-go-v2/aws"
  "github.com/aws/aws-sdk-go-v2/config"
  "github.com/sowens-csd/ftlambdas/awsproxy"
  "github.com/sowens-csd/ftlambdas/ftdb"
  "github.com/sowens-csd/ftlambdas/messaging"

  "github.com/aws/aws-sdk-go-v2/service/apigatewaymanagementapi"
  apigtypes "github.com/aws/aws-sdk-go-v2/service/apigatewaymanagementapi/types"
)

// cfg is the base or parent AWS configuration for this lambda.
var cfg aws.Config

// apiClient provides access to the Amazon API Gateway management functions. Once initialized, the instance is reused
// across subsequent AWS Lambda invocations. This potentially amortizes the instance creation over multiple executions
// of the AWS Lambda instance.
var apiClient *apigatewaymanagementapi.Client

// endpointHost holds the current value of the host name that can be used to reach the Amazon API Gateway endpoint.
// This value is only available in the initial connect call so it is stored in DynamoDB
// and then initialized in the initHost function
var endpointHost string

// InitHost get the domainName from the DB or set it if it doesn't exist
// This must either set a valid host or find one in the DB, anything else
// panics. 
func InitHost(ctx context.Context, domainName string) {
  if len(endpointHost) > 0 {
    // it's already been initialized okay to leave
    return
  }

  storedHost = ReadEndpointHostFromDB(ctx)
  if "" == storedHost {
    if Len(domainName)>0 {
      endpointHost = domainName
      WriteEndpointHostToDB(ctx, endpointHost)
    }
  } else {
    endpointHost = storedHost
  }
  if len(endpointHost) == 0 {
    panic("no endpoint host found")
  }
}

func initConfig(ctx context.Context) {
  if nil == apiClient {
    var err error
    cfg, err = config.LoadDefaultConfig(ctx)
    if err != nil {
      panic("unable to load SDK config")
    }
    apiClient = newAPIGatewayManagementClient(ctx, &cfg, endpoint, "dev")
  }
}

// newAPIGatewayManagementClient creates a new API Gateway Management Client instance from the provided parameters. The
// new client will have a custom endpoint that resolves to the application's deployed API.
func newAPIGatewayManagementClient(ctx context.Context, cfg *aws.Config, domain, stage string) *apigatewaymanagementapi.Client {
  cp := cfg.Copy()
  cp.EndpointResolver = aws.EndpointResolverFunc(func(service, region string) (aws.Endpoint, error) {

  var endpoint url.URL
  endpoint.Path = stage
  endpoint.Host = endpointHost
  endpoint.Scheme = "https"
  return aws.Endpoint{
    SigningRegion: region,
    URL: endpoint.String(),
    }, nil
  })

  return apigatewaymanagementapi.NewFromConfig(cp)
}

// publish publishes the provided data to the Amazon API Gateway connection ID. A common failure scenario which
// results in an error is if the connection ID is no longer valid. This can occur when a client disconnected from the
// Amazon API Gateway endpoint but the disconnect AWS Lambda was not invoked as it is not guaranteed to be invoked when
// clients disconnect.
func publish(ctx context.Context, id string, data []byte) error {
  initConfig(ftCtx)
  ctxTimeout, _ := context.WithTimeout(ctx, 10*time.Second)
  _, err := apiClient.PostToConnection(ctxTimeout, 
    &apigatewaymanagementapi.PostToConnectionInput{
      Data:         data,
      ConnectionId: &id,
    })

  return err
}

There’s a lot going on in these methods and I’m trying to make this a short post so I’ll just point out a few highlights.

  • The endpointHost is the value that we received from the connect request and then stored in the DB as pointed out above.
  • InitHost has to be called at least once in the lifetime of the Lambda to ensure that endpointHost is setup.
  • The apiClient, once configured, ‘knows’ the correct endpoint because it was configured with that endpoint as it was built. As far as I can tell you must set the endpoint for web socket calls or it won’t work and fails in a variety of confusing ways.
  • The key to the magic to make the apiClient know the endpoint is the cp.EndpointResolver which uses the aws.EndpointResolverFunc. You can see the details on that in the newAPIGatewayManagementClient function shown above.
  • The ReadEndpointHostFromDB and WriteEndpointHostToDB methods are left as an exercise for the reader.
  • publish uses the apiClient to send a message on the socket.

Leave a Reply