From 722e70f599b4c089d6b65177673bf08b98d22888 Mon Sep 17 00:00:00 2001
From: Siyuan Miao <i@xswan.net>
Date: Tue, 11 Feb 2025 13:48:03 +0100
Subject: [PATCH] feat: include device ip and ICE server list in device
 handshake payload

---
 .env.example   |  5 ++++-
 src/devices.ts |  5 +++--
 src/index.ts   |  4 ++++
 src/webrtc.ts  | 35 +++++++++++++++++++++++++++++------
 4 files changed, 40 insertions(+), 9 deletions(-)

diff --git a/.env.example b/.env.example
index 2f39949..86ac439 100644
--- a/.env.example
+++ b/.env.example
@@ -17,4 +17,7 @@ R2_SECRET_ACCESS_KEY=XXX    # Any S3 compatible secret access key
 R2_BUCKET=XXX               # Any S3 compatible bucket
 R2_CDN_URL=XXX              # Any S3 compatible CDN URL
 
-CORS_ORIGINS=https://app.jetkvm.com,http://localhost:5173 # Allowed CORS Origins, split by comma
\ No newline at end of file
+CORS_ORIGINS=https://app.jetkvm.com,http://localhost:5173 # Allowed CORS Origins, split by comma
+
+REAL_IP_HEADER=XXX          # Real IP Header for the reverse proxy (e.g. X-Real-IP), leave empty if not needed
+ICE_SERVERS=XXX             # ICE Servers for WebRTC, split by comma (e.g. stun:stun.l.google.com:19302,stun:stun1.l.google.com:19302)
\ No newline at end of file
diff --git a/src/devices.ts b/src/devices.ts
index 7aa6217..9a66247 100644
--- a/src/devices.ts
+++ b/src/devices.ts
@@ -119,8 +119,9 @@ export const Delete = async (req: express.Request, res: express.Response) => {
   await prisma.device.delete({ where: { id, user: { googleId: sub } } });
 
   // We just removed the device, so we should close any running open socket connections
-  const socket = activeConnections.get(id);
-  if (socket) {
+  const conn = activeConnections.get(id);
+  if (conn) {
+    const [socket] = conn;
     socket.send("Deregistered from server");
     socket.close();
   }
diff --git a/src/index.ts b/src/index.ts
index e8631be..1178ee3 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -38,6 +38,10 @@ declare global {
       R2_CDN_URL: string;
 
       CORS_ORIGINS: string;
+
+      // Real IP
+      REAL_IP_HEADER: string;
+      ICE_SERVERS: string;
     }
   }
 }
diff --git a/src/webrtc.ts b/src/webrtc.ts
index d2d3c3f..fd3d9aa 100644
--- a/src/webrtc.ts
+++ b/src/webrtc.ts
@@ -7,9 +7,19 @@ import { IncomingMessage } from "http";
 import { Socket } from "node:net";
 import { Device } from "@prisma/client";
 
-export const activeConnections: Map<string, WebSocket> = new Map();
+export const activeConnections: Map<string, [WebSocket, string]> = new Map();
 export const inFlight: Set<string> = new Set();
 
+function toICEServers(str: string) {
+  return str.split(",").filter(
+    (url) => url.startsWith("stun:")
+  );
+}
+
+export const iceServers = toICEServers(
+  process.env.ICE_SERVERS || "stun.cloudflare.com:3478,stun:stun.l.google.com:19302,stun:stun1.l.google.com:5349"
+);
+
 export const CreateSession = async (req: express.Request, res: express.Response) => {
   const idToken = req.session?.id_token;
   const { sub } = jose.decodeJwt(idToken);
@@ -35,12 +45,15 @@ export const CreateSession = async (req: express.Request, res: express.Response)
     );
   }
 
-  const ws = activeConnections.get(id);
-  if (!ws) {
+  const wsTuple = activeConnections.get(id);
+  if (!wsTuple) {
     console.log("No socket for id", id);
     throw new NotFoundError(`No socket for id found`, "kvm_socket_not_found");
   }
 
+  // extract the websocket and ip from the tuple
+  const [ws, ip] = wsTuple;
+
   let wsRes: ((value: unknown) => void) | null = null,
     wsRej: ((value: unknown) => void) | null = null;
 
@@ -63,7 +76,13 @@ export const CreateSession = async (req: express.Request, res: express.Response)
 
       // If the HTTP client closes the connection before the websocket response is received, reject the promise
       req.socket.on("close", wsRej);
-      ws.send(JSON.stringify({ sd, OidcGoogle: idToken }));
+
+      ws.send(JSON.stringify({
+        sd,
+        ip,
+        iceServers,
+        OidcGoogle: idToken
+      }));
     });
 
     return res.json(JSON.parse(resp.data));
@@ -170,7 +189,7 @@ export const registerWebsocketServer = (server: any) => {
       console.log(
         "Device already in active connection list. Terminating & deleting existing websocket.",
       );
-      activeConnections.get(device.id)?.terminate();
+      activeConnections.get(device.id)?.[0]?.terminate();
       activeConnections.delete(device.id);
     }
 
@@ -227,7 +246,11 @@ export const registerWebsocketServer = (server: any) => {
       return ws.close();
     }
 
-    activeConnections.set(id, ws);
+    const ip = (
+      process.env.REAL_IP_HEADER && req.headers[process.env.REAL_IP_HEADER]
+    ) || req.socket.remoteAddress;
+
+    activeConnections.set(id, [ws, `${ip}`]);
     console.log("New socket for id", id);
 
     ws.on("error", async () => {