Make LensProxy more injectable

Signed-off-by: Sebastian Malton <sebastian@malton.name>

src/common/terminal/channels.ts

@@ -3,6 +3,10 @@
  * Licensed under MIT License. See LICENSE in root directory for more information.
  */
 
+import type { TabId } from "../../renderer/components/dock/dock/store";
+import type { ClusterId } from "../cluster-types";
+import type { RequestChannel } from "../utils/channel/request-channel-listener-injection-token";
+
 
 export enum TerminalChannels {
   STDIN = "stdin",
@@ -29,3 +33,12 @@ export type TerminalMessage = {
 } | {
   type: TerminalChannels.PING;
 };
+
+export interface ClusterShellAuthenticationArgs {
+  clusterId: ClusterId;
+  tabId: TabId;
+}
+
+export const clusterShellAuthenticationChannel: RequestChannel<ClusterShellAuthenticationArgs, Uint8Array> = {
+  id: "cluster-shell-authentication-request",
+};

src/main/cluster/manager.ts

@@ -4,13 +4,11 @@
  */
 
 import "../../common/ipc/cluster";
-import type http from "http";
 import type { ObservableSet } from "mobx";
 import { action, makeObservable, observable, observe, reaction, toJS } from "mobx";
 import type { Cluster } from "../../common/cluster/cluster";
 import logger from "../logger";
-import { apiKubePrefix } from "../../common/vars";
-import { getClusterIdFromHost, isErrnoException } from "../../common/utils";
+import { isErrnoException } from "../../common/utils";
 import type { KubernetesClusterPrometheusMetrics } from "../../common/catalog-entities/kubernetes-cluster";
 import { isKubernetesCluster, KubernetesCluster, LensKubernetesClusterStatus } from "../../common/catalog-entities/kubernetes-cluster";
 import { ipcMainOn } from "../../common/ipc";
@@ -260,27 +258,6 @@ export class ClusterManager {
       cluster.disconnect();
     });
   }
