Websocket and whois restructured and added sections

pkg/web/websocket.go

@@ -1,6 +1,6 @@
 // Fail2ban UI - A Swiss made, management interface for Fail2ban.
 //
-// Copyright (C) 2025 Swissmakers GmbH (https://swissmakers.ch)
+// Copyright (C) 2026 Swissmakers GmbH (https://swissmakers.ch)
 //
 // Licensed under the GNU General Public License, Version 3 (GPL-3.0)
 // You may not use this file except in compliance with the License.
@@ -28,14 +28,27 @@ import (
 	"github.com/swissmakers/fail2ban-ui/internal/storage"
 )
 
+// =========================================================================
+//  Types and Constants
+// =========================================================================
+
+type Client struct {
+	hub  *Hub
+	conn *websocket.Conn
+	send chan []byte
+}
+
+type Hub struct {
+	clients    map[*Client]bool
+	broadcast  chan []byte
+	register   chan *Client
+	unregister chan *Client
+	mu         sync.RWMutex
+}
+
 const (
-	// Time allowed to write a message to the peer
+	writeWait  = 10 * time.Second
-	writeWait = 10 * time.Second
+	pongWait   = 60 * time.Second
-
-	// Time allowed to read the next pong message from the peer
-	pongWait = 60 * time.Second
-
-	// Send pings to peer with this period (must be less than pongWait)
 	pingPeriod = (pongWait * 9) / 10
 )
 
@@ -43,37 +56,15 @@ var upgrader = websocket.Upgrader{
 	ReadBufferSize:  1024,
 	WriteBufferSize: 1024,
 	CheckOrigin: func(r *http.Request) bool {
-		// Allow all origins for now - can be restricted in production
 		return true
 	},
 }
 
-// Client represents a WebSocket connection
+// =========================================================================
-type Client struct {
+//  Fail2ban-UI WebSocket Hub
-	hub  *Hub
+// =========================================================================
-	conn *websocket.Conn
-	send chan []byte
-}
 
