mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-09-01 15:02:58 +08:00
mt21625457
0f7f8a317e
Stop WebSocket follow-up turns from reusing a request-wide audit cache, scan client-controlled instruction fields, fail closed on stale weaker configs, and tighten preview/SSRF controls including persisted request stage. Co-authored-by: Cursor <cursoragent@cursor.com>
64 lines
2.5 KiB
Go
64 lines
2.5 KiB
Go
package handler
|
|
|
|
import (
|
|
"context"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"sync/atomic"
|
|
"testing"
|
|
|
|
"github.com/Wei-Shaw/sub2api/internal/securityaudit"
|
|
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func TestCachesSecurityAuditCompletionSkipsWebSocketStages(t *testing.T) {
|
|
require.True(t, cachesSecurityAuditCompletion("http"))
|
|
require.True(t, cachesSecurityAuditCompletion(""))
|
|
require.False(t, cachesSecurityAuditCompletion("first_turn"))
|
|
require.False(t, cachesSecurityAuditCompletion("subsequent_turn"))
|
|
}
|
|
|
|
func TestRunSecurityAuditDoesNotSkipSubsequentWebSocketTurns(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
engine := &turnCountingEngine{mode: securityaudit.ModeAsync}
|
|
coordinator := securityaudit.NewCoordinator(nil, engine)
|
|
|
|
recorder := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(recorder)
|
|
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
|
|
|
subject := middleware2.AuthSubject{UserID: 7, Concurrency: 1}
|
|
first := runSecurityAudit(c, nil, coordinator, nil, nil, subject, "openai_responses", "gpt-test",
|
|
[]byte(`{"type":"response.create","response":{"input":"benign"}}`), "first_turn")
|
|
require.NotNil(t, first)
|
|
require.True(t, first.AllowNextStage)
|
|
require.Equal(t, int64(1), engine.enqueues.Load())
|
|
_, cached := c.Get(securityAuditCompletedContextKey)
|
|
require.False(t, cached, "WebSocket stages must not set the HTTP completion cache")
|
|
|
|
// Even if an HTTP path previously cached completion on this Context, WS turns
|
|
// must still audit every response.create payload.
|
|
c.Set(securityAuditCompletedContextKey, true)
|
|
|
|
second := runSecurityAudit(c, nil, coordinator, nil, nil, subject, "openai_responses", "gpt-test",
|
|
[]byte(`{"type":"response.create","response":{"input":"malicious follow-up"}}`), "subsequent_turn")
|
|
require.NotNil(t, second)
|
|
require.Equal(t, int64(2), engine.enqueues.Load(), "subsequent WebSocket turns must be audited again")
|
|
}
|
|
|
|
type turnCountingEngine struct {
|
|
mode securityaudit.Mode
|
|
enqueues atomic.Int64
|
|
}
|
|
|
|
func (e *turnCountingEngine) EffectiveMode() securityaudit.Mode { return e.mode }
|
|
func (e *turnCountingEngine) Enqueue(context.Context, securityaudit.Request) error {
|
|
e.enqueues.Add(1)
|
|
return nil
|
|
}
|
|
func (e *turnCountingEngine) Evaluate(context.Context, securityaudit.Request) (*securityaudit.PromptDecision, error) {
|
|
return &securityaudit.PromptDecision{Kind: securityaudit.DecisionAllow, AllowNextStage: true}, nil
|
|
}
|