-
-  getClusterForRequest = (req: http.IncomingMessage): Cluster | undefined => {
-    if (!req.headers.host) {
-      return undefined;
-    }
-
-    // lens-server is connecting to 127.0.0.1:<port>/<uid>
-    if (req.url && req.headers.host.startsWith("127.0.0.1")) {
-      const clusterId = req.url.split("/")[1];
-      const cluster = this.dependencies.store.getById(clusterId);
-
-      if (cluster) {
-        // we need to swap path prefix so that request is proxied to kube api
-        req.url = req.url.replace(`/${clusterId}`, apiKubePrefix);
-      }
-
-      return cluster;
-    }
-
-    return this.dependencies.store.getById(getClusterIdFromHost(req.headers.host));
-  };
 }
 
 export function catalogEntityFromCluster(cluster: Cluster) {

src/main/lens-proxy/get-cluster-for-request.injectable.ts

@@ -0,0 +1,48 @@
+/**
+ * Copyright (c) OpenLens Authors. All rights reserved.
+ * Licensed under MIT License. See LICENSE in root directory for more information.
+ */
+import { getInjectable } from "@ogre-tools/injectable";
+import getClusterByIdInjectable from "../../common/cluster-store/get-by-id.injectable";
+import type { Cluster } from "../../common/cluster/cluster";
+import { getClusterIdFromHost } from "../../common/utils";
+import { apiKubePrefix } from "../../common/vars";
+import type { ServerIncomingMessage } from "./lens-proxy";
+
+export type GetClusterForRequest = (req: ServerIncomingMessage) => Cluster | undefined;
+
+const getClusterForRequestInjectable = getInjectable({
+  id: "get-cluster-for-request",
+  instantiate: (di): GetClusterForRequest => {
+    const getClusterById = di.inject(getClusterByIdInjectable);
+
+    return (req) => {
+      if (!req.headers.host) {
+        return undefined;
+      }
+
+      // lens-server is connecting to 127.0.0.1:<port>/<uid>
+      if (req.url && req.headers.host.startsWith("127.0.0.1")) {
+        const clusterId = req.url.split("/")[1];
+        const cluster = getClusterById(clusterId);
+
+        if (cluster) {
+        // we need to swap path prefix so that request is proxied to kube api
+          req.url = req.url.replace(`/${clusterId}`, apiKubePrefix);
+        }
+
+        return cluster;
+      }
+
+      const clusterId = getClusterIdFromHost(req.headers.host);
+
+      if (!clusterId) {
+        return undefined;
+      }
+
+      return getClusterById(clusterId);
+    };
+  },
+});
+
+export default getClusterForRequestInjectable;

src/main/lens-proxy/lens-proxy.injectable.ts

@@ -4,15 +4,15 @@
  */
 import { getInjectable } from "@ogre-tools/injectable";
 import { LensProxy } from "./lens-proxy";
-import { kubeApiUpgradeRequest } from "./proxy-functions";
 import routerInjectable from "../router/router.injectable";
 import httpProxy from "http-proxy";
-import clusterManagerInjectable from "../cluster/manager.injectable";
 import shellApiRequestInjectable from "./proxy-functions/shell-api-request/shell-api-request.injectable";
 import lensProxyPortInjectable from "./lens-proxy-port.injectable";
 import contentSecurityPolicyInjectable from "../../common/vars/content-security-policy.injectable";
 import emitAppEventInjectable from "../../common/app-event-bus/emit-event.injectable";
 import loggerInjectable from "../../common/logger.injectable";
+import getClusterForRequestInjectable from "./get-cluster-for-request.injectable";
+import { kubeApiUpgradeRequest } from "./proxy-functions/kube-api-upgrade-request";
 
 const lensProxyInjectable = getInjectable({
   id: "lens-proxy",
@@ -20,9 +20,9 @@ const lensProxyInjectable = getInjectable({
   instantiate: (di) => new LensProxy({
     router: di.inject(routerInjectable),
     proxy: httpProxy.createProxy(),
-    kubeApiUpgradeRequest,
-    shellApiRequest: di.inject(shellApiRequestInjectable),
-    getClusterForRequest: di.inject(clusterManagerInjectable).getClusterForRequest,
+    kubeApiUpgradeRequestHandler: kubeApiUpgradeRequest,
+    shellApiRequestHandler: di.inject(shellApiRequestInjectable),
+    getClusterForRequest: di.inject(getClusterForRequestInjectable),
     lensProxyPort: di.inject(lensProxyPortInjectable),
     contentSecurityPolicy: di.inject(contentSecurityPolicyInjectable),
     emitAppEvent: di.inject(emitAppEventInjectable),

src/main/lens-proxy/lens-proxy.ts

@@ -10,22 +10,29 @@ import type httpProxy from "http-proxy";
 import { apiPrefix, apiKubePrefix } from "../../common/vars";
 import type { Router } from "../router/router";
 import type { ClusterContextHandler } from "../context-handler/context-handler";
-import type { Cluster } from "../../common/cluster/cluster";
-import type { ProxyApiRequestArgs } from "./proxy-functions";
 import { getBoolean } from "../utils/parse-query";
 import assert from "assert";
 import type { SetRequired } from "type-fest";
 import type { EmitAppEvent } from "../../common/app-event-bus/emit-event.injectable";
 import type { Logger } from "../../common/logger";
-
-type GetClusterForRequest = (req: http.IncomingMessage) => Cluster | undefined;
+import type { GetClusterForRequest } from "./get-cluster-for-request.injectable";
+import type { Cluster } from "../../common/cluster/cluster";
 
 export type ServerIncomingMessage = SetRequired<http.IncomingMessage, "url" | "method">;
 
+export interface ProxyApiRequestArgs {
+  req: SetRequired<http.IncomingMessage, "url" | "method">;
+  socket: net.Socket;
+  head: Buffer;
+  cluster: Cluster;
+}
+
+export type ProxyApiRequestHandler = (args: ProxyApiRequestArgs) => Promise<void> | void;
+
 interface Dependencies {
   getClusterForRequest: GetClusterForRequest;
-  shellApiRequest: (args: ProxyApiRequestArgs) => void | Promise<void>;
-  kubeApiUpgradeRequest: (args: ProxyApiRequestArgs) => void | Promise<void>;
+  shellApiRequestHandler: ProxyApiRequestHandler;
+  kubeApiUpgradeRequestHandler: ProxyApiRequestHandler;
   emitAppEvent: EmitAppEvent;
   readonly router: Router;
   readonly proxy: httpProxy;
@@ -86,11 +93,18 @@ export class LensProxy {
           this.dependencies.logger.error(`[LENS-PROXY]: Could not find cluster for upgrade request from url=${req.url}`);
           socket.destroy();
         } else {
-          const isInternal = req.url.startsWith(`${apiPrefix}?`);
-          const reqHandler = isInternal ? dependencies.shellApiRequest : dependencies.kubeApiUpgradeRequest;
-
-          (async () => reqHandler({ req, socket, head, cluster }))()
-            .catch(error => this.dependencies.logger.error("[LENS-PROXY]: failed to handle proxy upgrade", error));
+          (async () => {
+            try {
+              if (req.url.startsWith(`${apiPrefix}?`)) {
+                // internal request
+                await this.dependencies.shellApiRequestHandler({ req, socket, head, cluster });
+              } else {
+                await this.dependencies.kubeApiUpgradeRequestHandler({ req, socket, head, cluster });
+              }
+            } catch (error) {
+              this.dependencies.logger.error("[LENS-PROXY]: failed to handle proxy upgrade", error);
+            }
+          })();
         }
       });
   }

src/main/lens-proxy/proxy-functions/index.ts

@@ -1,6 +0,0 @@
-/**
- * Copyright (c) OpenLens Authors. All rights reserved.
- * Licensed under MIT License. See LICENSE in root directory for more information.
- */
-export * from "./kube-api-upgrade-request";
-export * from "./types";

src/main/lens-proxy/proxy-functions/kube-api-upgrade-request.ts

@@ -8,11 +8,11 @@ import type { ConnectionOptions } from "tls";
 import { connect } from "tls";
 import url from "url";
 import { apiKubePrefix } from "../../../common/vars";
-import type { ProxyApiRequestArgs } from "./types";
+import type { ProxyApiRequestHandler } from "../lens-proxy";
 
 const skipRawHeaders = new Set(["Host", "Authorization"]);
 
-export async function kubeApiUpgradeRequest({ req, socket, head, cluster }: ProxyApiRequestArgs) {
+export const kubeApiUpgradeRequest: ProxyApiRequestHandler = async ({ req, socket, head, cluster }) => {
   const proxyUrl = await cluster.contextHandler.resolveAuthProxyUrl() + req.url.replace(apiKubePrefix, "");
   const proxyCa = cluster.contextHandler.resolveAuthProxyCa();
   const apiUrl = url.parse(cluster.apiUrl);
@@ -24,18 +24,14 @@ export async function kubeApiUpgradeRequest({ req, socket, head, cluster }: Prox
   };
 
   const proxySocket = connect(connectOpts, () => {
+    const headers = chunk(req.rawHeaders, 2)
+      .filter(([key]) => !skipRawHeaders.has(key))
+      .map(([key, value]) => `${key}: ${value}`)
+      .join("\r\n");
+
     proxySocket.write(`${req.method} ${pUrl.path} HTTP/1.1\r\n`);
     proxySocket.write(`Host: ${apiUrl.host}\r\n`);
-
-    for (const [key, value] of chunk(req.rawHeaders, 2)) {
-      if (skipRawHeaders.has(key)) {
-        continue;
-      }
-
-      proxySocket.write(`${key}: ${value}\r\n`);
-    }
-
-    proxySocket.write("\r\n");
+    proxySocket.write(`${headers}\r\n`);
     proxySocket.write(head);
   });
 
@@ -44,23 +40,13 @@ export async function kubeApiUpgradeRequest({ req, socket, head, cluster }: Prox
   proxySocket.setTimeout(0);
   socket.setTimeout(0);
 
-  proxySocket.on("data", function (chunk) {
-    socket.write(chunk);
-  });
-  proxySocket.on("end", function () {
-    socket.end();
-  });
-  proxySocket.on("error", function () {
+  proxySocket.on("data", chunk => socket.write(chunk));
+  proxySocket.on("end", () => socket.end());
+  proxySocket.on("error", () => {
     socket.write(`HTTP/${req.httpVersion} 500 Connection error\r\n\r\n`);
     socket.end();
   });
-  socket.on("data", function (chunk) {
-    proxySocket.write(chunk);
-  });
-  socket.on("end", function () {
-    proxySocket.end();
-  });
-  socket.on("error", function () {
-    proxySocket.end();
-  });
-}
+  socket.on("data", (chunk) => proxySocket.write(chunk));
+  socket.on("end", () => proxySocket.end());
+  socket.on("error", () => proxySocket.end());
+};

src/main/lens-proxy/proxy-functions/shell-api-request/authenticate.injectable.ts

@@ -0,0 +1,39 @@
+/**
+ * Copyright (c) OpenLens Authors. All rights reserved.
+ * Licensed under MIT License. See LICENSE in root directory for more information.
+ */
+import { getInjectable } from "@ogre-tools/injectable";
+import type { ClusterId } from "../../../../common/cluster-types";
+import shellRequestAuthTokensInjectable from "./shell-request-auth-tokens.injectable";
+import crypto from "crypto";
+
+export type AuthenticateShellRequest = (clusterId: ClusterId, tabId: string, token: string | null) => boolean;
+
+const authenticateShellRequestInjectable = getInjectable({
+  id: "authenticate-shell-request",
+  instantiate: (di): AuthenticateShellRequest => {
+    const shellRequestAuthTokens = di.inject(shellRequestAuthTokensInjectable);
+
+    return (clusterId, tabId, token) => {
+      const clusterTokens = shellRequestAuthTokens.get(clusterId);
+
+      if (!clusterTokens || !tabId || !token) {
+        return false;
+      }
+
+      const authToken = clusterTokens.get(tabId);
+      const buf = Uint8Array.from(Buffer.from(token, "base64"));
+
+      if (authToken instanceof Uint8Array && authToken.length === buf.length && crypto.timingSafeEqual(authToken, buf)) {
+        // remove the token because it is a single use token
+        clusterTokens.delete(tabId);
+
+        return true;
+      }
+
+      return false;
+    };
+  },
+});
+
+export default authenticateShellRequestInjectable;

src/main/lens-proxy/proxy-functions/shell-api-request/handler.injectable.ts

@@ -0,0 +1,28 @@
+/**
+ * Copyright (c) OpenLens Authors. All rights reserved.
+ * Licensed under MIT License. See LICENSE in root directory for more information.
+ */
+import { promisify } from "util";
+import { clusterShellAuthenticationChannel } from "../../../../common/terminal/channels";
+import { getRequestChannelListenerInjectable } from "../../../utils/channel/channel-listeners/listener-tokens";
+import shellRequestAuthTokensInjectable from "./shell-request-auth-tokens.injectable";
+import crypto from "crypto";
+import { getOrInsertMap, put } from "../../../../common/utils";
+
+const randomBytes = promisify(crypto.randomBytes);
+
+const clusterShellAuthenticationRequestHandlerInjectable = getRequestChannelListenerInjectable({
+  channel: clusterShellAuthenticationChannel,
+  handler: (di) => {
+    const shellRequestAuthTokens = di.inject(shellRequestAuthTokensInjectable);
+
+    return async ({ clusterId, tabId }) => {
+      const authToken = Uint8Array.from(await randomBytes(128));
+      const forCluster = getOrInsertMap(shellRequestAuthTokens, clusterId);
+
+      return put(forCluster, tabId, authToken);
+    };
+  },
+});
+
+export default clusterShellAuthenticationRequestHandlerInjectable;

src/main/lens-proxy/proxy-functions/shell-api-request/shell-api-request.injectable.ts

@@ -3,19 +3,42 @@
  * Licensed under MIT License. See LICENSE in root directory for more information.
  */
 import { getInjectable } from "@ogre-tools/injectable";
-import { shellApiRequest } from "./shell-api-request";
-import shellRequestAuthenticatorInjectable from "./shell-request-authenticator/shell-request-authenticator.injectable";
-import clusterManagerInjectable from "../../../cluster/manager.injectable";
 import openShellSessionInjectable from "../../../shell-session/create-shell-session.injectable";
+import getClusterForRequestInjectable from "../../get-cluster-for-request.injectable";
+import { WebSocketServer } from "ws";
+import { URL } from "url";
+import loggerInjectable from "../../../../common/logger.injectable";
+import type { ProxyApiRequestHandler } from "../../lens-proxy";
+import authenticateShellRequestInjectable from "./authenticate.injectable";
 
 const shellApiRequestInjectable = getInjectable({
   id: "shell-api-request",
 
-  instantiate: (di) => shellApiRequest({
-    openShellSession: di.inject(openShellSessionInjectable),
-    authenticateRequest: di.inject(shellRequestAuthenticatorInjectable).authenticate,
-    clusterManager: di.inject(clusterManagerInjectable),
-  }),
+  instantiate: (di): ProxyApiRequestHandler => {
+    const openShellSession = di.inject(openShellSessionInjectable);
+    const authenticateShellRequest = di.inject(authenticateShellRequestInjectable);
+    const getClusterForRequest = di.inject(getClusterForRequestInjectable);
+    const logger = di.inject(loggerInjectable);
+
+    return ({ req, socket, head }) => {
+      const cluster = getClusterForRequest(req);
+      const { searchParams } = new URL(req.url, "http://localhost");
+      const nodeName = searchParams.get("node");
+      const shellToken = searchParams.get("shellToken");
+      const tabId = searchParams.get("id");
+
+      if (!tabId || !cluster || !authenticateShellRequest(cluster.id, tabId, shellToken)) {
+        socket.write("Invalid shell request");
+        socket.end();
+      } else {
+        new WebSocketServer({ noServer: true })
+          .handleUpgrade(req, socket, head, (websocket) => {
+            openShellSession({ websocket, cluster, tabId, nodeName })
+              .catch(error => logger.error(`[SHELL-SESSION]: failed to open a ${nodeName ? "node" : "local"} shell`, error));
+          });
+      }
+    };
+  },
 });
 
 export default shellApiRequestInjectable;

src/main/lens-proxy/proxy-functions/shell-api-request/shell-api-request.ts

@@ -1,36 +0,0 @@
-/**
- * Copyright (c) OpenLens Authors. All rights reserved.
- * Licensed under MIT License. See LICENSE in root directory for more information.
- */
-
-import logger from "../../../logger";
-import { Server as WebSocketServer } from "ws";
-import type { ProxyApiRequestArgs } from "../types";
-import type { ClusterManager } from "../../../cluster/manager";
-import URLParse from "url-parse";
-import type { ClusterId } from "../../../../common/cluster-types";
-import type { OpenShellSession } from "../../../shell-session/create-shell-session.injectable";
-
-interface Dependencies {
-  authenticateRequest: (clusterId: ClusterId, tabId: string, shellToken: string | undefined) => boolean;
-  openShellSession: OpenShellSession;
-  clusterManager: ClusterManager;
-}
-
-export const shellApiRequest = ({ openShellSession, authenticateRequest, clusterManager }: Dependencies) => ({ req, socket, head }: ProxyApiRequestArgs): void => {
-  const cluster = clusterManager.getClusterForRequest(req);
-  const { query: { node: nodeName, shellToken, id: tabId }} = new URLParse(req.url, true);
-
-  if (!tabId || !cluster || !authenticateRequest(cluster.id, tabId, shellToken)) {
-    socket.write("Invalid shell request");
-
-    return void socket.end();
-  }
-
-  const ws = new WebSocketServer({ noServer: true });
-
-  ws.handleUpgrade(req, socket, head, (websocket) => {
-    openShellSession({ websocket, cluster, tabId, nodeName })
-      .catch(error => logger.error(`[SHELL-SESSION]: failed to open a ${nodeName ? "node" : "local"} shell`, error));
-  });
-};

src/main/lens-proxy/proxy-functions/shell-api-request/shell-request-auth-tokens.injectable.ts

@@ -0,0 +1,14 @@
+/**
+ * Copyright (c) OpenLens Authors. All rights reserved.
+ * Licensed under MIT License. See LICENSE in root directory for more information.
+ */
+import { getInjectable } from "@ogre-tools/injectable";
+import type { ClusterId } from "../../../../common/cluster-types";
+import type { TabId } from "../../../../renderer/components/dock/dock/store";
+
+const shellRequestAuthTokensInjectable = getInjectable({
+  id: "shell-request-auth-tokens",
+  instantiate: () => new Map<ClusterId, Map<TabId, Uint8Array>>(),
+});
+
+export default shellRequestAuthTokensInjectable;

src/main/lens-proxy/proxy-functions/shell-api-request/shell-request-authenticator/shell-request-authenticator.injectable.ts

@@ -1,20 +0,0 @@
-/**
- * Copyright (c) OpenLens Authors. All rights reserved.
- * Licensed under MIT License. See LICENSE in root directory for more information.
- */
-import { getInjectable } from "@ogre-tools/injectable";
-import { ShellRequestAuthenticator } from "./shell-request-authenticator";
-
-const shellRequestAuthenticatorInjectable = getInjectable({
-  id: "shell-request-authenticator",
-
-  instantiate: () => {
-    const authenticator = new ShellRequestAuthenticator();
-
-    authenticator.init();
-
-    return authenticator;
-  },
-});
-
-export default shellRequestAuthenticatorInjectable;

src/main/lens-proxy/proxy-functions/shell-api-request/shell-request-authenticator/shell-request-authenticator.ts

@@ -1,53 +0,0 @@
-/**
- * Copyright (c) OpenLens Authors. All rights reserved.
- * Licensed under MIT License. See LICENSE in root directory for more information.
- */
-import { getOrInsertMap } from "../../../../../common/utils";
-import type { ClusterId } from "../../../../../common/cluster-types";
-import { ipcMainHandle } from "../../../../../common/ipc";
-import crypto from "crypto";
-import { promisify } from "util";
-
-const randomBytes = promisify(crypto.randomBytes);
-
-export class ShellRequestAuthenticator {
-  private tokens = new Map<ClusterId, Map<string, Uint8Array>>();
-
-  init() {
-    ipcMainHandle("cluster:shell-api", async (event, clusterId, tabId) => {
-      const authToken = Uint8Array.from(await randomBytes(128));
-      const forCluster = getOrInsertMap(this.tokens, clusterId);
-
-      forCluster.set(tabId, authToken);
-
-      return authToken;
-    });
-  }
-
-  /**
-   * Authenticates a single use token for creating a new shell
-   * @param clusterId The `ClusterId` for the shell
-   * @param tabId The ID for the shell
-   * @param token The value that is being presented as a one time authentication token
-   * @returns `true` if `token` was valid, false otherwise
-   */
-  authenticate = (clusterId: ClusterId, tabId: string, token: string | undefined): boolean => {
-    const clusterTokens = this.tokens.get(clusterId);
-
-    if (!clusterTokens || !tabId || !token) {
-      return false;
-    }
-
-    const authToken = clusterTokens.get(tabId);
-    const buf = Uint8Array.from(Buffer.from(token, "base64"));
-
-    if (authToken instanceof Uint8Array && authToken.length === buf.length && crypto.timingSafeEqual(authToken, buf)) {
-      // remove the token because it is a single use token
-      clusterTokens.delete(tabId);
-
-      return true;
-    }
-
-    return false;
-  };
-}

src/main/lens-proxy/proxy-functions/types.ts

@@ -1,16 +0,0 @@
-/**
- * Copyright (c) OpenLens Authors. All rights reserved.
- * Licensed under MIT License. See LICENSE in root directory for more information.
- */
-
-import type http from "http";
-import type net from "net";
-import type { SetRequired } from "type-fest";
-import type { Cluster } from "../../../common/cluster/cluster";
-
-export interface ProxyApiRequestArgs {
-  req: SetRequired<http.IncomingMessage, "url" | "method">;
-  socket: net.Socket;
-  head: Buffer;
-  cluster: Cluster;
-}

src/main/shell-session/create-shell-session.injectable.ts

@@ -12,7 +12,7 @@ export interface OpenShellSessionArgs {
   websocket: WebSocket;
   cluster: Cluster;
   tabId: string;
-  nodeName?: string;
+  nodeName: string | null;
 }
 
 export type OpenShellSession = (args: OpenShellSessionArgs) => Promise<void>;