package app import ( "encoding/json" "net/http" "net/http/httptest" "strings" "testing" "time" "ccLoad/internal/model" "github.com/gin-gonic/gin" ) // ============================================================================ // 认证中间件测试 // 覆盖 RequireAPIAuth 和 RequireTokenAuth 的各种认证场景 // ============================================================================ // ============================================================================ // RequireAPIAuth 测试 // ============================================================================ func TestRequireAPIAuth_BearerToken(t *testing.T) { t.Parallel() svc := newTestAuthService(t) injectAPIToken(svc, "sk-test-123", 0, 1) // expiresAt=0 永不过期 req := httptest.NewRequest(http.MethodGet, "/test", nil) req.Header.Set("Authorization", "Bearer sk-test-123") w := runMiddleware(t, svc.RequireAPIAuth(), req) if w.Code != http.StatusOK { t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String()) } } func TestRequireAPIAuth_XAPIKey(t *testing.T) { t.Parallel() svc := newTestAuthService(t) injectAPIToken(svc, "key-abc", 0, 2) req := httptest.NewRequest(http.MethodGet, "/test", nil) req.Header.Set("X-API-Key", "key-abc") w := runMiddleware(t, svc.RequireAPIAuth(), req) if w.Code != http.StatusOK { t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String()) } } func TestRequireAPIAuth_GoogleKey(t *testing.T) { t.Parallel() svc := newTestAuthService(t) injectAPIToken(svc, "AIza-google-key", 0, 3) req := httptest.NewRequest(http.MethodGet, "/test", nil) req.Header.Set("x-goog-api-key", "AIza-google-key") w := runMiddleware(t, svc.RequireAPIAuth(), req) if w.Code != http.StatusOK { t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String()) } } func TestRequireAPIAuth_QueryParam(t *testing.T) { t.Parallel() svc := newTestAuthService(t) injectAPIToken(svc, "query-key-789", 0, 4) req := httptest.NewRequest(http.MethodGet, "/test?key=query-key-789", nil) w := runMiddleware(t, svc.RequireAPIAuth(), req) if w.Code != http.StatusOK { t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String()) } } func TestRequireAPIAuth_InvalidToken(t *testing.T) { t.Parallel() svc := newTestAuthService(t) injectAPIToken(svc, "real-token", 0, 1) req := httptest.NewRequest(http.MethodGet, "/test", nil) req.Header.Set("Authorization", "Bearer wrong-token") w := runMiddleware(t, svc.RequireAPIAuth(), req) if w.Code != http.StatusUnauthorized { t.Fatalf("expected 401, got %d: %s", w.Code, w.Body.String()) } } func TestRequireAPIAuth_NoToken(t *testing.T) { t.Parallel() svc := newTestAuthService(t) injectAPIToken(svc, "some-token", 0, 1) req := httptest.NewRequest(http.MethodGet, "/test", nil) w := runMiddleware(t, svc.RequireAPIAuth(), req) if w.Code != http.StatusUnauthorized { t.Fatalf("expected 401, got %d: %s", w.Code, w.Body.String()) } } func TestRequireAPIAuth_NoConfiguredTokens(t *testing.T) { t.Parallel() svc := newTestAuthService(t) // 不注入任何 token req := httptest.NewRequest(http.MethodGet, "/test", nil) req.Header.Set("Authorization", "Bearer any-token") w := runMiddleware(t, svc.RequireAPIAuth(), req) if w.Code != http.StatusUnauthorized { t.Fatalf("expected 401, got %d: %s", w.Code, w.Body.String()) } } func TestRequireAPIAuth_ExpiredToken(t *testing.T) { t.Parallel() svc := newTestAuthService(t) // 设置过期时间为过去(毫秒时间戳) expiredAt := time.Now().Add(-time.Hour).UnixMilli() injectAPIToken(svc, "expired-token", expiredAt, 5) req := httptest.NewRequest(http.MethodGet, "/test", nil) req.Header.Set("Authorization", "Bearer expired-token") w := runMiddleware(t, svc.RequireAPIAuth(), req) if w.Code != http.StatusUnauthorized { t.Fatalf("expected 401, got %d: %s", w.Code, w.Body.String()) } // 验证响应包含 "token expired" var resp map[string]string if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil { t.Fatalf("unmarshal response: %v", err) } if resp["error"] != "token expired" { t.Fatalf("expected 'token expired' error, got: %s", resp["error"]) } // 验证懒惰删除:token 应已从内存中移除 tokenHash := model.HashToken("expired-token") svc.authTokensMux.RLock() _, stillExists := svc.authTokens[tokenHash] svc.authTokensMux.RUnlock() if stillExists { t.Fatal("expected expired token to be lazily deleted from memory") } } func TestRequireAPIAuth_ContextValues(t *testing.T) { t.Parallel() svc := newTestAuthService(t) injectAPIToken(svc, "ctx-token", 0, 42) req := httptest.NewRequest(http.MethodGet, "/test", nil) req.Header.Set("Authorization", "Bearer ctx-token") w := runMiddleware(t, svc.RequireAPIAuth(), req) if w.Code != http.StatusOK { t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String()) } var resp map[string]any if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil { t.Fatalf("unmarshal response: %v", err) } // 验证 token_hash 被设置到 context expectedHash := model.HashToken("ctx-token") if got, ok := resp["token_hash"].(string); !ok || got != expectedHash { t.Fatalf("expected token_hash=%s, got=%v", expectedHash, resp["token_hash"]) } // 验证 token_id 被设置到 context if got, ok := resp["token_id"].(float64); !ok || int64(got) != 42 { t.Fatalf("expected token_id=42, got=%v", resp["token_id"]) } } func TestRequireAPIAuth_LastUsedUpdate(t *testing.T) { t.Parallel() svc := newTestAuthService(t) injectAPIToken(svc, "lu-token", 0, 1) req := httptest.NewRequest(http.MethodGet, "/test", nil) req.Header.Set("Authorization", "Bearer lu-token") _ = runMiddleware(t, svc.RequireAPIAuth(), req) // 验证 tokenHash 被发送到 lastUsedCh(非阻塞通道) expectedHash := model.HashToken("lu-token") select { case hash := <-svc.lastUsedCh: if hash != expectedHash { t.Fatalf("expected lastUsedCh to receive %s, got %s", expectedHash, hash) } case <-time.After(100 * time.Millisecond): t.Fatal("expected tokenHash to be sent to lastUsedCh") } } func TestRequireAPIAuth_TokenConcurrencyLimit(t *testing.T) { t.Parallel() svc := newTestAuthService(t) injectAPIToken(svc, "limited-token", 0, 11) injectAPIToken(svc, "other-token", 0, 12) limitedHash := model.HashToken("limited-token") otherHash := model.HashToken("other-token") svc.authTokensMux.Lock() svc.authTokenMaxConns[limitedHash] = 1 svc.authTokenMaxConns[otherHash] = 1 svc.authTokensMux.Unlock() release, _, _, ok := svc.acquireTokenConcurrencySlot(limitedHash) if !ok { t.Fatal("expected manual slot acquisition to succeed") } defer release() req := httptest.NewRequest(http.MethodGet, "/test", nil) req.Header.Set("Authorization", "Bearer limited-token") w := runMiddleware(t, svc.RequireAPIAuth(), req) if w.Code != http.StatusTooManyRequests { t.Fatalf("expected 429, got %d: %s", w.Code, w.Body.String()) } if !strings.Contains(w.Body.String(), "token_concurrency_exceeded") { t.Fatalf("expected token_concurrency_exceeded in response: %s", w.Body.String()) } otherReq := httptest.NewRequest(http.MethodGet, "/test", nil) otherReq.Header.Set("Authorization", "Bearer other-token") otherW := runMiddleware(t, svc.RequireAPIAuth(), otherReq) if otherW.Code != http.StatusOK { t.Fatalf("expected other token to pass, got %d: %s", otherW.Code, otherW.Body.String()) } } func TestRequireAPIAuth_TokenConcurrencyLimit_AppliesImmediatelyAfterUpdate(t *testing.T) { t.Parallel() svc := newTestAuthService(t) injectAPIToken(svc, "dynamic-token", 0, 21) tokenHash := model.HashToken("dynamic-token") gin.SetMode(gin.TestMode) started := make(chan struct{}) releaseHandler := make(chan struct{}) firstDone := make(chan struct{}) firstW := httptest.NewRecorder() _, engine := gin.CreateTestContext(firstW) engine.Any("/test", svc.RequireAPIAuth(), func(c *gin.Context) { close(started) <-releaseHandler c.JSON(http.StatusOK, gin.H{"passed": true}) }) firstReq := httptest.NewRequest(http.MethodGet, "/test", nil) firstReq.Header.Set("Authorization", "Bearer dynamic-token") go func() { defer close(firstDone) engine.ServeHTTP(firstW, firstReq) }() <-started svc.authTokensMux.Lock() svc.authTokenMaxConns[tokenHash] = 1 svc.authTokensMux.Unlock() secondReq := httptest.NewRequest(http.MethodGet, "/test", nil) secondReq.Header.Set("Authorization", "Bearer dynamic-token") secondW := httptest.NewRecorder() engine.ServeHTTP(secondW, secondReq) if secondW.Code != http.StatusTooManyRequests { t.Fatalf("expected 429 after enabling limit, got %d: %s", secondW.Code, secondW.Body.String()) } if !strings.Contains(secondW.Body.String(), "token_concurrency_exceeded") { t.Fatalf("expected token_concurrency_exceeded in response: %s", secondW.Body.String()) } close(releaseHandler) <-firstDone } // ============================================================================ // RequireTokenAuth 测试 // ============================================================================ func TestRequireTokenAuth_ValidBearer(t *testing.T) { t.Parallel() svc := newTestAuthService(t) injectAdminToken(svc, "admin-token-valid", time.Now().Add(time.Hour)) req := httptest.NewRequest(http.MethodGet, "/test", nil) req.Header.Set("Authorization", "Bearer admin-token-valid") w := runMiddleware(t, svc.RequireTokenAuth(), req) if w.Code != http.StatusOK { t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String()) } } func TestRequireTokenAuth_InvalidBearer(t *testing.T) { t.Parallel() svc := newTestAuthService(t) injectAdminToken(svc, "admin-token", time.Now().Add(time.Hour)) req := httptest.NewRequest(http.MethodGet, "/test", nil) req.Header.Set("Authorization", "Bearer wrong-admin-token") w := runMiddleware(t, svc.RequireTokenAuth(), req) if w.Code != http.StatusUnauthorized { t.Fatalf("expected 401, got %d: %s", w.Code, w.Body.String()) } } func TestRequireTokenAuth_MissingHeader(t *testing.T) { t.Parallel() svc := newTestAuthService(t) injectAdminToken(svc, "admin-token", time.Now().Add(time.Hour)) req := httptest.NewRequest(http.MethodGet, "/test", nil) w := runMiddleware(t, svc.RequireTokenAuth(), req) if w.Code != http.StatusUnauthorized { t.Fatalf("expected 401, got %d: %s", w.Code, w.Body.String()) } } func TestRequireTokenAuth_ExpiredToken(t *testing.T) { t.Parallel() svc := newTestAuthService(t) injectAdminToken(svc, "admin-expired", time.Now().Add(-time.Second)) req := httptest.NewRequest(http.MethodGet, "/test", nil) req.Header.Set("Authorization", "Bearer admin-expired") w := runMiddleware(t, svc.RequireTokenAuth(), req) if w.Code != http.StatusUnauthorized { t.Fatalf("expected 401, got %d: %s", w.Code, w.Body.String()) } // 验证过期 token 已从内存中删除 tokenHash := model.HashToken("admin-expired") svc.tokensMux.RLock() _, stillExists := svc.validTokens[tokenHash] svc.tokensMux.RUnlock() if stillExists { t.Fatal("expected expired admin token to be deleted from memory") } } func TestRequireTokenAuth_NoBearerPrefix(t *testing.T) { t.Parallel() svc := newTestAuthService(t) injectAdminToken(svc, "admin-token", time.Now().Add(time.Hour)) req := httptest.NewRequest(http.MethodGet, "/test", nil) req.Header.Set("Authorization", "admin-token") // 没有 Bearer 前缀 w := runMiddleware(t, svc.RequireTokenAuth(), req) if w.Code != http.StatusUnauthorized { t.Fatalf("expected 401, got %d: %s", w.Code, w.Body.String()) } } func TestRequireAPIAuth_HashDirectMatch(t *testing.T) { t.Parallel() svc := newTestAuthService(t) injectAPIToken(svc, "plaintext-token", 0, 10) // 计算hash,用hash值作为Bearer token发送 hash := model.HashToken("plaintext-token") req := httptest.NewRequest(http.MethodGet, "/test", nil) req.Header.Set("Authorization", "Bearer "+hash) w := runMiddleware(t, svc.RequireAPIAuth(), req) if w.Code != http.StatusOK { t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String()) } // 验证 context 中的 token_hash 和 token_id var resp map[string]any if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil { t.Fatalf("unmarshal response: %v", err) } if got, ok := resp["token_hash"].(string); !ok || got != hash { t.Fatalf("expected token_hash=%s, got=%v", hash, resp["token_hash"]) } if got, ok := resp["token_id"].(float64); !ok || int64(got) != 10 { t.Fatalf("expected token_id=10, got=%v", resp["token_id"]) } } func TestRequireAPIAuth_HashExpired(t *testing.T) { t.Parallel() svc := newTestAuthService(t) expiredAt := time.Now().Add(-time.Hour).UnixMilli() injectAPIToken(svc, "expired-plain", expiredAt, 20) // 用hash值作为Bearer token发送 hash := model.HashToken("expired-plain") req := httptest.NewRequest(http.MethodGet, "/test", nil) req.Header.Set("Authorization", "Bearer "+hash) w := runMiddleware(t, svc.RequireAPIAuth(), req) if w.Code != http.StatusUnauthorized { t.Fatalf("expected 401, got %d: %s", w.Code, w.Body.String()) } // 验证响应包含 "token expired" var resp map[string]string if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil { t.Fatalf("unmarshal response: %v", err) } if resp["error"] != "token expired" { t.Fatalf("expected 'token expired' error, got: %s", resp["error"]) } // 验证懒惰删除:hash应已从内存中移除 svc.authTokensMux.RLock() _, stillExists := svc.authTokens[hash] svc.authTokensMux.RUnlock() if stillExists { t.Fatal("expected expired token to be lazily deleted from memory") } } // TestRequireAPIAuth_TokenPriority 验证 token 提取优先级(Bearer > X-API-Key > x-goog-api-key > query) func TestRequireAPIAuth_TokenPriority(t *testing.T) { t.Parallel() svc := newTestAuthService(t) injectAPIToken(svc, "bearer-token", 0, 1) // 同时设置 Bearer 和 X-API-Key,Bearer 应优先 req := httptest.NewRequest(http.MethodGet, "/test", nil) req.Header.Set("Authorization", "Bearer bearer-token") req.Header.Set("X-API-Key", "wrong-key") w := runMiddleware(t, svc.RequireAPIAuth(), req) if w.Code != http.StatusOK { t.Fatalf("expected 200 (Bearer should take priority), got %d: %s", w.Code, w.Body.String()) } }