| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249 |
- package securityaudit
- import (
- "bytes"
- "context"
- "encoding/json"
- "errors"
- "net/http"
- "net/http/httptest"
- "strings"
- "testing"
- infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
- servermiddleware "github.com/Wei-Shaw/sub2api/internal/server/middleware"
- "github.com/gin-gonic/gin"
- "github.com/stretchr/testify/require"
- )
- type fakePromptAdminService struct {
- config PublicConfig
- save func(context.Context, UpdateConfigRequest, int64) (PublicConfig, error)
- probe func(context.Context, ProbeRequest) ProbeResult
- runtime RuntimeSnapshot
- list func(context.Context, EventFilter, int, int) (*EventPage, error)
- get func(context.Context, int64) (*Event, error)
- deleteOne func(context.Context, int64) (*DeleteResult, error)
- deleteIDs func(context.Context, []int64) (*DeleteResult, error)
- preview func(context.Context, EventFilter, int64) (*DeletePreview, error)
- deleteFilter func(context.Context, DeleteByFilterRequest, int64) (*DeleteResult, error)
- }
- func (s *fakePromptAdminService) GetConfig() (PublicConfig, error) {
- return s.config, nil
- }
- func (s *fakePromptAdminService) SaveConfig(ctx context.Context, req UpdateConfigRequest, actorID int64) (PublicConfig, error) {
- if s.save == nil {
- return PublicConfig{}, errors.New("unexpected SaveConfig call")
- }
- return s.save(ctx, req, actorID)
- }
- func (s *fakePromptAdminService) Probe(ctx context.Context, req ProbeRequest) ProbeResult {
- if s.probe == nil {
- return ProbeResult{}
- }
- return s.probe(ctx, req)
- }
- func (s *fakePromptAdminService) Runtime(context.Context) RuntimeSnapshot { return s.runtime }
- func (s *fakePromptAdminService) ListEvents(ctx context.Context, filter EventFilter, page, pageSize int) (*EventPage, error) {
- if s.list == nil {
- return &EventPage{}, nil
- }
- return s.list(ctx, filter, page, pageSize)
- }
- func (s *fakePromptAdminService) GetEvent(ctx context.Context, id int64) (*Event, error) {
- if s.get == nil {
- return nil, ErrEventNotFound
- }
- return s.get(ctx, id)
- }
- func (s *fakePromptAdminService) DeleteEvent(ctx context.Context, id int64) (*DeleteResult, error) {
- if s.deleteOne == nil {
- return &DeleteResult{}, nil
- }
- return s.deleteOne(ctx, id)
- }
- func (s *fakePromptAdminService) DeleteEventsByIDs(ctx context.Context, ids []int64) (*DeleteResult, error) {
- if s.deleteIDs == nil {
- return &DeleteResult{}, nil
- }
- return s.deleteIDs(ctx, ids)
- }
- func (s *fakePromptAdminService) PreviewDelete(ctx context.Context, filter EventFilter, actorID int64) (*DeletePreview, error) {
- if s.preview == nil {
- return &DeletePreview{}, nil
- }
- return s.preview(ctx, filter, actorID)
- }
- func (s *fakePromptAdminService) DeleteByFilter(ctx context.Context, req DeleteByFilterRequest, actorID int64) (*DeleteResult, error) {
- if s.deleteFilter == nil {
- return &DeleteResult{}, nil
- }
- return s.deleteFilter(ctx, req, actorID)
- }
- func promptAdminRouter(service PromptAdminService) *gin.Engine {
- gin.SetMode(gin.TestMode)
- router := gin.New()
- router.Use(func(c *gin.Context) {
- c.Set(string(servermiddleware.ContextKeyUser), servermiddleware.AuthSubject{UserID: 42})
- c.Set(string(servermiddleware.ContextKeyUserRole), "admin")
- c.Next()
- })
- handler := NewPromptAdminHandler(service)
- group := router.Group("/admin/prompt-audit")
- group.GET("/config", handler.GetConfig)
- group.PUT("/config", handler.UpdateConfig)
- group.POST("/endpoints/probe", handler.ProbeEndpoint)
- group.GET("/runtime", handler.GetRuntime)
- group.GET("/events", handler.ListEvents)
- group.GET("/events/:id", handler.GetEvent)
- group.DELETE("/events/:id", handler.DeleteEvent)
- group.POST("/events/batch-delete", handler.BatchDelete)
- group.POST("/events/delete-preview", handler.DeletePreview)
- group.POST("/events/delete-by-filter", handler.DeleteByFilter)
- return router
- }
- func promptAdminRequest(t *testing.T, router http.Handler, method, path string, body any) *httptest.ResponseRecorder {
- t.Helper()
- var reader *bytes.Reader
- if body == nil {
- reader = bytes.NewReader(nil)
- } else {
- raw, err := json.Marshal(body)
- require.NoError(t, err)
- reader = bytes.NewReader(raw)
- }
- req := httptest.NewRequest(method, path, reader)
- req.Header.Set("Content-Type", "application/json")
- recorder := httptest.NewRecorder()
- router.ServeHTTP(recorder, req)
- return recorder
- }
- func TestPromptAdminConfigRequiresVersionMapsConflictAndNeverEchoesToken(t *testing.T) {
- const canary = "prompt-admin-token-canary"
- t.Run("missing expected version", func(t *testing.T) {
- router := promptAdminRouter(&fakePromptAdminService{})
- response := promptAdminRequest(t, router, http.MethodPut, "/admin/prompt-audit/config", map[string]any{})
- require.Equal(t, http.StatusBadRequest, response.Code)
- require.Contains(t, response.Body.String(), "prompt_audit_invalid_config_request")
- })
- t.Run("CAS conflict", func(t *testing.T) {
- service := &fakePromptAdminService{save: func(context.Context, UpdateConfigRequest, int64) (PublicConfig, error) {
- return PublicConfig{}, infraerrors.Conflict(ErrorCodeConfigConflict, "配置已被更新")
- }}
- response := promptAdminRequest(t, promptAdminRouter(service), http.MethodPut, "/admin/prompt-audit/config", validHandlerUpdateRequest(canary))
- require.Equal(t, http.StatusConflict, response.Code)
- require.Contains(t, response.Body.String(), ErrorCodeConfigConflict)
- require.NotContains(t, response.Body.String(), canary)
- })
- t.Run("success public DTO", func(t *testing.T) {
- service := &fakePromptAdminService{save: func(_ context.Context, req UpdateConfigRequest, actorID int64) (PublicConfig, error) {
- require.Equal(t, int64(42), actorID)
- require.Equal(t, canary, req.Endpoints[0].Token)
- return PublicConfig{ConfigVersion: 8, Endpoints: []PublicEndpoint{{ID: "guard-1", HasToken: true, TokenStatus: "configured"}}}, nil
- }}
- response := promptAdminRequest(t, promptAdminRouter(service), http.MethodPut, "/admin/prompt-audit/config", validHandlerUpdateRequest(canary))
- require.Equal(t, http.StatusOK, response.Code)
- body := response.Body.String()
- require.NotContains(t, body, canary)
- require.NotContains(t, body, "token_ciphertext")
- require.NotContains(t, body, `"token":`)
- require.Contains(t, body, `"has_token":true`)
- })
- }
- func TestPromptAdminGetConfigReturnsSecretFreeUnavailableError(t *testing.T) {
- const canary = "persisted-config-secret-canary"
- repository := &switchableSettingRepository{loadErr: errors.New("failed to load token " + canary)}
- manager := NewConfigManager(nil, repository, nil, prefixEncryptor{}, testTotpKeyConfig())
- require.Error(t, manager.Reload(context.Background()))
- service := &PromptService{config: manager}
- response := promptAdminRequest(t, promptAdminRouter(service), http.MethodGet, "/admin/prompt-audit/config", nil)
- require.Equal(t, http.StatusServiceUnavailable, response.Code)
- require.Contains(t, response.Body.String(), ErrorCodeConfigUnavailable)
- require.NotContains(t, response.Body.String(), canary)
- require.NotContains(t, response.Body.String(), `"config_version"`)
- require.NotContains(t, response.Body.String(), `"token"`)
- }
- func TestPromptAdminProbeSupportsTemporaryOrSavedTokenWithoutEcho(t *testing.T) {
- const canary = "probe-token-canary"
- for _, tc := range []struct {
- name string
- token string
- tokenApplied bool
- }{
- {name: "temporary token", token: canary, tokenApplied: true},
- {name: "saved token", token: "", tokenApplied: true},
- } {
- t.Run(tc.name, func(t *testing.T) {
- service := &fakePromptAdminService{probe: func(_ context.Context, req ProbeRequest) ProbeResult {
- require.Equal(t, tc.token, req.Endpoint.Token)
- return ProbeResult{OK: true, Status: "healthy", Message: "ok", TokenApplied: tc.tokenApplied}
- }}
- endpoint := validHandlerUpdateRequest(tc.token).Endpoints[0]
- response := promptAdminRequest(t, promptAdminRouter(service), http.MethodPost, "/admin/prompt-audit/endpoints/probe", ProbeRequest{Endpoint: endpoint})
- require.Equal(t, http.StatusOK, response.Code)
- require.NotContains(t, response.Body.String(), canary)
- require.NotContains(t, response.Body.String(), `"token":`)
- require.Contains(t, response.Body.String(), `"token_applied":true`)
- })
- }
- }
- func TestPromptAdminRejectsInvalidEventIDsTimesAndPagination(t *testing.T) {
- router := promptAdminRouter(&fakePromptAdminService{})
- for _, tc := range []struct {
- method string
- path string
- body any
- reason string
- }{
- {http.MethodGet, "/admin/prompt-audit/events/not-a-number", nil, "prompt_audit_invalid_event_id"},
- {http.MethodDelete, "/admin/prompt-audit/events/-1", nil, "prompt_audit_invalid_event_id"},
- {http.MethodGet, "/admin/prompt-audit/events?group_id=bad", nil, "prompt_audit_invalid_filter_id"},
- {http.MethodGet, "/admin/prompt-audit/events?start_at=not-time", nil, "prompt_audit_invalid_time"},
- {http.MethodGet, "/admin/prompt-audit/events?page=0", nil, "prompt_audit_invalid_pagination"},
- {http.MethodPost, "/admin/prompt-audit/events/batch-delete", map[string]any{"ids": []int64{1, -2}}, "prompt_audit_invalid_event_id"},
- } {
- response := promptAdminRequest(t, router, tc.method, tc.path, tc.body)
- require.Equalf(t, http.StatusBadRequest, response.Code, "%s %s", tc.method, tc.path)
- require.Contains(t, response.Body.String(), tc.reason)
- }
- }
- func validHandlerUpdateRequest(token string) UpdateConfigRequest {
- return UpdateConfigRequest{
- ExpectedConfigVersion: 7,
- Strategy: "priority",
- WorkerCount: 1,
- QueueCapacity: 10,
- Scanners: []string{"pii"},
- AllGroups: true,
- Endpoints: []UpdateEndpoint{{
- ID: "guard-1", Name: "Guard One", Protocol: "openai_compatible",
- BaseURL: "http://127.0.0.1:18080", Model: DefaultGuardModel, Token: token,
- TimeoutMS: 1000, InputLimit: 1024, Enabled: true,
- }},
- }
- }
- func TestPromptAdminDeleteConfirmationErrorsStayGeneric(t *testing.T) {
- service := &fakePromptAdminService{deleteFilter: func(context.Context, DeleteByFilterRequest, int64) (*DeleteResult, error) {
- return nil, errors.New("sensitive-token-or-filter-detail")
- }}
- response := promptAdminRequest(t, promptAdminRouter(service), http.MethodPost, "/admin/prompt-audit/events/delete-by-filter", DeleteByFilterRequest{
- SnapshotMaxID: 3, FilterHash: strings.Repeat("a", 64), ConfirmationToken: "secret-confirmation", Confirm: true,
- })
- require.Equal(t, http.StatusBadRequest, response.Code)
- require.Contains(t, response.Body.String(), "prompt_audit_delete_confirmation_invalid")
- require.NotContains(t, response.Body.String(), "sensitive-token")
- require.NotContains(t, response.Body.String(), "secret-confirmation")
- }
|