Fix feedback from code review - remove duplicate code and clean up architecture

Co-authored-by: erikdubbelboer <522870+erikdubbelboer@users.noreply.github.com>

Author: Copilot

internal/signaling/handler.go

@@ -58,12 +58,6 @@ func Handler(ctx context.Context, store stores.Store, cloudflare *cloudflare.Cre
 		logger := logging.GetLogger(ctx)
 		logger.Debug("upgrading connection")
 
-		// Extract remote address for rate limiting
-		remoteAddr := r.RemoteAddr
-		if r.Header.Get("X-Forwarded-For") != "" {
-			remoteAddr = strings.TrimSpace(strings.Split(r.Header.Get("X-Forwarded-For"), ",")[0])
-		}
-
 		ctx, cancel := context.WithCancel(ctx)
 		defer cancel()
 
@@ -85,7 +79,6 @@ func Handler(ctx context.Context, store stores.Store, cloudflare *cloudflare.Cre
 
 			retrievedIDCallback: manager.Reconnected,
 			rateLimiter:         passwordRateLimiter,
-			remoteAddr:          remoteAddr,
 		}
 		defer func() {
 			logger.Info("peer websocket closed", zap.String("peer", peer.ID), zap.String("game", peer.Game), zap.String("origin", r.Header.Get("Origin")))

internal/signaling/peer.go

@@ -27,7 +27,6 @@ type Peer struct {
 
 	retrievedIDCallback func(context.Context, string, string, string) (bool, []string, error)
 	rateLimiter         *util.PasswordRateLimiter
-	remoteAddr          string
 
 	ID     string
 	Secret string
@@ -464,9 +463,8 @@ func (p *Peer) HandleJoinPacket(ctx context.Context, packet JoinPacket) error {
 
 	// Check rate limit for password attempts if rate limiter is available
 	if p.rateLimiter != nil {
-		logger.Debug("checking rate limit", zap.String("remote_addr", p.remoteAddr))
-		if !p.rateLimiter.IsAllowed(ctx, p.remoteAddr) {
-			logger.Warn("rate limit exceeded for password attempt", zap.String("remote_addr", p.remoteAddr))
+		if !p.rateLimiter.IsAllowed(ctx, "") {
+			logger.Warn("rate limit exceeded for password attempt")
 			util.ReplyError(ctx, p.conn, util.ErrorWithCode(fmt.Errorf("too many password attempts"), "rate-limited"))
 			return nil
 		}
@@ -480,8 +478,7 @@ func (p *Peer) HandleJoinPacket(ctx context.Context, packet JoinPacket) error {
 		} else if err == stores.ErrInvalidPassword {
 			// Record failed password attempt for rate limiting
 			if p.rateLimiter != nil {
-				logger.Debug("recording failed password attempt", zap.String("remote_addr", p.remoteAddr))
-				p.rateLimiter.RecordFailedAttempt(ctx, p.remoteAddr)
+				p.rateLimiter.RecordFailedAttempt(ctx, "")
 			}
 			util.ReplyError(ctx, p.conn, util.ErrorWithCode(err, "invalid-password"))
 			return nil

internal/util/config_test.go

@@ -6,13 +6,14 @@ import (
 	"time"
 )
 
-// GetRemoteAddr extracts the remote IP address from context 
-// This duplicates the logic from metrics package to avoid circular imports
-func GetRemoteAddr(ctx context.Context) string {
-	// Use the same context key pattern as metrics package
-	type metricsContextKey int
-	remoteAddrKey := metricsContextKey(1)
-	
+// metricsContextKey matches the key type used in metrics package to avoid import cycle
+type metricsContextKey int
+
+const remoteAddrKey = metricsContextKey(1)
+
+// getRemoteAddr extracts the remote IP address from context
+// This matches the implementation in metrics package to avoid circular imports
+func getRemoteAddr(ctx context.Context) string {
 	if addr, ok := ctx.Value(remoteAddrKey).(string); ok {
 		return addr
 	}
@@ -48,8 +49,12 @@ func NewPasswordRateLimiter(maxAttempts int, windowSize time.Duration) *Password
 }
 
 // IsAllowed checks if a password attempt from the given IP is allowed
+// If remoteAddr is empty, it will be extracted from the context
 // Returns true if attempt is allowed, false if rate limited
 func (rl *PasswordRateLimiter) IsAllowed(ctx context.Context, remoteAddr string) bool {
+	if remoteAddr == "" {
+		remoteAddr = getRemoteAddr(ctx)
+	}
 	if remoteAddr == "" {
 		return true // Allow if we can't determine IP
 	}
@@ -77,7 +82,11 @@ func (rl *PasswordRateLimiter) IsAllowed(ctx context.Context, remoteAddr string)
 }
 
 // RecordFailedAttempt records a failed password attempt for the given IP
+// If remoteAddr is empty, it will be extracted from the context
 func (rl *PasswordRateLimiter) RecordFailedAttempt(ctx context.Context, remoteAddr string) {
+	if remoteAddr == "" {
+		remoteAddr = getRemoteAddr(ctx)
+	}
 	if remoteAddr == "" {
 		return // Nothing to record if we can't determine IP
 	}

internal/util/ratelimit.go

@@ -142,4 +142,60 @@ func TestPasswordRateLimiter_EmptyIP(t *testing.T) {
 	if !rl.IsAllowed(ctx, "") {
 		t.Error("Expected empty IP to still be allowed")
 	}
+}
+
+func TestPasswordRateLimiter_Configuration(t *testing.T) {
+	// Test that different configurations work
+	tests := []struct {
+		name        string
+		maxAttempts int
+		windowSize  time.Duration
+		attempts    int
+		shouldBlock bool
+	}{
+		{
+			name:        "strict limit - 2 attempts per minute",
+			maxAttempts: 2,
+			windowSize:  time.Minute,
+			attempts:    2,
+			shouldBlock: true,
+		},
+		{
+			name:        "lenient limit - 10 attempts per minute",
+			maxAttempts: 10,
+			windowSize:  time.Minute,
+			attempts:    5,
+			shouldBlock: false,
+		},
+		{
+			name:        "very strict - 1 attempt per minute",
+			maxAttempts: 1,
+			windowSize:  time.Minute,
+			attempts:    1,
+			shouldBlock: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			rl := NewPasswordRateLimiter(tt.maxAttempts, tt.windowSize)
+			defer rl.Close()
+
+			ctx := context.Background()
+			ip := "192.168.1.100"
+
+			// Record the specified number of attempts
+			for i := 0; i < tt.attempts; i++ {
+				rl.RecordFailedAttempt(ctx, ip)
+			}
+
+			// Check if next attempt should be blocked
+			allowed := rl.IsAllowed(ctx, ip)
+			blocked := !allowed
+
+			if blocked != tt.shouldBlock {
+				t.Errorf("Expected blocked=%v but got blocked=%v", tt.shouldBlock, blocked)
+			}
+		})
+	}
 }