openteam/server/internal/proxy/authcache_test.go

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