Sub2API v1.0 - AI API 网关（二开初始版本，基于上游 Wei-Shaw/sub2api）

backend/internal/server/middleware/invalid_auth_abuse_test.go

@@ -0,0 +1,147 @@
+//go:build unit
+
+package middleware
+
+import (
+	"context"
+	"errors"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/Wei-Shaw/sub2api/internal/config"
+	"github.com/Wei-Shaw/sub2api/internal/service"
+	"github.com/gin-gonic/gin"
+	"github.com/stretchr/testify/require"
+)
+
+func invalidAuthAbuseTestConfig(threshold int) *config.Config {
+	return &config.Config{
+		RunMode: config.RunModeSimple,
+		APIKeyAuth: config.APIKeyAuthCacheConfig{InvalidAbuse: config.InvalidAuthAbuseConfig{
+			Enabled: true, Threshold: threshold, WindowSeconds: 60, BlockSeconds: 60, Capacity: 256,
+		}},
+	}
+}
+
+func TestAPIKeyAuthInvalidAbuseReturns429BeforeRepository(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+	repoCalls := 0
+	repo := &stubApiKeyRepo{getByKey: func(context.Context, string) (*service.APIKey, error) {
+		repoCalls++
+		return nil, service.ErrAPIKeyNotFound
+	}}
+	cfg := invalidAuthAbuseTestConfig(3)
+	svc := service.NewAPIKeyService(repo, nil, nil, nil, nil, nil, cfg)
+	r := gin.New()
+	var reason IngressRejectReason
+	r.Use(func(c *gin.Context) { c.Next(); reason, _ = GetIngressRejectReason(c) })
+	r.Use(gin.HandlerFunc(NewAPIKeyAuthMiddleware(svc, nil, cfg)))
+	r.POST("/v1/messages", func(c *gin.Context) { c.Status(http.StatusOK) })
+
+	requests := []*http.Request{
+		httpRequest(t, "/v1/messages", "", ""),
+		httpRequest(t, "/v1/messages", "Basic malformed", ""),
+		httpRequest(t, "/v1/messages", "", "random-invalid-key"),
+	}
+	for _, req := range requests {
+		w := httptest.NewRecorder()
+		r.ServeHTTP(w, req)
+		require.NotEqual(t, http.StatusTooManyRequests, w.Code)
+	}
+
+	w := httptest.NewRecorder()
+	r.ServeHTTP(w, httpRequest(t, "/v1/messages", "", "another-random-key"))
+	require.Equal(t, http.StatusTooManyRequests, w.Code)
+	require.Equal(t, "60", w.Header().Get("Retry-After"))
+	require.Contains(t, w.Body.String(), "INVALID_AUTH_RATE_LIMITED")
+	require.Equal(t, IngressRejectInvalidAuthRateLimited, reason)
+	require.Equal(t, 1, repoCalls, "rate-limited request must not reach the repository")
+}
+
+func TestGoogleAPIKeyAuthInvalidAbuseReturnsProtocol429(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+	repoCalls := 0
+	repo := fakeAPIKeyRepo{getByKey: func(context.Context, string) (*service.APIKey, error) {
+		repoCalls++
+		return nil, service.ErrAPIKeyNotFound
+	}}
+	cfg := invalidAuthAbuseTestConfig(2)
+	svc := service.NewAPIKeyService(repo, nil, nil, nil, nil, nil, cfg)
+	r := gin.New()
+	var reason IngressRejectReason
+	r.Use(func(c *gin.Context) { c.Next(); reason, _ = GetIngressRejectReason(c) })
+	r.Use(APIKeyAuthGoogle(svc, cfg))
+	r.POST("/v1beta/models/test:generateContent", func(c *gin.Context) { c.Status(http.StatusOK) })
+	for _, key := range []string{"random-1", "random-2"} {
+		w := httptest.NewRecorder()
+		req := httpRequest(t, "/v1beta/models/test:generateContent", "", key)
+		req.Header.Del("x-api-key")
+		req.Header.Set("x-goog-api-key", key)
+		r.ServeHTTP(w, req)
+		require.Equal(t, http.StatusUnauthorized, w.Code)
+	}
+	w := httptest.NewRecorder()
+	req := httpRequest(t, "/v1beta/models/test:generateContent", "", "random-3")
+	req.Header.Del("x-api-key")
+	req.Header.Set("x-goog-api-key", "random-3")
+	r.ServeHTTP(w, req)
+	require.Equal(t, http.StatusTooManyRequests, w.Code)
+	require.Equal(t, "60", w.Header().Get("Retry-After"))
+	require.Contains(t, w.Body.String(), "RESOURCE_EXHAUSTED")
+	require.Equal(t, IngressRejectInvalidAuthRateLimited, reason)
+	require.Equal(t, 2, repoCalls)
+}
+
+func TestInvalidAuthAbuseDoesNotCountValidOrOperationalFailures(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+	user := &service.User{ID: 1, Status: service.StatusActive, Role: service.RoleUser, Balance: 1}
+	repo := &stubApiKeyRepo{getByKey: func(_ context.Context, key string) (*service.APIKey, error) {
+		switch key {
+		case "valid-key":
+			return &service.APIKey{ID: 1, UserID: 1, Key: key, Status: service.StatusActive, User: user}, nil
+		case "db-error":
+			return nil, errors.New("database unavailable")
+		default:
+			return nil, service.ErrAPIKeyNotFound
+		}
+	}}
+	cfg := invalidAuthAbuseTestConfig(10)
+	svc := service.NewAPIKeyService(repo, nil, nil, nil, nil, nil, cfg)
+	r := gin.New()
+	r.Use(gin.HandlerFunc(NewAPIKeyAuthMiddleware(svc, nil, cfg)))
+	r.POST("/t", func(c *gin.Context) { c.Status(http.StatusOK) })
+
+	for _, tc := range []struct {
+		key  string
+		want int
+	}{{"invalid", 401}, {"valid-key", 200}, {"db-error", 500}, {"db-error", 500}} {
+		w := httptest.NewRecorder()
+		r.ServeHTTP(w, httpRequest(t, "/t", "", tc.key))
+		require.Equal(t, tc.want, w.Code)
+	}
+	w := httptest.NewRecorder()
+	req := httpRequest(t, "/t", "", "")
+	req.Header.Set("x-goog-api-key", "valid-key")
+	r.ServeHTTP(w, req)
+	require.Equal(t, http.StatusOK, w.Code)
+	require.Equal(t, uint64(1), svc.InvalidAuthAbuseHealth().Recorded)
+}
+
+func TestNormalizeIngressRejectIPGroupsIPv6By64(t *testing.T) {
+	require.Equal(t, "2001:db8:abcd:1234::", normalizeIngressRejectIP("2001:db8:abcd:1234:1111::1"))
+	require.Equal(t, normalizeIngressRejectIP("2001:db8:abcd:1234:1111::1"), normalizeIngressRejectIP("2001:db8:abcd:1234:ffff::2"))
+}
+
+func httpRequest(t *testing.T, path, authorization, apiKey string) *http.Request {
+	t.Helper()
+	req := httptest.NewRequest(http.MethodPost, path, nil)
+	req.RemoteAddr = "203.0.113.10:12345"
+	if authorization != "" {
+		req.Header.Set("Authorization", authorization)
+	}
+	if apiKey != "" {
+		req.Header.Set("x-api-key", apiKey)
+	}
+	return req
+}