Websocket con API Gateway y AWS Lambda

En el siguiente código muestro como utilizar Java para crear un Lambda de AWS que permite utilizar Web Socket desde API Gateway.

public class WebSocketLambda implements RequestStreamHandler {

  private static final AWSCredentials CREDENCIALES
          = new EnvironmentVariableCredentialsProvider().getCredentials();

  @Override
  public void handleRequest(InputStream in, OutputStream out, Context cntxt) throws IOException {
    // obtenemos la representación como JSON del request generado por aws
    JsonObject json = Json.createReader(in).readObject();
    if (json.getJsonObject("requestContext").getString("routeKey").equals("$default")) {
      final String payload = json.getString("body"); // contenido que el cliente nos envia
      Request request = new DefaultRequest<>("execute-api");
      request.setHttpMethod(HttpMethodName.POST);
      request.setEndpoint(URI.create(new StringBuilder()
              .append(System.getenv("url_api")) // el API gateway la define
              .append("/")
              // identificador del cliente al que le respondemos
              .append(json.getJsonObject("requestContext").getString("connectionId")) 
              .toString()));
      // escribimos el mismo mensaje que el cliente nos envía
      request.setContent(new ByteArrayInputStream(payload.getBytes(StandardCharsets.UTF_8)));
      AWS4Signer signer = new AWS4Signer();
      signer.setServiceName(request.getServiceName());
      signer.sign(request, CREDENCIALES);
      // enviar request al api gateway
      new AmazonHttpClient(new ClientConfiguration())
              .requestExecutionBuilder()
              .executionContext(new ExecutionContext(true))
              .request(request)
              .errorResponseHandler(new HttpResponseHandler() {
                @Override
                public AmazonClientException handle(HttpResponse rsp) throws Exception {
                  return new AmazonClientException(
                          new StringBuilder()
                                  .append("error al enviar mensaje")
                                  .append(System.getenv("url_api"))
                                  .append(rsp.getStatusCode())
                                  .append(rsp.getStatusText())
                                  .append(leerError(rsp.getContent())).toString());
                }

                @Override
                public boolean needsConnectionLeftOpen() {
                  return false;
                }
              }).execute(new HttpResponseHandler() {
        @Override
        public Void handle(HttpResponse response) {
          return null;
        }

        @Override
        public boolean needsConnectionLeftOpen() {
          return false;
        }
      });
    }
    try (PrintWriter pw = new PrintWriter(out)) {
      JsonObject response = Json.createObjectBuilder()
              .add("statusCode", 200)
              .build();
      pw.write(response.toString());
    }
  }

  private static String leerError(InputStream inputStream) throws IOException {
    final ByteArrayOutputStream bos = new ByteArrayOutputStream();
    byte[] buffer = new byte[1024];
    int length;
    while ((length = inputStream.read(buffer)) != -1) {
      bos.write(buffer, 0, length);
    }
    return bos.toString(StandardCharsets.UTF_8.name());
  }
}

Algunas notas sobre el código:

• Recibe conexiones y responde el mismo mensaje recibido al cliente utilizando la URL que entrega el API Gateway.
• La URL Connection que entrega el API Gateway se define como variable de ambiente (url_api).
• En connectionId esta definido el identificador del cliente entregado por AWS.

En github he dejado el ejemplo completo utilizando maven para construirlo y algunos comentarios adicionales sobre la configuración en AWS.

Algunos enlaces MUY útiles para conocer los Web Sockets en API Gateway:

• https://aws.amazon.com/es/blogs/compute/announcing-websocket-apis-in-amazon-api-gateway/
• https://docs.aws.amazon.com/es_es/apigateway/latest/developerguide/apigateway-websocket-api.html