feat: websocket server proxying to provider websocket

refs #1024

Author: jzsfkzm

apps/api/package.json

@@ -57,6 +57,7 @@
     "@cosmjs/tendermint-rpc": "^0.32.4",
     "@dotenvx/dotenvx": "^1.9.0",
     "@hono/node-server": "1.13.7",
+    "@hono/node-ws": "^1.2.0",
     "@hono/otel": "^0.1.0",
     "@hono/swagger-ui": "0.4.1",
     "@hono/zod-openapi": "0.18.4",
@@ -90,6 +91,7 @@
     "markdown-to-txt": "^2.0.1",
     "memory-cache": "^0.2.0",
     "murmurhash": "^2.0.1",
+    "node-forge": "^1.3.1",
     "pg": "^8.12.0",
     "pg-hstore": "^2.3.4",
     "postgres": "^3.4.4",
@@ -125,6 +127,7 @@
     "@types/pg": "^8.11.6",
     "@types/semver": "^7.5.2",
     "@types/uuid": "^8.3.1",
+    "@types/ws": "^8.18.1",
     "@typescript-eslint/eslint-plugin": "^7.12.0",
     "alias-hq": "^5.1.6",
     "copy-webpack-plugin": "^12.0.2",
@@ -148,6 +151,7 @@
     "ts-loader": "^9.5.2",
     "type-fest": "^4.26.1",
     "typescript": "~5.8.2",
+    "wait-for-expect": "^4.0.0",
     "webpack": "^5.91.0",
     "webpack-cli": "4.10.0",
     "webpack-node-externals": "^3.0.0"

apps/api/src/app.ts

@@ -3,11 +3,13 @@ import "reflect-metadata";
 import { LoggerService } from "@akashnetwork/logging";
 import { HttpLoggerIntercepter } from "@akashnetwork/logging/hono";
 import { serve } from "@hono/node-server";
+import { createNodeWebSocket } from "@hono/node-ws";
 import { otel } from "@hono/otel";
 import { swaggerUI } from "@hono/swagger-ui";
 import { Hono } from "hono";
 import { cors } from "hono/cors";
 import once from "lodash/once";
+import type { AddressInfo } from "net";
 import { container } from "tsyringe";
 
 import { AuthInterceptor } from "@src/auth/services/auth.interceptor";
@@ -38,6 +40,7 @@ import { userRouter } from "./routers/userRouter";
 import { web3IndexRouter } from "./routers/web3indexRouter";
 import { env } from "./utils/env";
 import { bytesToHumanReadableSize } from "./utils/files";
+import { initLeaseWebsocketRoute } from "./websocket/routes/websocket/websocket.router";
 import { addressRouter } from "./address";
 import { sendVerificationEmailRouter } from "./auth";
 import {
@@ -160,6 +163,9 @@ for (const handler of openApiHonoHandlers) {
   appHono.route("/", handler);
 }
 
+const { injectWebSocket, upgradeWebSocket } = createNodeWebSocket({ app: appHono });
+initLeaseWebsocketRoute(appHono, upgradeWebSocket);
+
 appHono.route("/", notificationsApiProxy);
 
 appHono.route("/", healthzRouter);
@@ -196,22 +202,31 @@ const appLogger = LoggerService.forContext("APP");
  * Start scheduler
  * Start server
  */
-export async function initApp() {
+export async function initApp(port: number = Number(PORT)): Promise<AppServer> {
   try {
     await Promise.all([initDb(), container.resolve(FeatureFlagsService).initialize()]);
     startScheduler();
 
-    appLogger.info({ event: "SERVER_STARTING", url: `http://localhost:${PORT}`, NODE_OPTIONS: process.env.NODE_OPTIONS });
+    appLogger.info({ event: "SERVER_STARTING", url: `http://localhost:${port}`, NODE_OPTIONS: process.env.NODE_OPTIONS });
     const server = serve({
       fetch: appHono.fetch,
-      port: typeof PORT === "string" ? parseInt(PORT, 10) : PORT
+      port: typeof port === "string" ? parseInt(port, 10) : port
     });
+    injectWebSocket(server);
     const shutdown = once(() => shutdownServer(server, appLogger, container.dispose.bind(container)));
 
     process.on("SIGTERM", shutdown);
     process.on("SIGINT", shutdown);
+
+    return {
+      host: `http://localhost:${(server.address() as AddressInfo).port}`,
+      async close() {
+        await shutdown();
+      }
+    };
   } catch (error) {
     appLogger.error({ event: "APP_INIT_ERROR", error });
+    throw error;
   }
 }
 
@@ -236,3 +251,8 @@ export async function initDb() {
 }
 
 export { appHono as app };
+
+export interface AppServer {
+  host: string;
+  close(): Promise<void>;
+}

apps/api/src/core/lib/telemetry.ts

@@ -0,0 +1,12 @@
+import type { Span } from "@opentelemetry/api";
+import { context } from "@opentelemetry/api";
+import { trace } from "@opentelemetry/api";
+
+export function traceActiveSpan<T extends (span: Span) => any>(name: string, callback: T): ReturnType<T> {
+  return trace.getTracer("default").startActiveSpan(name, callback);
+}
+
+export function propagateTracingContext<T extends (...args: any[]) => any>(callback: T): T {
+  const currentContext = context.active();
+  return ((...args) => context.with(currentContext, () => callback(...args))) as T;
+}

apps/api/src/lib/shutdown-server/shutdown-server.ts

@@ -0,0 +1,17 @@
+export async function shutdownServer(server: ClosableServer | undefined): Promise<void> {
+  if (!server || !server.listening) return;
+  return new Promise<void>((resolve, reject) => {
+    server.close(error => {
+      if (error) {
+        reject(error);
+      } else {
+        resolve();
+      }
+    });
+  });
+}
+
+export interface ClosableServer {
+  close: (cb?: (error?: Error) => void) => void;
+  listening: boolean;
+}

apps/api/src/provider/services/jwt-token/jwt-token.service.spec.ts

@@ -82,12 +82,12 @@ describe("JwtTokenService", () => {
     });
   });
 
