89 lines
2.8 KiB
Go
89 lines
2.8 KiB
Go
package management
|
|
|
|
import (
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
|
"github.com/router-for-me/CLIProxyAPI/v7/internal/pluginhost"
|
|
)
|
|
|
|
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)
|
|
}
|
|
}
|
|
|
|
func TestMiddlewareSetsSupportPluginHeader(t *testing.T) {
|
|
|
|
h := &Handler{
|
|
cfg: &config.Config{},
|
|
failedAttempts: make(map[string]*attemptInfo),
|
|
envSecret: "test-secret",
|
|
}
|
|
middleware := h.Middleware()
|
|
|
|
t.Run("invalid key", func(t *testing.T) {
|
|
rec := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(rec)
|
|
c.Request = httptest.NewRequest(http.MethodGet, "/v0/management/config", nil)
|
|
c.Request.RemoteAddr = "127.0.0.1:12345"
|
|
c.Request.Header.Set("X-Management-Key", "wrong-secret")
|
|
|
|
middleware(c)
|
|
|
|
if rec.Code != http.StatusUnauthorized {
|
|
t.Fatalf("status = %d, want %d", rec.Code, http.StatusUnauthorized)
|
|
}
|
|
if got := rec.Header().Get("X-CPA-SUPPORT-PLUGIN"); got != pluginhost.SupportPluginHeaderValue() {
|
|
t.Fatalf("X-CPA-SUPPORT-PLUGIN = %q, want %q", got, pluginhost.SupportPluginHeaderValue())
|
|
}
|
|
})
|
|
|
|
t.Run("valid key", func(t *testing.T) {
|
|
engine := gin.New()
|
|
engine.GET("/v0/management/config", middleware, func(c *gin.Context) {
|
|
c.Status(http.StatusOK)
|
|
})
|
|
|
|
rec := httptest.NewRecorder()
|
|
req := httptest.NewRequest(http.MethodGet, "/v0/management/config", nil)
|
|
req.RemoteAddr = "127.0.0.1:12345"
|
|
req.Header.Set("X-Management-Key", "test-secret")
|
|
engine.ServeHTTP(rec, req)
|
|
|
|
if rec.Code != http.StatusOK {
|
|
t.Fatalf("status = %d, want %d", rec.Code, http.StatusOK)
|
|
}
|
|
if got := rec.Header().Get("X-CPA-SUPPORT-PLUGIN"); got != pluginhost.SupportPluginHeaderValue() {
|
|
t.Fatalf("X-CPA-SUPPORT-PLUGIN = %q, want %q", got, pluginhost.SupportPluginHeaderValue())
|
|
}
|
|
})
|
|
}
|