From 698fafb6badcb77bf15a144e916794222e4d0ed4 Mon Sep 17 00:00:00 2001
From: sophie <sofija@duck.com>
Date: Fri, 19 Jul 2024 04:52:39 +0300
Subject: [PATCH] support reverse proxies

---
 README.md     |  2 ++
 src/Client.ts |  4 ++--
 src/Server.ts | 14 ++++++++++++--
 3 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index 70b7a66..96ff17a 100644
--- a/README.md
+++ b/README.md
@@ -14,6 +14,8 @@ To use this server [use a patched client](https://git.sad.ovh/sophie/pianoverse_
 ```
 HASH="somesecurestring"
 PORT=8081
+# Only if you're running pianoverse_server under a proxy like Caddy, nginx or Apache
+TRUST_PROXY=true
 ```
 
 ![chat](assets/screenshot1.png)
diff --git a/src/Client.ts b/src/Client.ts
index e302e82..399d52b 100644
--- a/src/Client.ts
+++ b/src/Client.ts
@@ -9,9 +9,9 @@ export class Client {
   uniqWsID: string;
   private ws: ServerWebSocket<Client>;
 
-  constructor(ws: ServerWebSocket<Client>) {
+  constructor(ws: ServerWebSocket<Client>, ip: string) {
     this.ws = ws;
-    this.id = [...Bun.SHA256.hash(ws.remoteAddress + process.env.HASH)]
+    this.id = [...Bun.SHA256.hash(ip + process.env.HASH)]
       .slice(0, 7)
       .map((z) => z.toString(16))
       .join("");
diff --git a/src/Server.ts b/src/Server.ts
index e94e71e..4f119d1 100644
--- a/src/Server.ts
+++ b/src/Server.ts
@@ -41,7 +41,14 @@ export class Server {
     Bun.serve({
       // #region WS upgrading
       fetch(req, server) {
-        if (server.upgrade(req, { data: req.headers.get("User-Agent") })) {
+        let data;
+        if(process.env.TRUST_PROXY) {
+          data = req.headers.get("X-Forwarded-For")?.split(",")[0]?.trim()
+          if(!data) {
+            console.log('Trust proxy is enabled, but XFF is empty. Spoofing / server issue?')
+          }
+        }
+        if (server.upgrade(req, { data })) {
           return;
         }
         return new Response("Upgrade failed", { status: 500 });
@@ -297,7 +304,10 @@ export class Server {
           }
         },
         open(ws: ServerWebSocket<Client>) {
-          const client = new Client(ws);
+          let ip: string = ws.remoteAddress;
+          if(process.env.TRUST_PROXY) ip = ws.data as unknown as string;
+
+          const client = new Client(ws, ip);
           if (getSiteBan(client.id)) {
             ws.close();
             return;