package proxy import ( "context" "io" "net/http" "net/http/httptest" "sync" "testing" "time" "github.com/openteam/server/internal/channel" "github.com/openteam/server/internal/store" "gorm.io/gorm" gormlogger "gorm.io/gorm/logger" ) // 第二条请求应完全命中鉴权缓存:不再有 api_keys / users 查询。 func TestAuthCacheSkipsQueries(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/json") _, _ = io.WriteString(w, `{"id":"x","model":"m1","choices":[{"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1}}`) })) defer srv.Close() h := newHarness(t, srv.URL, 0, true, nil) h.gw.auth = newAuthCache(time.Minute) tr := &sqlTracer{Interface: gormlogger.Discard} traced := h.db.Session(&gorm.Session{Logger: tr}) h.gw.db = traced h.gw.ch = channel.NewService(traced, h.gw.enc) body := `{"model":"m1","messages":[{"role":"user","content":"hi"}]}` for i := 0; i < 2; i++ { if w := h.run(t, nil, body); w.Code != http.StatusOK { t.Fatalf("第 %d 次 code=%d body=%s", i+1, w.Code, w.Body.String()) } tr.take() } // 第三次:缓存已热,不应再有 SQL if w := h.run(t, nil, body); w.Code != http.StatusOK { t.Fatalf("code=%d", w.Code) } if sqls := tr.take(); len(sqls) != 0 { t.Fatalf("缓存命中后不应再查库,实际 %d 条:%v", len(sqls), sqls) } } // 缓存值可能是 TTL 内的旧状态:密钥/用户被停用后,缓存命中也要被拒。 func TestAuthCacheRechecksStatus(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { _, _ = io.WriteString(w, `{"id":"x","choices":[{"message":{"content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1}}`) })) defer srv.Close() h := newHarness(t, srv.URL, 0, true, nil) h.gw.auth = newAuthCache(time.Minute) body := `{"model":"m1","messages":[{"role":"user","content":"hi"}]}` if w := h.run(t, nil, body); w.Code != http.StatusOK { t.Fatalf("首选请求应成功:%d", w.Code) } // 直接改缓存里的状态(模拟"库里已停用、缓存尚未过期") h.gw.auth.mu.Lock() for k, e := range h.gw.auth.m { e.user.Status = store.UserStatusDisabled h.gw.auth.m[k] = e } h.gw.auth.mu.Unlock() if w := h.run(t, nil, body); w.Code != http.StatusForbidden { t.Fatalf("缓存里的停用用户也应被拒,实际 %d", w.Code) } } // 密钥过期按当前时间判定,不受缓存影响。 func TestAuthCacheRespectsKeyExpiry(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { _, _ = io.WriteString(w, `{"id":"x","choices":[{"message":{"content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1}}`) })) defer srv.Close() h := newHarness(t, srv.URL, 0, true, nil) h.gw.auth = newAuthCache(time.Minute) body := `{"model":"m1","messages":[{"role":"user","content":"hi"}]}` if w := h.run(t, nil, body); w.Code != http.StatusOK { t.Fatalf("首选请求应成功:%d", w.Code) } past := time.Now().Add(-time.Minute) h.gw.auth.mu.Lock() for k, e := range h.gw.auth.m { e.key.ExpiresAt = &past h.gw.auth.m[k] = e } h.gw.auth.mu.Unlock() if w := h.run(t, nil, body); w.Code != http.StatusUnauthorized { t.Fatalf("过期密钥应 401,实际 %d", w.Code) } } // InvalidateAuth 之后必须重新查库(吊销立刻生效)。 func TestAuthCacheInvalidate(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { _, _ = io.WriteString(w, `{"id":"x","choices":[{"message":{"content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1}}`) })) defer srv.Close() h := newHarness(t, srv.URL, 0, true, nil) h.gw.auth = newAuthCache(time.Minute) body := `{"model":"m1","messages":[{"role":"user","content":"hi"}]}` if w := h.run(t, nil, body); w.Code != http.StatusOK { t.Fatalf("首选请求应成功:%d", w.Code) } // 吊销密钥后清缓存:下一次请求应 401 h.db.Model(&store.APIKey{}).Where("id > 0").Update("status", store.KeyStatusRevoked) h.gw.InvalidateAuth() if w := h.run(t, nil, body); w.Code != http.StatusUnauthorized { t.Fatalf("吊销并清缓存后应 401,实际 %d", w.Code) } } // TTL 到期后重新查库:直接改库(绕过失效)也能在一个 TTL 内生效。 func TestAuthCacheTTLExpiry(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { _, _ = io.WriteString(w, `{"id":"x","choices":[{"message":{"content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1}}`) })) defer srv.Close() h := newHarness(t, srv.URL, 0, true, nil) h.gw.auth = newAuthCache(50 * time.Millisecond) body := `{"model":"m1","messages":[{"role":"user","content":"hi"}]}` if w := h.run(t, nil, body); w.Code != http.StatusOK { t.Fatalf("首选请求应成功:%d", w.Code) } h.db.Model(&store.APIKey{}).Where("id > 0").Update("status", store.KeyStatusRevoked) time.Sleep(70 * time.Millisecond) if w := h.run(t, nil, body); w.Code != http.StatusUnauthorized { t.Fatalf("TTL 到期后应重新查库并 401,实际 %d", w.Code) } } // 关闭缓存(TTL=0)时行为不变:每次都查库。 func TestAuthCacheDisabled(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { _, _ = io.WriteString(w, `{"id":"x","choices":[{"message":{"content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1}}`) })) defer srv.Close() h := newHarness(t, srv.URL, 0, true, nil) h.gw.auth = newAuthCache(0) body := `{"model":"m1","messages":[{"role":"user","content":"hi"}]}` if w := h.run(t, nil, body); w.Code != http.StatusOK { t.Fatalf("code=%d", w.Code) } if h.gw.auth.on() { t.Fatal("TTL=0 时缓存应关闭") } } // 并发访问(-race 验证) func TestAuthCacheConcurrent(t *testing.T) { c := newAuthCache(time.Minute) var wg sync.WaitGroup for i := 0; i < 8; i++ { wg.Add(1) go func(n int) { defer wg.Done() for j := 0; j < 200; j++ { h := "hash" c.put(h, store.APIKey{ID: uint64(n)}, store.User{ID: uint64(n)}) c.get(h) if j%50 == 0 { c.clear() } } }(i) } wg.Wait() } var _ = context.Background