-  describe("getGranularLeases", () => {
+  describe("getScopedLeases", () => {
     it("returns leases for a provider, scoped", () => {
       const { jwtTokenService } = setup();
 
       const provider = createAkashAddress();
-      const result = jwtTokenService.getGranularLeases({ provider, scope: ["status"] });
+      const result = jwtTokenService.getScopedLeases({ provider, scope: ["status"] });
 
       expect(result).toEqual({
         access: "granular",
@@ -102,6 +102,25 @@ describe("JwtTokenService", () => {
     });
   });
 
+  describe("getFullAccessLeases", () => {
+    it("returns leases for a provider, with full access", () => {
+      const { jwtTokenService } = setup();
+
+      const provider = createAkashAddress();
+      const result = jwtTokenService.getFullAccessLeases({ provider });
+
+      expect(result).toEqual({
+        access: "granular",
+        permissions: [
+          {
+            provider,
+            access: "full"
+          }
+        ]
+      });
+    });
+  });
+
   function setup(): {
     jwtTokenService: JwtTokenService;
     mockBillingConfig: MockProxy<BillingConfig>;

apps/api/src/provider/services/jwt-token/jwt-token.service.ts

@@ -51,10 +51,17 @@ export class JwtTokenService {
     return { jwtToken, address: akashWallet.address };
   }
 
-  getGranularLeases({ provider, scope }: { provider: string; scope: AccessScope[] }): JwtTokenPayload["leases"] {
+  getScopedLeases({ provider, scope }: { provider: string; scope: AccessScope[] }): JwtTokenPayload["leases"] {
     return {
       access: "granular",
       permissions: [{ provider, access: "scoped", scope }]
     };
   }
+
+  getFullAccessLeases({ provider }: { provider: string }): JwtTokenPayload["leases"] {
+    return {
+      access: "granular",
+      permissions: [{ provider, access: "full" }]
+    };
+  }
 }

apps/api/src/provider/services/provider/provider.service.spec.ts

@@ -35,7 +35,7 @@ describe(ProviderService.name, () => {
 
       providerHttpService.getProvider.mockResolvedValue(mockProviderResponse);
       jwtTokenService.generateJwtToken.mockResolvedValue(jwtToken);
-      jwtTokenService.getGranularLeases.mockReturnValue(leases);
+      jwtTokenService.getScopedLeases.mockReturnValue(leases);
       providerHttpService.sendManifest.mockResolvedValue({ success: true });
 
       const result = await service.sendManifest({ provider: providerAddress, dseq, manifest, walletId });
@@ -187,7 +187,7 @@ describe(ProviderService.name, () => {
 
       providerHttpService.getProvider.mockResolvedValue(mockProviderResponse);
       jwtTokenService.generateJwtToken.mockResolvedValue(jwtToken);
-      jwtTokenService.getGranularLeases.mockReturnValue(leases);
+      jwtTokenService.getScopedLeases.mockReturnValue(leases);
       providerHttpService.getLeaseStatus.mockResolvedValue(mockLeaseStatus);
 
       const result = await service.getLeaseStatus(providerAddress, dseq, gseq, oseq, walletId);

apps/api/src/provider/services/provider/provider.service.ts

@@ -61,7 +61,7 @@ export class ProviderService {
       try {
         const jwtToken = await this.jwtTokenService.generateJwtToken({
           walletId,
-          leases: this.jwtTokenService.getGranularLeases({
+          leases: this.jwtTokenService.getScopedLeases({
             provider: providerIdentity.owner,
             scope: ["send-manifest"]
           })
@@ -89,7 +89,7 @@ export class ProviderService {
 
     const jwtToken = await this.jwtTokenService.generateJwtToken({
       walletId,
-      leases: this.jwtTokenService.getGranularLeases({
+      leases: this.jwtTokenService.getScopedLeases({
         provider,
         scope: ["status"]
       })

apps/api/src/websocket/controllers/websocket/websocket.controller.ts

@@ -0,0 +1,159 @@
+import { LoggerService } from "@akashnetwork/logging";
+import { Attributes, Span } from "@opentelemetry/api";
+import { WSContext } from "hono/ws";
+import { singleton } from "tsyringe";
+
+import { AuthService } from "@src/auth/services/auth.service";
+import { traceActiveSpan } from "@src/core/lib/telemetry";
+import { MESSAGE_SCHEMA, WsMessage } from "@src/websocket/http-schemas/websocket.schema";
+import { ProviderWebsocketService } from "@src/websocket/services/provider-websocket/provider-websocket.service";
+import { ClientWebSocketStats, WebSocketUsage } from "@src/websocket/services/websocket-stats/websocket-stats.service";
+
+const logger = LoggerService.forContext("LeaseWebsocketController");
+
+@singleton()
+export class WebsocketController {
+  constructor(
+    private readonly authService: AuthService,
+    private readonly providerWebsocketService: ProviderWebsocketService
+  ) {}
+
+  async handleOpen(stats: ClientWebSocketStats) {
+    const { currentUser, ability } = this.authService;
+
+    logger.info({
+      event: "WEBSOCKET_CONNECTION_OPENED",
+      userId: currentUser?.id,
+      hasAbility: !!ability
+    });
+
+    stats.setUserIfExists(currentUser, ability);
+  }
+
+  async handleMessage(messageStr: string, ws: WSContext, stats: ClientWebSocketStats) {
+    logger.debug({
+      event: "WEBSOCKET_MESSAGE_RECEIVED",
+      messageLength: messageStr?.length
+    });
+
+    const userInfo = stats.getUser();
+    if (userInfo) {
+      logger.debug({
+        event: "WEBSOCKET_USER_INFO_RETRIEVED",
+        userId: userInfo.currentUser?.id
+      });
+    }
+
+    traceActiveSpan("ws.message", async span => {
+      let message: WsMessage | undefined;
+      try {
+        message = typeof messageStr === "string" ? JSON.parse(messageStr) : undefined;
+      } catch (error) {
+        logger.error({
+          event: "CLIENT_MESSAGE_INVALID_JSON",
+          message: "Received message is not a JSON string",
+          messageLength: messageStr?.length,
+          messageType: typeof messageStr,
+          messageConstructor: messageStr?.constructor?.name
+        });
+      }
+
+      if (!message) {
+        return ws.send(
+          JSON.stringify({
+            type: "websocket",
+            message: "Message is not a JSON string",
+            error: "Invalid message format"
+          })
+        );
+      }
+
+      const parsedMessage = MESSAGE_SCHEMA.safeParse(message);
+      if (parsedMessage.error) {
+        logger.error({
+          event: "CLIENT_MESSAGE_INVALID_JSON",
+          message: "Message doesn't match expected schema",
+          error: parsedMessage.error
+        });
+        return ws.send(
+          JSON.stringify({
+            type: "websocket",
+            message: "Message doesn't match expected schema",
+            error: "Invalid message format"
+          })
+        );
+      }
+
+      const attributes: Attributes = {
+        type: message.type
+      };
+      if (message.type === "websocket") {
+        attributes.providerUrl = message.url;
+        attributes.providerAddress = message.providerAddress;
+        attributes.chainNetwork = message.chainNetwork;
+        attributes.function = getWebSocketUsage(message);
+      }
+
+      span.setAttributes(attributes);
+      logger.info({
+        event: "NEW_WEBSOCKET_MESSAGE",
+        attributes
+      });
+
+      stats.setUsage(getWebSocketUsage(message));
+
+      try {
+        if (message.type === "ping") {
+          ws.send(
+            JSON.stringify({
+              type: "pong"
+            })
+          );
+        } else if (message.type === "websocket") {
+          await this.providerWebsocketService.proxyMessageToProvider(message, ws, stats);
+        }
+      } catch (err) {
+        logger.error({
+          event: "CLIENT_MESSAGE_SEND_ERROR",
+          error: err
+        });
+        ws.send(
+          JSON.stringify({
+            id: message.id,
+            error: "Unable to send message to provider socket",
+            type: message.type
+          })
+        );
+      }
+
+      span.end();
+    });
+  }
+
+  async handleClose(_event: any, stats: ClientWebSocketStats, span?: Span) {
+    logger.info("Closing socket");
+    stats.close();
+
+    this.providerWebsocketService.closeProviderSocket(stats.id);
+
+    span?.end();
+  }
+
+  async handleError(event: Event) {
+    logger.error({
+      event: "CLIENT_WEBSOCKET_ERROR",
+      error: event
+    });
+  }
+}
+
+function getWebSocketUsage(message: any): WebSocketUsage {
+  if (message.type === "websocket") {
+    if (message.url.includes("logs?follow=false&tail=10000000")) return "DownloadLogs";
+    if (message.url.includes("logs?follow=true")) return "StreamLogs";
+    if (message.url.includes("kubeevents?follow=true")) return "StreamEvents";
+    if (message.url.includes("/shell?stdin=")) return "Shell";
+  }
+
+  return "Unknown";
+}