feat(security): implement IP ban for repeated management key and Redis AUTH failures

- Added IP ban logic to `AuthenticateManagementKey` and Redis protocol handlers, blocking requests after multiple failed attempts. - Introduced unit tests to validate IP ban behavior across localhost and remote clients. - Synchronized Redis protocol's authentication policy with management key validation.
2026-04-25 21:39:58 +08:00
parent 28d78273e4
commit 2c626efc59
4 changed files with 328 additions and 76 deletions
@@ -207,84 +207,42 @@ func (h *Handler) AuthenticateManagementKey(clientIP string, localClient bool, p
 	}
 	envSecret := h.envSecret

-	fail := func() {}
-	if !localClient {
+	now := time.Now()
+	h.attemptsMu.Lock()
+	ai := h.failedAttempts[clientIP]
+	if ai != nil && !ai.blockedUntil.IsZero() {
+		if now.Before(ai.blockedUntil) {
+			remaining := ai.blockedUntil.Sub(now).Round(time.Second)
+			h.attemptsMu.Unlock()
+			return false, http.StatusForbidden, fmt.Sprintf("IP banned due to too many failed attempts. Try again in %s", remaining)
+		}
+		// Ban expired, reset state
+		ai.blockedUntil = time.Time{}
+		ai.count = 0
+	}
+	h.attemptsMu.Unlock()
+
+	if !localClient && !allowRemote {
+		return false, http.StatusForbidden, "remote management disabled"
+	}
+
+	fail := func() {
 		h.attemptsMu.Lock()
-		ai := h.failedAttempts[clientIP]
-		if ai != nil {
-			if !ai.blockedUntil.IsZero() {
-				if time.Now().Before(ai.blockedUntil) {
-					remaining := time.Until(ai.blockedUntil).Round(time.Second)
-					h.attemptsMu.Unlock()
-					return false, http.StatusForbidden, fmt.Sprintf("IP banned due to too many failed attempts. Try again in %s", remaining)
-				}
-				// Ban expired, reset state
-				ai.blockedUntil = time.Time{}
-				ai.count = 0
-			}
+		aip := h.failedAttempts[clientIP]
+		if aip == nil {
+			aip = &attemptInfo{}
+			h.failedAttempts[clientIP] = aip
+		}
+		aip.count++
+		aip.lastActivity = time.Now()
+		if aip.count >= maxFailures {
+			aip.blockedUntil = time.Now().Add(banDuration)
+			aip.count = 0
 		}
 		h.attemptsMu.Unlock()
-
-		if !allowRemote {
-			return false, http.StatusForbidden, "remote management disabled"
-		}
-
-		fail = func() {
-			h.attemptsMu.Lock()
-			aip := h.failedAttempts[clientIP]
-			if aip == nil {
-				aip = &attemptInfo{}
-				h.failedAttempts[clientIP] = aip
-			}
-			aip.count++
-			aip.lastActivity = time.Now()
-			if aip.count >= maxFailures {
-				aip.blockedUntil = time.Now().Add(banDuration)
-				aip.count = 0
-			}
-			h.attemptsMu.Unlock()
-		}
 	}

-	if secretHash == "" && envSecret == "" {
-		return false, http.StatusForbidden, "remote management key not set"
-	}
-
-	if provided == "" {
-		if !localClient {
-			fail()
-		}
-		return false, http.StatusUnauthorized, "missing management key"
-	}
-
-	if localClient {
-		if lp := h.localPassword; lp != "" {
-			if subtle.ConstantTimeCompare([]byte(provided), []byte(lp)) == 1 {
-				return true, 0, ""
-			}
-		}
-	}
-
-	if envSecret != "" && subtle.ConstantTimeCompare([]byte(provided), []byte(envSecret)) == 1 {
-		if !localClient {
-			h.attemptsMu.Lock()
-			if ai := h.failedAttempts[clientIP]; ai != nil {
-				ai.count = 0
-				ai.blockedUntil = time.Time{}
-			}
-			h.attemptsMu.Unlock()
-		}
-		return true, 0, ""
-	}
-
-	if secretHash == "" || bcrypt.CompareHashAndPassword([]byte(secretHash), []byte(provided)) != nil {
-		if !localClient {
-			fail()
-		}
-		return false, http.StatusUnauthorized, "invalid management key"
-	}
-
-	if !localClient {
+	reset := func() {
 		h.attemptsMu.Lock()
 		if ai := h.failedAttempts[clientIP]; ai != nil {
 			ai.count = 0
@@ -293,6 +251,36 @@ func (h *Handler) AuthenticateManagementKey(clientIP string, localClient bool, p
 		h.attemptsMu.Unlock()
 	}

+	if secretHash == "" && envSecret == "" {
+		return false, http.StatusForbidden, "remote management key not set"
+	}
+
+	if provided == "" {
+		fail()
+		return false, http.StatusUnauthorized, "missing management key"
+	}
+
+	if localClient {
+		if lp := h.localPassword; lp != "" {
+			if subtle.ConstantTimeCompare([]byte(provided), []byte(lp)) == 1 {
+				reset()
+				return true, 0, ""
+			}
+		}
+	}
+
+	if envSecret != "" && subtle.ConstantTimeCompare([]byte(provided), []byte(envSecret)) == 1 {
+		reset()
+		return true, 0, ""
+	}
+
+	if secretHash == "" || bcrypt.CompareHashAndPassword([]byte(secretHash), []byte(provided)) != nil {
+		fail()
+		return false, http.StatusUnauthorized, "invalid management key"
+	}
+
+	reset()
+
 	return true, 0, ""
 }

@@ -0,0 +1,38 @@
+package management
+
+import (
+	"net/http"
+	"strings"
+	"testing"
+
+	"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
+)
+
+func TestAuthenticateManagementKey_LocalhostIPBan_BlocksCorrectKeyDuringBan(t *testing.T) {
+	h := &Handler{
+		cfg:            &config.Config{},
+		failedAttempts: make(map[string]*attemptInfo),
+		envSecret:      "test-secret",
+	}
+
+	for i := 0; i < 5; i++ {
+		allowed, statusCode, errMsg := h.AuthenticateManagementKey("127.0.0.1", true, "wrong-secret")
+		if allowed {
+			t.Fatalf("expected auth to be denied at attempt %d", i+1)
+		}
+		if statusCode != http.StatusUnauthorized || errMsg != "invalid management key" {
+			t.Fatalf("unexpected auth failure at attempt %d: status=%d msg=%q", i+1, statusCode, errMsg)
+		}
+	}
+
+	allowed, statusCode, errMsg := h.AuthenticateManagementKey("127.0.0.1", true, "test-secret")
+	if allowed {
+		t.Fatalf("expected correct key to be denied while banned")
+	}
+	if statusCode != http.StatusForbidden {
+		t.Fatalf("expected forbidden status while banned, got %d", statusCode)
+	}
+	if !strings.HasPrefix(errMsg, "IP banned due to too many failed attempts. Try again in") {
+		t.Fatalf("unexpected banned message: %q", errMsg)
+	}
+}