From 44a64bfce4bc7435ad476071f8e0fe96e0e43d69 Mon Sep 17 00:00:00 2001
From: Sebastian Mark <smark@posteo.net>
Date: Sun, 3 Nov 2024 10:00:14 +0100
Subject: [PATCH] feat: handle X-Forward-For for log output

- add Gorilla handlers package for enhanced HTTP handling
- refactor HTTP server to use a new ServeMux for routing
- update ListenAndServe to utilize ProxyHeaders for better proxy support
- add RealIP field to client model
  - use RealIP fild for connect/disconnect log output
---
 cmd/server/main.go                      | 6 ++++--
 go.mod                                  | 2 ++
 go.sum                                  | 4 ++++
 internal/websocket/client_commands.go   | 2 +-
 internal/websocket/handle_connection.go | 7 ++++---
 pkg/models/client.go                    | 3 ++-
 6 files changed, 17 insertions(+), 7 deletions(-)

diff --git a/cmd/server/main.go b/cmd/server/main.go
index 126ec2f..c30f3e4 100644
--- a/cmd/server/main.go
+++ b/cmd/server/main.go
@@ -4,6 +4,7 @@ import (
 	"flag"
 	"fmt"
 	"github.com/charmbracelet/log"
+	"github.com/gorilla/handlers"
 	"net/http"
 	"os"
 
@@ -43,14 +44,15 @@ func Start() {
 	listen := fmt.Sprintf("%s:%d", serverConfig.ListenAddress, serverConfig.ListenPort)
 
 	// start connection handler and broadcast
-	http.HandleFunc("/ws", websocket.HandleConnection)
+	r := http.NewServeMux()
+	r.HandleFunc("/ws", websocket.HandleConnection)
 	go websocket.SendPermanentBroadCastMessage()
 
 	log.Info("GoTomato started", "version", metadata.GoTomatoVersion)
 	log.Info("Websocket listening", "address", listen)
 
 	// start the listener
-	err := http.ListenAndServe(listen, nil)
+	err := http.ListenAndServe(listen, handlers.ProxyHeaders(r))
 	if err != nil {
 		log.Fatal("Error starting server:", "msg", err)
 	}
diff --git a/go.mod b/go.mod
index 148683e..e78b89d 100644
--- a/go.mod
+++ b/go.mod
@@ -4,12 +4,14 @@ go 1.23
 
 require (
 	github.com/charmbracelet/log v0.4.0
+	github.com/gorilla/handlers v1.5.2
 	github.com/gorilla/websocket v1.5.3
 )
 
 require (
 	github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect
 	github.com/charmbracelet/lipgloss v0.10.0 // indirect
+	github.com/felixge/httpsnoop v1.0.3 // indirect
 	github.com/go-logfmt/logfmt v0.6.0 // indirect
 	github.com/lucasb-eyer/go-colorful v1.2.0 // indirect
 	github.com/mattn/go-isatty v0.0.18 // indirect
diff --git a/go.sum b/go.sum
index 9437970..77f660c 100644
--- a/go.sum
+++ b/go.sum
@@ -6,8 +6,12 @@ github.com/charmbracelet/log v0.4.0 h1:G9bQAcx8rWA2T3pWvx7YtPTPwgqpk7D68BX21IRW8
 github.com/charmbracelet/log v0.4.0/go.mod h1:63bXt/djrizTec0l11H20t8FDSvA4CRZJ1KH22MdptM=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/felixge/httpsnoop v1.0.3 h1:s/nj+GCswXYzN5v2DpNMuMQYe+0DDwt5WVCU6CWBdXk=
+github.com/felixge/httpsnoop v1.0.3/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U=
 github.com/go-logfmt/logfmt v0.6.0 h1:wGYYu3uicYdqXVgoYbvnkrPVXkuLM1p1ifugDMEdRi4=
 github.com/go-logfmt/logfmt v0.6.0/go.mod h1:WYhtIu8zTZfxdn5+rREduYbwxfcBr/Vr6KEVveWlfTs=
+github.com/gorilla/handlers v1.5.2 h1:cLTUSsNkgcwhgRqvCNmdbRWG0A3N4F+M2nWKdScwyEE=
+github.com/gorilla/handlers v1.5.2/go.mod h1:dX+xVpaxdSw+q0Qek8SSsl3dfMk3jNddUkMzo0GtH0w=
 github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg=
 github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
 github.com/lucasb-eyer/go-colorful v1.2.0 h1:1nnpGOrhyZZuNyfu1QjKiUICQ74+3FNCN69Aj6K7nkY=
diff --git a/internal/websocket/client_commands.go b/internal/websocket/client_commands.go
index 2fba584..1a66198 100644
--- a/internal/websocket/client_commands.go
+++ b/internal/websocket/client_commands.go
@@ -21,8 +21,8 @@ func handleClientCommands(c models.WebsocketClient) {
 
 		_, message, err := ws.ReadMessage()
 		if err != nil {
-			log.Info("Client disconnected:", "msg", err, "host", ws.RemoteAddr(), "clients", len(Clients)-1)
 			delete(Clients, ws.LocalAddr())
+			log.Info("Client disconnected:", "msg", err, "host", c.RealIP, "clients", len(Clients))
 			break
 		}
 
diff --git a/internal/websocket/handle_connection.go b/internal/websocket/handle_connection.go
index 6abb339..6ee22bc 100644
--- a/internal/websocket/handle_connection.go
+++ b/internal/websocket/handle_connection.go
@@ -29,16 +29,17 @@ func HandleConnection(w http.ResponseWriter, r *http.Request) {
 	}
 	defer ws.Close()
 
-	log.Info("Client connected", "host", ws.RemoteAddr(), "clients", len(Clients)+1)
-
 	// Register the new client
 	client := models.WebsocketClient{
-		Conn: ws,
+		Conn:   ws,
+		RealIP: r.RemoteAddr,
 	}
 	mu.Lock()
 	Clients[ws.LocalAddr()] = &client
 	mu.Unlock()
 
+	log.Info("Client connected", "host", client.RealIP, "clients", len(Clients))
+
 	// Listen for commands from the connected client
 	handleClientCommands(client)
 }
diff --git a/pkg/models/client.go b/pkg/models/client.go
index c0998f8..aec2f3c 100644
--- a/pkg/models/client.go
+++ b/pkg/models/client.go
@@ -14,7 +14,8 @@ type ClientCommand struct {
 
 // Represents a single client
 type WebsocketClient struct {
-	Conn *websocket.Conn
+	Conn   *websocket.Conn
+	RealIP string
 }
 
 // Sends a message to the websocket.