From b381101a868595bb3adedf343383e0634c10b622 Mon Sep 17 00:00:00 2001
From: M Rasyid Hakim <rasyid.xyz@gmail.com>
Date: Thu, 23 Sep 2021 14:57:18 +0700
Subject: [PATCH] fix: checkorigin toggle

---
 docs/reference/configurations.md |  2 +-
 websocket/server.go              | 14 +++++++++-----
 2 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/docs/reference/configurations.md b/docs/reference/configurations.md
index 46767a4e..39f39eb8 100644
--- a/docs/reference/configurations.md
+++ b/docs/reference/configurations.md
@@ -78,7 +78,7 @@ Number of goroutine spawn to Ping clients.
 
 ### `SERVER_WEBSOCKET_CHECK_ORIGIN`
 
-Toggle CORS check function. Keep this to true, currently false will lead to rejection on every request.
+Toggle CORS check function. Set `true` to check each request origin. Set `false` to disable check origin and allow every request. Check origin function check against `Origin` header.
 
 * Type: `Optional`
 * Default value: `true`
diff --git a/websocket/server.go b/websocket/server.go
index 5b65db3c..38acbbaf 100644
--- a/websocket/server.go
+++ b/websocket/server.go
@@ -71,7 +71,7 @@ func CreateServer() (*Server, chan EventsBatch) {
 	pingChannel := make(chan connection, config.ServerWs.ServerMaxConn)
 	user := NewUserStore(config.ServerWs.ServerMaxConn)
 	wsHandler := &Handler{
-		websocketUpgrader: getWebSocketUpgrader(config.ServerWs.ReadBufferSize, config.ServerWs.WriteBufferSize, config.ServerWs.CheckOrigin),
+		websocketUpgrader: newWebSocketUpgrader(config.ServerWs.ReadBufferSize, config.ServerWs.WriteBufferSize, config.ServerWs.CheckOrigin),
 		bufferChannel:     bufferChannel,
 		user:              user,
 		PongWaitInterval:  config.ServerWs.PongWaitInterval,
@@ -101,13 +101,17 @@ func Router(h *Handler) http.Handler {
 	return router
 }
 
-func getWebSocketUpgrader(readBufferSize int, writeBufferSize int, checkOrigin bool) websocket.Upgrader {
+func newWebSocketUpgrader(readBufferSize int, writeBufferSize int, checkOrigin bool) websocket.Upgrader {
+	var checkOriginFunc func(r *http.Request) bool
+	if checkOrigin == false {
+		checkOriginFunc = func(r *http.Request) bool {
+			return true
+		}
+	}
 	ug := websocket.Upgrader{
 		ReadBufferSize:  readBufferSize,
 		WriteBufferSize: writeBufferSize,
-		CheckOrigin: func(r *http.Request) bool {
-			return checkOrigin
-		},
+		CheckOrigin: checkOriginFunc,
 	}
 	return ug
 }