initial websocket

Author: cnuss

Diff for: package.json

@@ -33,11 +33,13 @@
     "@types/aws-lambda": "^8.10.138",
     "@types/node": "18",
     "@types/which": "^3.0.3",
+    "@types/ws": "^8.5.10",
     "aws-lambda": "^1.0.7",
     "axios": "^1.7.2",
     "execa": "^9.1.0",
     "typescript": "5",
-    "which": "^4.0.0"
+    "which": "^4.0.0",
+    "ws": "^8.17.0"
   },
   "devDependencies": {
     "source-map": "^0.7.4",
@@ -48,4 +50,4 @@
     "webpack": "^5.91.0",
     "webpack-cli": "^5.1.4"
   }
-}
+}

Diff for: src/index.ts

@@ -8,8 +8,9 @@ import {
   EndpointProxyRequest,
   EndpointResponse,
 } from "./internal/types";
+import { WebsocketProxy } from "./internal/websocket";
 
-const { _HANDLER, AWS_LAMBDA_RUNTIME_API } = process.env;
+const { _HANDLER, AWS_LAMBDA_RUNTIME_API, _WEBSOCKET_ROUTE } = process.env;
 
 export const run = async (): Promise<void> => {
   if (!AWS_LAMBDA_RUNTIME_API) {
@@ -20,16 +21,18 @@ export const run = async (): Promise<void> => {
     throw new Error("No handler specified");
   }
 
-  log("Bootstraping", { _HANDLER, AWS_LAMBDA_RUNTIME_API });
+  log("Bootstraping", { _HANDLER, AWS_LAMBDA_RUNTIME_API, _WEBSOCKET_ROUTE });
 
   const { childProcess, bin, endpoint } = await endpointSpawn(
     _HANDLER,
     process.env
   );
 
+  const websocketProxy = new WebsocketProxy(endpoint, _WEBSOCKET_ROUTE);
+
   try {
     log("Polling for events", { bin, endpoint });
-    await pollForEvents(AWS_LAMBDA_RUNTIME_API, bin, endpoint);
+    await pollForEvents(AWS_LAMBDA_RUNTIME_API, bin, endpoint, websocketProxy);
   } catch (e) {
     if (childProcess) {
       log("Killing child process", { pid: childProcess.pid });

Diff for: src/internal/endpoints.ts

@@ -14,19 +14,6 @@ import { info, log } from "./log";
 import { ChildProcess, spawn } from "child_process";
 import { APIGatewayProxyEventV2 } from "aws-lambda";
 
-function wsify(url?: URL): URL | undefined {
-  if (!url) return undefined;
-
-  const wsUrl = new URL(url.toString());
-  wsUrl.protocol = url.protocol.replace("http", "ws");
-
-  if (!wsUrl.protocol.startsWith("ws")) {
-    return undefined;
-  }
-
-  return wsUrl;
-}
-
 function convertHeaders(
   headers: RawAxiosResponseHeaders | AxiosResponseHeaders
 ): { [header: string]: boolean | number | string } | undefined {
@@ -92,8 +79,7 @@ const waitForEndpoint = async (
 
 export const endpointSpawn = async (
   handler: string,
-  env?: NodeJS.ProcessEnv,
-  detached: boolean = true
+  env?: NodeJS.ProcessEnv
 ): Promise<SpawnResult> => {
   // handler is in the format of
   // - `{some-bin}@http://localhost:{the-bins-port} (will start some-bin, and forward requests to the http server)
@@ -120,7 +106,7 @@ export const endpointSpawn = async (
 
     const cmds = bin.split(" ");
     childProcess = spawn(cmds[0], cmds.slice(1), {
-      detached,
+      detached: true,
       stdio: "inherit",
       env: env,
     });
@@ -134,7 +120,6 @@ export const endpointSpawn = async (
     childProcess,
     bin,
     endpoint,
-    wsEndpoint: wsify(endpoint),
   };
 };
 
@@ -175,9 +160,12 @@ export const endpointProxy = async ({
   endpoint,
   event,
   deadline,
+  wsProxy,
 }: EndpointProxyRequest): Promise<EndpointResponse> => {
   const rawEvent = JSON.parse(event) as Partial<APIGatewayProxyEventV2>;
 
+  log("!!! Received event", { requestId, rawEvent });
+
   const {
     requestContext,
     rawPath,
@@ -202,6 +190,10 @@ export const endpointProxy = async ({
     );
   }
 
+  if (wsProxy) {
+    wsProxy.init();
+  }
+
   if (!rawPath) {
     throw new Error("No path found in event");
   }

Diff for: src/internal/events.ts

@@ -2,11 +2,13 @@ import { log } from "./log";
 import { EndpointProxyRequest, EndpointExecRequest } from "./types";
 import { endpointExec, endpointProxy } from "./endpoints";
 import { getRuntimeEvent, postRuntimeEventResponse } from "./runtime";
+import { WebsocketProxy } from "./websocket";
 
 export const pollForEvents = async (
   runtimeApi: string,
   bin?: string,
-  endpoint?: URL
+  endpoint?: URL,
+  wsProxy?: WebsocketProxy
 ): Promise<void> => {
   log("Waiting for next event from Lambda Runtime API", { runtimeApi });
 
@@ -15,7 +17,7 @@ export const pollForEvents = async (
   let payload: any | undefined = undefined;
 
   if (bin && !endpoint) {
-    log("No endpoint specified, executing bin", { bin });
+    log("Executing bin", { bin });
 
     const request: EndpointExecRequest = {
       requestId,
@@ -26,20 +28,21 @@ export const pollForEvents = async (
 
     payload = (await endpointExec(request)).payload;
 
-    log("Bin execution complete", { bin });
+    log("Bin execution complete");
   } else if (endpoint) {
-    log("Endpoint specified, proxying request", { endpoint });
+    log("Proxying request", { endpoint });
 
     const request: EndpointProxyRequest = {
       requestId,
       endpoint,
       event,
       deadline,
+      wsProxy,
     };
 
     payload = (await endpointProxy(request)).payload;
 
-    log("Proxy request complete", { endpoint });
+    log("Proxy request complete");
   } else {
     throw new Error(
       `
@@ -57,5 +60,5 @@ Expected format: {bin}@{endpoint} or {bin} or {endpoint}:
 
   log("Response sent to Lambda Runtime API", { runtimeApi, requestId });
 
-  return pollForEvents(runtimeApi, bin, endpoint);
+  return pollForEvents(runtimeApi, bin, endpoint, wsProxy);
 };

Diff for: src/internal/types.ts

@@ -1,11 +1,11 @@
 import { APIGatewayProxyResult } from "aws-lambda";
 import { ChildProcess } from "child_process";
+import { WebsocketProxy } from "./websocket";
 
 export type SpawnResult = {
   childProcess?: ChildProcess;
   bin?: string;
   endpoint?: URL;
-  wsEndpoint?: URL;
 };
 
 export type RuntimeEvent = {
@@ -26,6 +26,7 @@ export type EndpointProxyRequest = {
   endpoint: URL;
   event: string;
   deadline: number;
+  wsProxy?: WebsocketProxy;
 };
 
 export type EndpointResponse = {

Diff for: src/internal/websocket.ts

@@ -0,0 +1,75 @@
+import WebSocket from "ws";
+import { log } from "./log";
+
+function wsify(url?: string | URL): URL | undefined {
+  if (!url) return undefined;
+  if (typeof url === "string") url = new URL(url);
+
+  const wsUrl = new URL(url.toString());
+  wsUrl.protocol = url.protocol.replace("http", "ws");
+
+  if (!wsUrl.protocol.startsWith("ws")) {
+    return undefined;
+  }
+
+  return wsUrl;
+}
+
+export class WebsocketProxy {
+  private ws?: WebSocket;
+  public readonly url?: URL;
+  constructor(url?: URL, private route?: string) {
+    this.url = wsify(url);
+  }
+
+  init() {
+    const { url, route, ws } = this;
+
+    if (!url || !route || ws) {
+      log("Skipping Websocket Initialization", {
+        url,
+        route,
+        initialized: !!ws,
+      });
+      return;
+    }
+
+    this.ws = new WebSocket(url);
+    this.ws.on("open", this.handleOpen.bind(this));
+    this.ws.on("message", this.handleMessage.bind(this));
+    this.ws.on("error", this.handleError.bind(this));
+    this.ws.on("ping", this.handlePing.bind(this));
+    this.ws.on("pong", this.handlePong.bind(this));
+  }
+
+  send(data: any) {
+    log("Sending message", { data });
+    if (!this.ws) {
+      throw new Error("WebSocket is not initialized");
+    }
+    this.ws.send(data);
+  }
+
+  handleOpen() {
+    log("Connected to WebSocket", {
+      url: this.url,
+      route: this.route,
+    });
+  }
+
+  handlePing(data: Buffer) {
+    log("Received Ping", { data });
+  }
+
+  handlePong(data: Buffer) {
+    log("Received Pong", { data });
+  }
+
+  handleMessage(data: any) {
+    log("Received Message", { data });
+  }
+
+  handleError(error: Error) {
+    log("WebSocket error", { error });
+  }
+}

Diff for: yarn.lock

@@ -142,6 +142,13 @@
   resolved "https://registry.yarnpkg.com/@types/which/-/which-3.0.3.tgz#41142ed5a4743128f1bc0b69c46890f0453ddb89"
   integrity sha512-2C1+XoY0huExTbs8MQv1DuS5FS86+SEjdM9F/+GS61gg5Hqbtj8ZiDSx8MfWcyei907fIPbfPGCOrNUTnVHY1g==
 
+"@types/ws@^8.5.10":
+  version "8.5.10"
+  resolved "https://registry.yarnpkg.com/@types/ws/-/ws-8.5.10.tgz#4acfb517970853fa6574a3a6886791d04a396787"
+  integrity sha512-vmQSUcfalpIq0R9q7uTo2lXs6eGIpt9wtnLdMv9LVpIjCA/+ufZRozlVoVelIYixx1ugCBKDhn89vnsEGOCx9A==
+  dependencies:
+    "@types/node" "*"
+
 "@webassemblyjs/[email protected]", "@webassemblyjs/ast@^1.12.1":
   version "1.12.1"
   resolved "https://registry.yarnpkg.com/@webassemblyjs/ast/-/ast-1.12.1.tgz#bb16a0e8b1914f979f45864c23819cc3e3f0d4bb"
@@ -1472,6 +1479,11 @@ wildcard@^2.0.0:
   resolved "https://registry.yarnpkg.com/wildcard/-/wildcard-2.0.1.tgz#5ab10d02487198954836b6349f74fff961e10f67"
   integrity sha512-CC1bOL87PIWSBhDcTrdeLo6eGT7mCFtrg0uIJtqJUFyK+eJnzl8A1niH56uu7KMa5XFrtiV+AQuHO3n7DsHnLQ==
 
+ws@^8.17.0:
+  version "8.17.0"
+  resolved "https://registry.yarnpkg.com/ws/-/ws-8.17.0.tgz#d145d18eca2ed25aaf791a183903f7be5e295fea"
+  integrity sha512-uJq6108EgZMAl20KagGkzCKfMEjxmKvZHG7Tlq0Z6nOky7YF7aq4mOx6xK8TJ/i1LeK4Qus7INktacctDgY8Ow==
+
 [email protected]:
   version "0.6.2"
   resolved "https://registry.yarnpkg.com/xml2js/-/xml2js-0.6.2.tgz#dd0b630083aa09c161e25a4d0901e2b2a929b499"