From 4e264d3c40b4acaf47afd8e87d506b274d491a84 Mon Sep 17 00:00:00 2001
From: Sebastian Malton <sebastian@malton.name>
Date: Mon, 8 Nov 2021 09:55:18 -0500
Subject: [PATCH] Block shell websocket connections from non-lens renderers

Signed-off-by: Sebastian Malton <sebastian@malton.name>
---
 src/main/lens-proxy.ts                        |  2 +-
 src/main/proxy-functions/shell-api-request.ts | 35 ++++++++++++-------
 src/main/shell-session/node-shell-session.ts  |  2 +-
 src/main/shell-session/shell-session.ts       |  2 +-
 4 files changed, 25 insertions(+), 16 deletions(-)

diff --git a/src/main/lens-proxy.ts b/src/main/lens-proxy.ts
index a3bee528cd..d22981ddf4 100644
--- a/src/main/lens-proxy.ts
+++ b/src/main/lens-proxy.ts
@@ -82,7 +82,7 @@ export class LensProxy extends Singleton {
         const reqHandler = isInternal ? shellApiRequest : kubeApiRequest;
 
         (async () => reqHandler({ req, socket, head }))()
-          .catch(error => logger.error(logger.error(`[LENS-PROXY]: failed to handle proxy upgrade: ${error}`)));
+          .catch(error => logger.error("[LENS-PROXY]: failed to handle proxy upgrade", error));
       });
   }
 
diff --git a/src/main/proxy-functions/shell-api-request.ts b/src/main/proxy-functions/shell-api-request.ts
index 3edaecff81..4ad75cd94e 100644
--- a/src/main/proxy-functions/shell-api-request.ts
+++ b/src/main/proxy-functions/shell-api-request.ts
@@ -19,29 +19,38 @@
  * CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
  */
 
-import type http from "http";
 import url from "url";
 import logger from "../logger";
-import * as WebSocket from "ws";
+import { Server as WebSocketServer } from "ws";
 import { NodeShellSession, LocalShellSession } from "../shell-session";
 import type { ProxyApiRequestArgs } from "./types";
 import { ClusterManager } from "../cluster-manager";
+import { appName, appSemVer } from "../../common/vars";
+import { LensProxy } from "../lens-proxy";
 
-export function shellApiRequest({ req, socket, head }: ProxyApiRequestArgs) {
-  const ws = new WebSocket.Server({ noServer: true });
+export function shellApiRequest({ req, socket, head }: ProxyApiRequestArgs): void {
+  const cluster = ClusterManager.getInstance().getClusterForRequest(req);
 
-  ws.on("connection", ((socket: WebSocket, req: http.IncomingMessage) => {
-    const cluster = ClusterManager.getInstance().getClusterForRequest(req);
+  if (
+    socket.remoteAddress !== "127.0.0.1"
+    || !req.headers["user-agent"]?.includes(` ${appName}/${appSemVer} `)
+    || !req.headers.origin?.endsWith(`.localhost:${LensProxy.getInstance().port}`)
+    || !cluster
+  ) {
+    socket.write("Invalid shell request");
+
+    return void socket.end();
+  }
+
+  const ws = new WebSocketServer({ noServer: true });
+
+  ws.handleUpgrade(req, socket, head, (webSocket) => {
     const nodeParam = url.parse(req.url, true).query["node"]?.toString();
     const shell = nodeParam
-      ? new NodeShellSession(socket, cluster, nodeParam)
-      : new LocalShellSession(socket, cluster);
+      ? new NodeShellSession(webSocket, cluster, nodeParam)
+      : new LocalShellSession(webSocket, cluster);
 
     shell.open()
-      .catch(error => logger.error(`[SHELL-SESSION]: failed to open: ${error}`, { error }));
-  }));
-
-  ws.handleUpgrade(req, socket, head, (con) => {
-    ws.emit("connection", con, req);
+      .catch(error => logger.error(`[SHELL-SESSION]: failed to open a ${nodeParam ? "node" : "local"} shell`, error));
   });
 }
diff --git a/src/main/shell-session/node-shell-session.ts b/src/main/shell-session/node-shell-session.ts
index 494137de7d..db0bb89bd3 100644
--- a/src/main/shell-session/node-shell-session.ts
+++ b/src/main/shell-session/node-shell-session.ts
@@ -19,7 +19,7 @@
  * CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
  */
 
-import type * as WebSocket from "ws";
+import type WebSocket from "ws";
 import { v4 as uuid } from "uuid";
 import * as k8s from "@kubernetes/client-node";
 import type { KubeConfig } from "@kubernetes/client-node";
diff --git a/src/main/shell-session/shell-session.ts b/src/main/shell-session/shell-session.ts
index 987a00afdd..ba05a22dc5 100644
--- a/src/main/shell-session/shell-session.ts
+++ b/src/main/shell-session/shell-session.ts
@@ -22,7 +22,7 @@
 import fse from "fs-extra";
 import type { Cluster } from "../cluster";
 import { Kubectl } from "../kubectl";
-import type * as WebSocket from "ws";
+import type WebSocket from "ws";
 import { shellEnv } from "../utils/shell-env";
 import { app } from "electron";
 import { clearKubeconfigEnvVars } from "../utils/clear-kube-env-vars";