An example Websocket endpoint on AWS Lambda that echos back the incoming text using streams

ws_stream_echo.js

// To test this code, first create a new lambda with the AWS Cloudformation with the template below;
// Template URL: https://gist.github.com/rikka0w0/53a38add3c17e61744ff1806080e180d
// It not only handles the lambda creation, but also takes care of the required permissions.
// Then, create a Websocket service in the AWS API Gateway and point the $defsult handler to the lambda function we just created.
// Finally, replace the content of index.js of the lambda function with this file.
// Dont forget to DEPLOY it each time after you make changes. Test with 'wscat -c'

const AWS = require('aws-sdk');
const stream = require('stream');
const util = require('util');
const promised_pipeline = util.promisify(stream.pipeline);

exports.handler = async function (event, context, callback) {
  console.log('Enter Websocket Endpoint Handler');

  let connectionInfo;
  let connectionId = event.requestContext.connectionId;

  const callbackAPI = new AWS.ApiGatewayManagementApi({
    apiVersion: '2018-11-29',
    endpoint: event.requestContext.domainName + '/' + event.requestContext.stage,
  });

  try {
    connectionInfo = await callbackAPI.getConnection({ ConnectionId: connectionId }).promise();
  } catch (e) {
    console.log(e);
  }

  connectionInfo.connectionID = connectionId;

  const sendToWebsocket = async (message) => {
    try {
      await callbackAPI.postToConnection({
        ConnectionId: connectionId,
        Data: message,
      }).promise();
    } catch (error) {
      throw error;
    }
  };

  const readableWebSocketStream = new stream.Readable();
  readableWebSocketStream.push(event.body);
  readableWebSocketStream.push(null);

  const writeableWebSocketStream = new stream.Writable({
    async write (chunk, encoding, callback) {
      console.log('!!!!!!!'+chunk.constructor.name);

      try {
        await sendToWebsocket(chunk);
        callback();
      } catch (error) {
        callback(error);
      }
    },
  });

  try {
    await promised_pipeline(readableWebSocketStream, writeableWebSocketStream);
  } catch (error) {
    console.error('Stream pipe error:', error);
    throw error;
  }
  console.log('done');

  return {
    statusCode: 200,
  };
};