move proxy option to server.

m1k1o · Nov 19, 2023 · 683b750 · 683b750
1 parent 3c4d7b9
commit 683b750
Show file tree

Hide file tree

Showing 5 changed files with 11 additions and 19 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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,
 	}