Files
sub2api/backend/internal/service/openai_cyber_session_block_test.go
T
李建琦 6d655c9903
Release / update-version (push) Has been cancelled
Release / build-frontend (push) Has been cancelled
Release / release (push) Has been cancelled
Release / sync-version-file (push) Has been cancelled
CI / shell (push) Canceled after 0s
CI / test (push) Canceled after 0s
CI / frontend (push) Canceled after 0s
CI / golangci-lint (push) Canceled after 0s
Security Scan / backend-security (push) Canceled after 0s
Security Scan / frontend-security (push) Canceled after 0s
Sub2API v1.0 - AI API 网关(二开初始版本,基于上游 Wei-Shaw/sub2api)
2026-08-21 18:30:13 +08:00

215 lines
7.6 KiB
Go

package service
import (
"context"
"errors"
"net/http/httptest"
"strings"
"testing"
"time"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
func newCyberBlockTestCtx(headers map[string]string, body string) (*gin.Context, []byte) {
gin.SetMode(gin.TestMode)
c, _ := gin.CreateTestContext(httptest.NewRecorder())
req := httptest.NewRequest("POST", "/openai/v1/responses", strings.NewReader(body))
for k, v := range headers {
req.Header.Set(k, v)
}
c.Request = req
return c, []byte(body)
}
// TestCyberSessionBlockKey verifies F5a key derivation: explicit session signals
// only (header session_id/conversation_id or body prompt_cache_key), apiKey
// isolated, and EMPTY when no explicit signal (no content-derived fallback —
// "不退化" decision).
func TestCyberSessionBlockKey(t *testing.T) {
c1, b1 := newCyberBlockTestCtx(map[string]string{"session_id": "sess-abc"}, `{}`)
k1 := CyberSessionBlockKey(101, c1, b1)
require.NotEmpty(t, k1)
// Same session, different apiKey → different key (isolation).
c2, b2 := newCyberBlockTestCtx(map[string]string{"session_id": "sess-abc"}, `{}`)
require.NotEqual(t, k1, CyberSessionBlockKey(202, c2, b2))
// Same session + same apiKey → stable key.
c3, b3 := newCyberBlockTestCtx(map[string]string{"session_id": "sess-abc"}, `{}`)
require.Equal(t, k1, CyberSessionBlockKey(101, c3, b3))
// prompt_cache_key in body counts as explicit.
c4, b4 := newCyberBlockTestCtx(nil, `{"prompt_cache_key":"pck-1"}`)
require.NotEmpty(t, CyberSessionBlockKey(101, c4, b4))
// No explicit signal → empty key → caller must skip blocking entirely.
c5, b5 := newCyberBlockTestCtx(nil, `{"input":"hello world"}`)
require.Empty(t, CyberSessionBlockKey(101, c5, b5))
// conversation_id header counts as explicit; key is stable and non-empty.
c6, b6 := newCyberBlockTestCtx(map[string]string{"conversation_id": "conv-xyz"}, `{}`)
k6 := CyberSessionBlockKey(101, c6, b6)
require.NotEmpty(t, k6)
c6b, b6b := newCyberBlockTestCtx(map[string]string{"conversation_id": "conv-xyz"}, `{}`)
require.Equal(t, k6, CyberSessionBlockKey(101, c6b, b6b), "conversation_id key must be stable")
}
// --- fakes ---
type fakeCyberBlockStore struct {
blocked map[string]bool
}
var _ CyberSessionBlockStore = (*fakeCyberBlockStore)(nil)
func (f *fakeCyberBlockStore) SetCyberSessionBlocked(_ context.Context, key string, _ time.Duration) error {
if f.blocked == nil {
f.blocked = map[string]bool{}
}
f.blocked[key] = true
return nil
}
func (f *fakeCyberBlockStore) IsCyberSessionBlocked(_ context.Context, key string) (bool, error) {
return f.blocked[key], nil
}
// fakeSettingRepo is a minimal SettingRepository stub for unit tests.
// Only GetValue is exercised by GetCyberSessionBlockRuntime; all other methods
// panic so accidental calls are caught immediately.
type fakeSettingRepo struct {
vals map[string]string
}
func (r *fakeSettingRepo) GetValue(_ context.Context, key string) (string, error) {
v, ok := r.vals[key]
if !ok {
return "", ErrSettingNotFound
}
return v, nil
}
func (r *fakeSettingRepo) Get(_ context.Context, _ string) (*Setting, error) {
panic("fakeSettingRepo.Get not implemented")
}
func (r *fakeSettingRepo) Set(_ context.Context, _, _ string) error {
panic("fakeSettingRepo.Set not implemented")
}
func (r *fakeSettingRepo) GetMultiple(_ context.Context, _ []string) (map[string]string, error) {
panic("fakeSettingRepo.GetMultiple not implemented")
}
func (r *fakeSettingRepo) SetMultiple(_ context.Context, _ map[string]string) error {
panic("fakeSettingRepo.SetMultiple not implemented")
}
func (r *fakeSettingRepo) GetAll(_ context.Context) (map[string]string, error) {
panic("fakeSettingRepo.GetAll not implemented")
}
func (r *fakeSettingRepo) Delete(_ context.Context, _ string) error {
panic("fakeSettingRepo.Delete not implemented")
}
var _ SettingRepository = (*fakeSettingRepo)(nil)
// comboCacheAndStore implements both GatewayCache (no-op stubs) and
// CyberSessionBlockStore (delegates to fakeCyberBlockStore) so it can be
// injected as s.cache and successfully type-asserted to CyberSessionBlockStore.
type comboCacheAndStore struct {
store fakeCyberBlockStore
}
var _ GatewayCache = (*comboCacheAndStore)(nil)
var _ CyberSessionBlockStore = (*comboCacheAndStore)(nil)
func (c *comboCacheAndStore) GetSessionAccountID(_ context.Context, _ int64, _ string) (int64, error) {
return 0, errors.New("stub")
}
func (c *comboCacheAndStore) SetSessionAccountID(_ context.Context, _ int64, _ string, _ int64, _ time.Duration) error {
return nil
}
func (c *comboCacheAndStore) RefreshSessionTTL(_ context.Context, _ int64, _ string, _ time.Duration) error {
return nil
}
func (c *comboCacheAndStore) DeleteSessionAccountID(_ context.Context, _ int64, _ string) error {
return nil
}
func (c *comboCacheAndStore) SetGrokVideoPendingBilling(_ context.Context, _ string, _ []byte, _ time.Duration) error {
return nil
}
func (c *comboCacheAndStore) GetGrokVideoPendingBilling(_ context.Context, _ string) ([]byte, error) {
return nil, nil
}
func (c *comboCacheAndStore) ClaimGrokVideoBilled(_ context.Context, _ string, _ time.Duration) (bool, error) {
return true, nil
}
func (c *comboCacheAndStore) ReleaseGrokVideoBilled(_ context.Context, _ string) error {
return nil
}
func (c *comboCacheAndStore) SetReasoningContent(_ context.Context, _ string, _ string, _ time.Duration) error {
return nil
}
func (c *comboCacheAndStore) GetReasoningContent(_ context.Context, _ string) (string, error) {
return "", ErrReasoningContentNotFound
}
func (c *comboCacheAndStore) SetCyberSessionBlocked(ctx context.Context, key string, ttl time.Duration) error {
return c.store.SetCyberSessionBlocked(ctx, key, ttl)
}
func (c *comboCacheAndStore) IsCyberSessionBlocked(ctx context.Context, key string) (bool, error) {
return c.store.IsCyberSessionBlocked(ctx, key)
}
// --- tests ---
// TestIsCyberSessionBlocked_EmptyKeyAndNilService covers the fail-open paths:
// empty key, nil service, store missing → always false / no panic.
func TestIsCyberSessionBlocked_EmptyKeyAndNilService(t *testing.T) {
var nilSvc *OpenAIGatewayService
require.False(t, nilSvc.IsCyberSessionBlocked(context.Background(), "k"))
require.NotPanics(t, func() { nilSvc.MarkCyberSessionBlocked(context.Background(), "k") })
svc := &OpenAIGatewayService{}
require.False(t, svc.IsCyberSessionBlocked(context.Background(), ""))
require.False(t, svc.IsCyberSessionBlocked(context.Background(), "k"), "no store + no settings → fail-open false")
}
// TestCyberSessionBlock_RoundTrip exercises the type-assertion success path:
// mark a session blocked via a combo cache+store, then confirm IsCyberSessionBlocked
// returns true, and an unrelated key returns false.
func TestCyberSessionBlock_RoundTrip(t *testing.T) {
// SettingService with only settingRepo set — GetCyberSessionBlockRuntime needs
// nothing else (cfg/proxyRepo/etc. are not touched by this code path).
settingSvc := &SettingService{
settingRepo: &fakeSettingRepo{
vals: map[string]string{
SettingKeyCyberSessionBlockEnabled: "true",
SettingKeyCyberSessionBlockTTLSeconds: "60",
},
},
}
combo := &comboCacheAndStore{}
svc := &OpenAIGatewayService{
cache: combo,
settingService: settingSvc,
}
ctx := context.Background()
const testKey = "deadbeef1234"
// Before marking: not blocked.
require.False(t, svc.IsCyberSessionBlocked(ctx, testKey))
// Mark as blocked.
svc.MarkCyberSessionBlocked(ctx, testKey)
// After marking: blocked.
require.True(t, svc.IsCyberSessionBlocked(ctx, testKey))
// Different key: still not blocked.
require.False(t, svc.IsCyberSessionBlocked(ctx, "other-key"))
}