From 7f511abb4f4f458aa2f82ce70a33e33d26c64e19 Mon Sep 17 00:00:00 2001
From: Bui <bui@bui.pm>
Date: Thu, 24 Mar 2022 02:32:31 +0000
Subject: [PATCH] allow 2p url config, origin limiting for server

---
 app.py              |  3 ++-
 config.example.py   |  3 +++
 public/src/js/p2.js |  2 +-
 server.py           | 10 ++++++++--
 4 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/app.py b/app.py
index d26d7fa..4c6c3e9 100644
--- a/app.py
+++ b/app.py
@@ -135,7 +135,8 @@ def get_config(credentials=False):
         'accounts': take_config('ACCOUNTS'),
         'custom_js': take_config('CUSTOM_JS'),
         'plugins': take_config('PLUGINS') and [x for x in take_config('PLUGINS') if x['url']],
-        'preview_type': take_config('PREVIEW_TYPE') or 'mp3'
+        'preview_type': take_config('PREVIEW_TYPE') or 'mp3',
+        'multiplayer_url': take_config('MULTIPLAYER_URL')
     }
     if credentials:
         google_credentials = take_config('GOOGLE_CREDENTIALS')
diff --git a/config.example.py b/config.example.py
index c515839..9a03ae5 100644
--- a/config.example.py
+++ b/config.example.py
@@ -4,6 +4,9 @@ ASSETS_BASEURL = '/assets/'
 # The full URL base song URL, with trailing slash.
 SONGS_BASEURL = '/songs/'
 
+# Multiplayer websocket URL. Defaults to /p2 if blank.
+MULTIPLAYER_URL = ''
+
 # The email address to display in the "About Simulator" menu.
 EMAIL = None
 
diff --git a/public/src/js/p2.js b/public/src/js/p2.js
index 3025334..99ced66 100644
--- a/public/src/js/p2.js
+++ b/public/src/js/p2.js
@@ -32,7 +32,7 @@ class P2Connection{
 		if(this.closed && !this.disabled){
 			this.closed = false
 			var wsProtocol = location.protocol == "https:" ? "wss:" : "ws:"
-			this.socket = new WebSocket(wsProtocol + "//" + location.host + "/p2")
+			this.socket = new WebSocket(gameConfig.multiplayer_url ? gameConfig.multiplayer_url : wsProtocol + "//" + location.host + "/p2")
 			pageEvents.race(this.socket, "open", "close").then(response => {
 				if(response.type === "open"){
 					return this.openEvent()
diff --git a/server.py b/server.py
index 665b391..5f4016e 100644
--- a/server.py
+++ b/server.py
@@ -1,11 +1,17 @@
 #!/usr/bin/env python3
 
+import argparse
 import asyncio
 import websockets
 import json
 import random
 import sys
 
+parser = argparse.ArgumentParser(description='Run the taiko-web multiplayer server.')
+parser.add_argument('port', type=int, metavar='PORT', nargs='?', default=34802, help='Port to listen on.')
+parser.add_argument('-o', '--allow-origin', action='append', help='Limit incoming connections to the specified origin. Can be specified multiple times.')
+args = parser.parse_args()
+
 server_status = {
 	"waiting": {},
 	"users": [],
@@ -372,11 +378,11 @@ async def connection(ws, path):
 		elif user["action"] == "invite" and user["session"] in server_status["invites"]:
 			del server_status["invites"][user["session"]]
 
-port = int(sys.argv[1]) if len(sys.argv) > 1 else 34802
+port = args.port
 print('Starting server on port %d' % port)
 loop = asyncio.get_event_loop()
 tasks = asyncio.gather(
-	websockets.serve(connection, "localhost", port)
+	websockets.serve(connection, "localhost", port, origins=args.allow_origin)
 )
 try:
 	loop.run_until_complete(tasks)