From ababbf57b5a8fbec6cac530f5b2d026fc2ac9ccf Mon Sep 17 00:00:00 2001
From: MDMCK10 <21245760+MDMCK10@users.noreply.github.com>
Date: Thu, 25 May 2023 15:59:16 +0200
Subject: [PATCH] Add (optional) Origin header check

---
 config.example.toml |  4 ++++
 src/IConfig.ts      |  2 ++
 src/WSServer.ts     | 30 ++++++++++++++++++++++++++----
 3 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/config.example.toml b/config.example.toml
index 5c6c635..40ad536 100644
--- a/config.example.toml
+++ b/config.example.toml
@@ -6,6 +6,10 @@ proxying = false
 # IPs allowed to access the server in proxy mode.
 # 99% of the time this will only be 127.0.0.1
 proxyAllowedIps = ["127.0.0.1"]
+# Whether the Origin header check is enabled.
+origin = false
+# Origins to accept connections from.
+originAllowedDomains = ["computernewb.com"]
 
 [vm]
 qemuArgs = "qemu-system-x86_64"
diff --git a/src/IConfig.ts b/src/IConfig.ts
index ab8ec6e..59aa0cb 100644
--- a/src/IConfig.ts
+++ b/src/IConfig.ts
@@ -4,6 +4,8 @@ export default interface IConfig {
         port : number;
         proxying : boolean;
         proxyAllowedIps : string[];
+        origin : boolean;
+        originAllowedDomains : string[];
     };
     vm : {
         qemuArgs : string;
diff --git a/src/WSServer.ts b/src/WSServer.ts
index 5aec822..7d2497f 100644
--- a/src/WSServer.ts
+++ b/src/WSServer.ts
@@ -85,13 +85,35 @@ export default class WSServer {
             socket.write("HTTP/1.1 400 Bad Request\n\n400 Bad Request");
             socket.destroy();
         }
-        if (
-            req.headers['sec-websocket-protocol'] !== "guacamole" 
-            // || req.headers['origin']?.toLocaleLowerCase() !== "https://computernewb.com"
-        ) {
+
+        if (req.headers['sec-websocket-protocol'] !== "guacamole") {
             killConnection();
             return;
         }
+
+        if (this.Config.http.origin) {
+            // If the client is not sending an Origin header, kill the connection.
+            if(!req.headers.origin) {
+                killConnection();
+                return;
+            }
+
+            // Try to parse the Origin header sent by the client, if it fails, kill the connection.
+            var _host;
+            try {
+                _host = new URL(req.headers.origin.toLowerCase()).hostname;
+            } catch {
+                killConnection();
+                return;
+            }
+
+            // If the domain name is not in the list of allowed origins, kill the connection.
+            if(!this.Config.http.originAllowedDomains.includes(_host)) {
+                killConnection();
+                return;
+            }
+        }
+
         if (this.Config.http.proxying) {
             // If the requesting IP isn't allowed to proxy, kill it
             //@ts-ignore