From bc85f11c854975119e22d92516a8280b41d3c7e8 Mon Sep 17 00:00:00 2001
From: Adrian Gallagher <thrasher@addictionsoftware.com>
Date: Tue, 12 Sep 2017 14:29:33 +1000
Subject: [PATCH] Add websocket connection limit, defaults to 1 and can be
 modified via config Add WebsocketAllowInsecureOrigin config setting, default
 to false

---
 config/config.go | 15 +++++++++++----
 websocket.go     | 20 +++++++++++++++++++-
 2 files changed, 30 insertions(+), 5 deletions(-)

diff --git a/config/config.go b/config/config.go
index 2e877c36..f59039ae 100644
--- a/config/config.go
+++ b/config/config.go
@@ -52,10 +52,12 @@ var (
 
 // WebserverConfig struct holds the prestart variables for the webserver.
 type WebserverConfig struct {
-	Enabled       bool
-	AdminUsername string
-	AdminPassword string
-	ListenAddress string
+	Enabled                      bool
+	AdminUsername                string
+	AdminPassword                string
+	ListenAddress                string
+	WebsocketConnectionLimit     int
+	WebsocketAllowInsecureOrigin bool
 }
 
 // SMSGlobalConfig structure holds all the variables you need for instant
@@ -239,6 +241,11 @@ func (c *Config) CheckWebserverConfigValues() error {
 	if port < 1 || port > 65355 {
 		return errors.New(WarningWebserverListenAddressInvalid)
 	}
+
+	if c.Webserver.WebsocketConnectionLimit <= 0 {
+		c.Webserver.WebsocketConnectionLimit = 1
+	}
+
 	return nil
 }
 
diff --git a/websocket.go b/websocket.go
index ef11548e..0d93dc56 100644
--- a/websocket.go
+++ b/websocket.go
@@ -54,11 +54,27 @@ var WebsocketClientHub []WebsocketClient
 // WebsocketClientHandler upgrades the HTTP connection to a websocket
 // compatible one
 func WebsocketClientHandler(w http.ResponseWriter, r *http.Request) {
+	connectionLimit := bot.config.Webserver.WebsocketConnectionLimit
+	numClients := len(WebsocketClientHub)
+
+	if numClients >= connectionLimit {
+		log.Printf("Websocket client rejected due to websocket client limit reached. Number of clients %d. Limit %d.",
+			numClients, connectionLimit)
+		w.WriteHeader(http.StatusForbidden)
+		return
+	}
+
 	upgrader := websocket.Upgrader{
 		WriteBufferSize: 1024,
 		ReadBufferSize:  1024,
 	}
 
+	// Allow insecure origin if the Origin request header is present and not
+	// equal to the Host request header. Default to false
+	if bot.config.Webserver.WebsocketAllowInsecureOrigin {
+		upgrader.CheckOrigin = func(r *http.Request) bool { return true }
+	}
+
 	newClient := WebsocketClient{
 		ID: len(WebsocketClientHub),
 	}
@@ -71,7 +87,9 @@ func WebsocketClientHandler(w http.ResponseWriter, r *http.Request) {
 
 	newClient.Conn = conn
 	WebsocketClientHub = append(WebsocketClientHub, newClient)
-	log.Println("New websocket client connected.")
+	numClients++
+	log.Printf("New websocket client connected. Connected clients: %d. Limit %d.",
+		numClients, connectionLimit)
 }
 
 // DisconnectWebsocketClient disconnects a websocket client