上游对单次请求下发 error.code=cyber_policy 硬阻断时,网关在所有端点 (/v1/responses、/v1/chat/completions、/v1/messages、WebSocket)及流式/ 非流式路径下,将该结果原样透传给客户端,绝不 failover、换号或同步拦截; 命中后异步完成审计与计费: - 风控中心记录 cyber_policy 留痕并发送通知邮件,落库先于发信,SMTP 阻塞 不影响留痕 - ops 错误请求记录,状态码对齐客户端实际接收(流式 200 / 非流式 400) - 用量明细标记 request_type=cyber,按上游真实 token 计费,HTTP 与 WebSocket 计费口径统一,零 token 命中不误扣 - 会话级自动屏蔽(管理员开关,默认关):命中的会话在可配 TTL 内本地拦截 不再发往上游,仅屏蔽该会话不影响同 Key 其他会话 - 封号计数排除开关:可选让 cyber 命中不计入自动封号,命中当次不判定且 历史行在违规计数中一并排除 WebSocket 多轮连接下 cyber 标记按 turn 生命周期管理,逐轮独立检测与记录; 透传的错误响应不被兜底逻辑追加内容污染。 Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
193 lines
6.8 KiB
Go
193 lines
6.8 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) 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"))
|
|
}
|