-// Hub maintains the set of active clients and broadcasts messages to them
+// Broadcasts the console log message to all connected clients.
-type Hub struct {
-	// Registered clients
-	clients map[*Client]bool
-
-	// Inbound messages from clients
-	broadcast chan []byte
-
-	// Register requests from clients
-	register chan *Client
-
-	// Unregister requests from clients
-	unregister chan *Client
-
-	// Mutex for thread-safe operations
-	mu sync.RWMutex
-}
-
-// BroadcastConsoleLog broadcasts a console log message to all connected clients
 func (h *Hub) BroadcastConsoleLog(message string) {
 	logMsg := map[string]interface{}{
 		"type":    "console_log",
@@ -89,11 +80,11 @@ func (h *Hub) BroadcastConsoleLog(message string) {
 	select {
 	case h.broadcast <- data:
 	default:
-		// Channel full, drop message
+		log.Printf("Broadcast channel full, dropping console log")
 	}
 }
 
-// NewHub creates a new WebSocket hub
+// Creates new Hub instance.
 func NewHub() *Hub {
 	return &Hub{
 		clients:    make(map[*Client]bool),
@@ -103,9 +94,8 @@ func NewHub() *Hub {
 	}
 }
 
-// Run starts the hub's main loop
+// Runs the Hub.
 func (h *Hub) Run() {
-	// Start heartbeat ticker
 	ticker := time.NewTicker(30 * time.Second)
 	defer ticker.Stop()
 
@@ -139,13 +129,12 @@ func (h *Hub) Run() {
 			h.mu.RUnlock()
 
 		case <-ticker.C:
-			// Send heartbeat to all clients
 			h.sendHeartbeat()
 		}
 	}
 }
 
-// sendHeartbeat sends a heartbeat message to all connected clients
+// Sends heartbeat message to all connected clients.
 func (h *Hub) sendHeartbeat() {
 	message := map[string]interface{}{
 		"type":   "heartbeat",
@@ -170,7 +159,10 @@ func (h *Hub) sendHeartbeat() {
 	h.mu.RUnlock()
 }
 
-// BroadcastBanEvent broadcasts a ban event to all connected clients
+// =========================================================================
+//  Broadcast Ban Event
+// =========================================================================
+
 func (h *Hub) BroadcastBanEvent(event storage.BanEventRecord) {
 	message := map[string]interface{}{
 		"type": "ban_event",
@@ -189,7 +181,10 @@ func (h *Hub) BroadcastBanEvent(event storage.BanEventRecord) {
 	}
 }
 
-// BroadcastUnbanEvent broadcasts an unban event to all connected clients
+// =========================================================================
+//  Broadcast Unban Event
+// =========================================================================
+
 func (h *Hub) BroadcastUnbanEvent(event storage.BanEventRecord) {
 	message := map[string]interface{}{
 		"type": "unban_event",
@@ -208,7 +203,11 @@ func (h *Hub) BroadcastUnbanEvent(event storage.BanEventRecord) {
 	}
 }
 
-// readPump pumps messages from the WebSocket connection to the hub
+// =========================================================================
+//  WebSocket Helper Functions
+// =========================================================================
+
+// Reads messages from the WebSocket connection.
 func (c *Client) readPump() {
 	defer func() {
 		c.hub.unregister <- c
@@ -232,7 +231,7 @@ func (c *Client) readPump() {
 	}
 }
 
-// writePump pumps messages from the hub to the WebSocket connection
+// Writes messages to the WebSocket connection.
 func (c *Client) writePump() {
 	ticker := time.NewTicker(pingPeriod)
 	defer func() {
@@ -255,7 +254,6 @@ func (c *Client) writePump() {
 			}
 			w.Write(message)
 
-			// Add queued messages to the current websocket message
 			n := len(c.send)
 			for i := 0; i < n; i++ {
 				w.Write([]byte{'\n'})
@@ -275,7 +273,7 @@ func (c *Client) writePump() {
 	}
 }
 
-// serveWS handles WebSocket requests from clients
+// Serves the WebSocket connection.
 func serveWS(hub *Hub, c *gin.Context) {
 	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
 	if err != nil {
@@ -291,12 +289,11 @@ func serveWS(hub *Hub, c *gin.Context) {
 
 	client.hub.register <- client
 
-	// Allow collection of memory referenced by the caller by doing all work in new goroutines
 	go client.writePump()
 	go client.readPump()
 }
 
-// WebSocketHandler is the Gin handler for WebSocket connections
+// This is called from routes.go and returns the Gin handler for WebSocket connections.
 func WebSocketHandler(hub *Hub) gin.HandlerFunc {
 	return func(c *gin.Context) {
 		serveWS(hub, c)

pkg/web/whois.go

@@ -1,6 +1,6 @@
 // Fail2ban UI - A Swiss made, management interface for Fail2ban.
 //
-// Copyright (C) 2025 Swissmakers GmbH (https://swissmakers.ch)
+// Copyright (C) 2026 Swissmakers GmbH (https://swissmakers.ch)
 //
 // Licensed under the GNU General Public License, Version 3 (GPL-3.0)
 // You may not use this file except in compliance with the License.
@@ -25,21 +25,26 @@ import (
 	"github.com/likexian/whois"
 )
 
-var (
+// =========================================================================
-	whoisCache      = make(map[string]cachedWhois)
+//  Types and Constants
-	whoisCacheMutex sync.RWMutex
+// =========================================================================
-	cacheExpiry     = 24 * time.Hour
-)
 
 type cachedWhois struct {
 	data      string
 	timestamp time.Time
 }
 
-// lookupWhois performs a whois lookup for the given IP address.
+var (
-// It uses caching to avoid repeated queries for the same IP.
+	whoisCache      = make(map[string]cachedWhois)
+	whoisCacheMutex sync.RWMutex
+	cacheExpiry     = 24 * time.Hour
+)
+
+// =========================================================================
+//  Lookup Whois Data
+// =========================================================================
+
 func lookupWhois(ip string) (string, error) {
-	// Check cache first
 	whoisCacheMutex.RLock()
 	if cached, ok := whoisCache[ip]; ok {
 		if time.Since(cached.timestamp) < cacheExpiry {
@@ -49,7 +54,6 @@ func lookupWhois(ip string) (string, error) {
 	}
 	whoisCacheMutex.RUnlock()
 
-	// Perform whois lookup with timeout
 	done := make(chan string, 1)
 	errChan := make(chan error, 1)
 
@@ -65,20 +69,17 @@ func lookupWhois(ip string) (string, error) {
 	var whoisData string
 	select {
 	case whoisData = <-done:
-		// Success - cache will be updated below
 	case err := <-errChan:
 		return "", fmt.Errorf("whois lookup failed: %w", err)
 	case <-time.After(10 * time.Second):
 		return "", fmt.Errorf("whois lookup timeout after 10 seconds")
 	}
 
-	// Cache the result
 	whoisCacheMutex.Lock()
 	whoisCache[ip] = cachedWhois{
 		data:      whoisData,
 		timestamp: time.Now(),
 	}
-	// Clean old cache entries if cache is getting large
 	if len(whoisCache) > 1000 {
 		now := time.Now()
 		for k, v := range whoisCache {
@@ -92,15 +93,16 @@ func lookupWhois(ip string) (string, error) {
 	return whoisData, nil
 }
 
-// extractCountryFromWhois attempts to extract country code from whois data.
+// =========================================================================
-// This is a fallback if GeoIP lookup fails.
+//  Extract Country from Whois Data
+// =========================================================================
+
 func extractCountryFromWhois(whoisData string) string {
 	lines := strings.Split(whoisData, "\n")
 	for _, line := range lines {
 		line = strings.TrimSpace(line)
 		lineLower := strings.ToLower(line)
 
-		// Look for country field
 		if strings.HasPrefix(lineLower, "country:") {
 			parts := strings.SplitN(line, ":", 2)
 			if len(parts) == 2 {
@@ -110,7 +112,6 @@ func extractCountryFromWhois(whoisData string) string {
 				}
 			}
 		}
-		// Alternative format
 		if strings.HasPrefix(lineLower, "country code:") {
 			parts := strings.SplitN(line, ":", 2)
 			if len(parts) == 2 {