From 584513de9b2b883a2a759ec88f432a0325000664 Mon Sep 17 00:00:00 2001
From: Craig <craigbett@gmail.com>
Date: Mon, 6 Apr 2020 20:14:30 +0000
Subject: [PATCH] reverse proxy mode

---
 server/internal/types/config/websocket.go | 7 +++++++
 server/internal/websocket/websocket.go    | 8 ++++++--
 2 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/server/internal/types/config/websocket.go b/server/internal/types/config/websocket.go
index b93dc59..099f34b 100644
--- a/server/internal/types/config/websocket.go
+++ b/server/internal/types/config/websocket.go
@@ -8,6 +8,7 @@ import (
 type WebSocket struct {
 	Password      string
 	AdminPassword string
+	Proxy         bool
 }
 
 func (WebSocket) Init(cmd *cobra.Command) error {
@@ -21,10 +22,16 @@ func (WebSocket) Init(cmd *cobra.Command) error {
 		return err
 	}
 
+	cmd.PersistentFlags().Bool("proxy", false, "enable reverse proxy mode")
+	if err := viper.BindPFlag("proxy", cmd.PersistentFlags().Lookup("proxy")); err != nil {
+		return err
+	}
+
 	return nil
 }
 
 func (s *WebSocket) Set() {
 	s.Password = viper.GetString("password")
 	s.AdminPassword = viper.GetString("password_admin")
+	s.Proxy = viper.GetBool("proxy")
 }
diff --git a/server/internal/websocket/websocket.go b/server/internal/websocket/websocket.go
index 1fe4598..fc6b4e3 100644
--- a/server/internal/websocket/websocket.go
+++ b/server/internal/websocket/websocket.go
@@ -70,7 +70,7 @@ func (ws *WebSocketHandler) Start() error {
 		}
 	})
 
-	ws.sessions.OnDestroy(func(id string) {
+	ws.sessions.OnDestroy(func(id string, session types.Session) {
 		if err := ws.handler.SessionDestroyed(id); err != nil {
 			ws.logger.Warn().Str("id", id).Err(err).Msg("session destroyed with and error")
 		} else {
@@ -191,7 +191,11 @@ func (ws *WebSocketHandler) Upgrade(w http.ResponseWriter, r *http.Request) erro
 }
 
 func (ws *WebSocketHandler) authenticate(r *http.Request) (string, string, bool, error) {
-	ip := utils.ReadUserIP(r)
+	ip := r.RemoteAddr
+
+	if ws.conf.Proxy {
+		ip = utils.ReadUserIP(r)
+	}
 
 	id, err := utils.NewUID(32)
 	if err != nil {