From 683b750189d86b91d7f7b00899063db0e7b77623 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Miroslav=20=C5=A0ediv=C3=BD?= <sedivy.miro@gmail.com>
Date: Sun, 19 Nov 2023 14:50:23 +0100
Subject: [PATCH] move proxy option to server.

---
 server/internal/config/server.go       |  7 +++++++
 server/internal/config/websocket.go    |  7 -------
 server/internal/http/http.go           |  3 +++
 server/internal/utils/ip.go            | 11 -----------
 server/internal/websocket/websocket.go |  2 +-
 5 files changed, 11 insertions(+), 19 deletions(-)

diff --git a/server/internal/config/server.go b/server/internal/config/server.go
index f39766e..a06f5a8 100644
--- a/server/internal/config/server.go
+++ b/server/internal/config/server.go
@@ -14,6 +14,7 @@ type Server struct {
 	Cert       string
 	Key        string
 	Bind       string
+	Proxy      bool
 	Static     string
 	PathPrefix string
 	CORS       []string
@@ -35,6 +36,11 @@ func (Server) 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
+	}
+
 	cmd.PersistentFlags().String("static", "./www", "path to neko client files to serve")
 	if err := viper.BindPFlag("static", cmd.PersistentFlags().Lookup("static")); err != nil {
 		return err
@@ -57,6 +63,7 @@ func (s *Server) Set() {
 	s.Cert = viper.GetString("cert")
 	s.Key = viper.GetString("key")
 	s.Bind = viper.GetString("bind")
+	s.Proxy = viper.GetBool("proxy")
 	s.Static = viper.GetString("static")
 	s.PathPrefix = path.Join("/", path.Clean(viper.GetString("path_prefix")))
 
diff --git a/server/internal/config/websocket.go b/server/internal/config/websocket.go
index 7e534fd..9bf9738 100644
--- a/server/internal/config/websocket.go
+++ b/server/internal/config/websocket.go
@@ -10,7 +10,6 @@ import (
 type WebSocket struct {
 	Password      string
 	AdminPassword string
-	Proxy         bool
 	Locks         []string
 
 	ControlProtection bool
@@ -30,11 +29,6 @@ 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
-	}
-
 	cmd.PersistentFlags().StringSlice("locks", []string{}, "resources, that will be locked when starting (control, login)")
 	if err := viper.BindPFlag("locks", cmd.PersistentFlags().Lookup("locks")); err != nil {
 		return err
@@ -63,7 +57,6 @@ func (WebSocket) Init(cmd *cobra.Command) error {
 func (s *WebSocket) Set() {
 	s.Password = viper.GetString("password")
 	s.AdminPassword = viper.GetString("password_admin")
-	s.Proxy = viper.GetBool("proxy")
 	s.Locks = viper.GetStringSlice("locks")
 
 	s.ControlProtection = viper.GetBool("control_protection")
diff --git a/server/internal/http/http.go b/server/internal/http/http.go
index d9f81a7..a013b42 100644
--- a/server/internal/http/http.go
+++ b/server/internal/http/http.go
@@ -35,6 +35,9 @@ func New(conf *config.Server, webSocketHandler types.WebSocketHandler, desktop t
 
 	router := chi.NewRouter()
 	router.Use(middleware.RequestID) // Create a request ID for each request
+	if conf.Proxy {
+		router.Use(middleware.RealIP)
+	}
 	router.Use(middleware.RequestLogger(&logformatter{logger}))
 	router.Use(middleware.Recoverer) // Recover from panics without crashing server
 	router.Use(middleware.Compress(5, "application/octet-stream"))
diff --git a/server/internal/utils/ip.go b/server/internal/utils/ip.go
index ae21211..4023935 100644
--- a/server/internal/utils/ip.go
+++ b/server/internal/utils/ip.go
@@ -38,14 +38,3 @@ func GetIP(serverUrl string) (string, error) {
 
 	return string(bytes.TrimSpace(buf)), nil
 }
-
-func GetHttpRequestIP(r *http.Request, proxy bool) string {
-	IPAddress := r.Header.Get("X-Real-Ip")
-	if IPAddress == "" {
-		IPAddress = r.Header.Get("X-Forwarded-For")
-	}
-	if IPAddress == "" || !proxy {
-		IPAddress = r.RemoteAddr
-	}
-	return IPAddress
-}
diff --git a/server/internal/websocket/websocket.go b/server/internal/websocket/websocket.go
index 6958c6f..e22ab0d 100644
--- a/server/internal/websocket/websocket.go
+++ b/server/internal/websocket/websocket.go
@@ -290,7 +290,7 @@ func (ws *WebSocketHandler) Upgrade(w http.ResponseWriter, r *http.Request) erro
 	socket := &WebSocket{
 		id:         id,
 		ws:         ws,
-		address:    utils.GetHttpRequestIP(r, ws.conf.Proxy),
+		address:    r.RemoteAddr,
 		connection: connection,
 	}