prompt_handler_test.go 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249
  1. package securityaudit
  2. import (
  3. "bytes"
  4. "context"
  5. "encoding/json"
  6. "errors"
  7. "net/http"
  8. "net/http/httptest"
  9. "strings"
  10. "testing"
  11. infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
  12. servermiddleware "github.com/Wei-Shaw/sub2api/internal/server/middleware"
  13. "github.com/gin-gonic/gin"
  14. "github.com/stretchr/testify/require"
  15. )
  16. type fakePromptAdminService struct {
  17. config PublicConfig
  18. save func(context.Context, UpdateConfigRequest, int64) (PublicConfig, error)
  19. probe func(context.Context, ProbeRequest) ProbeResult
  20. runtime RuntimeSnapshot
  21. list func(context.Context, EventFilter, int, int) (*EventPage, error)
  22. get func(context.Context, int64) (*Event, error)
  23. deleteOne func(context.Context, int64) (*DeleteResult, error)
  24. deleteIDs func(context.Context, []int64) (*DeleteResult, error)
  25. preview func(context.Context, EventFilter, int64) (*DeletePreview, error)
  26. deleteFilter func(context.Context, DeleteByFilterRequest, int64) (*DeleteResult, error)
  27. }
  28. func (s *fakePromptAdminService) GetConfig() (PublicConfig, error) {
  29. return s.config, nil
  30. }
  31. func (s *fakePromptAdminService) SaveConfig(ctx context.Context, req UpdateConfigRequest, actorID int64) (PublicConfig, error) {
  32. if s.save == nil {
  33. return PublicConfig{}, errors.New("unexpected SaveConfig call")
  34. }
  35. return s.save(ctx, req, actorID)
  36. }
  37. func (s *fakePromptAdminService) Probe(ctx context.Context, req ProbeRequest) ProbeResult {
  38. if s.probe == nil {
  39. return ProbeResult{}
  40. }
  41. return s.probe(ctx, req)
  42. }
  43. func (s *fakePromptAdminService) Runtime(context.Context) RuntimeSnapshot { return s.runtime }
  44. func (s *fakePromptAdminService) ListEvents(ctx context.Context, filter EventFilter, page, pageSize int) (*EventPage, error) {
  45. if s.list == nil {
  46. return &EventPage{}, nil
  47. }
  48. return s.list(ctx, filter, page, pageSize)
  49. }
  50. func (s *fakePromptAdminService) GetEvent(ctx context.Context, id int64) (*Event, error) {
  51. if s.get == nil {
  52. return nil, ErrEventNotFound
  53. }
  54. return s.get(ctx, id)
  55. }
  56. func (s *fakePromptAdminService) DeleteEvent(ctx context.Context, id int64) (*DeleteResult, error) {
  57. if s.deleteOne == nil {
  58. return &DeleteResult{}, nil
  59. }
  60. return s.deleteOne(ctx, id)
  61. }
  62. func (s *fakePromptAdminService) DeleteEventsByIDs(ctx context.Context, ids []int64) (*DeleteResult, error) {
  63. if s.deleteIDs == nil {
  64. return &DeleteResult{}, nil
  65. }
  66. return s.deleteIDs(ctx, ids)
  67. }
  68. func (s *fakePromptAdminService) PreviewDelete(ctx context.Context, filter EventFilter, actorID int64) (*DeletePreview, error) {
  69. if s.preview == nil {
  70. return &DeletePreview{}, nil
  71. }
  72. return s.preview(ctx, filter, actorID)
  73. }
  74. func (s *fakePromptAdminService) DeleteByFilter(ctx context.Context, req DeleteByFilterRequest, actorID int64) (*DeleteResult, error) {
  75. if s.deleteFilter == nil {
  76. return &DeleteResult{}, nil
  77. }
  78. return s.deleteFilter(ctx, req, actorID)
  79. }
  80. func promptAdminRouter(service PromptAdminService) *gin.Engine {
  81. gin.SetMode(gin.TestMode)
  82. router := gin.New()
  83. router.Use(func(c *gin.Context) {
  84. c.Set(string(servermiddleware.ContextKeyUser), servermiddleware.AuthSubject{UserID: 42})
  85. c.Set(string(servermiddleware.ContextKeyUserRole), "admin")
  86. c.Next()
  87. })
  88. handler := NewPromptAdminHandler(service)
  89. group := router.Group("/admin/prompt-audit")
  90. group.GET("/config", handler.GetConfig)
  91. group.PUT("/config", handler.UpdateConfig)
  92. group.POST("/endpoints/probe", handler.ProbeEndpoint)
  93. group.GET("/runtime", handler.GetRuntime)
  94. group.GET("/events", handler.ListEvents)
  95. group.GET("/events/:id", handler.GetEvent)
  96. group.DELETE("/events/:id", handler.DeleteEvent)
  97. group.POST("/events/batch-delete", handler.BatchDelete)
  98. group.POST("/events/delete-preview", handler.DeletePreview)
  99. group.POST("/events/delete-by-filter", handler.DeleteByFilter)
  100. return router
  101. }
  102. func promptAdminRequest(t *testing.T, router http.Handler, method, path string, body any) *httptest.ResponseRecorder {
  103. t.Helper()
  104. var reader *bytes.Reader
  105. if body == nil {
  106. reader = bytes.NewReader(nil)
  107. } else {
  108. raw, err := json.Marshal(body)
  109. require.NoError(t, err)
  110. reader = bytes.NewReader(raw)
  111. }
  112. req := httptest.NewRequest(method, path, reader)
  113. req.Header.Set("Content-Type", "application/json")
  114. recorder := httptest.NewRecorder()
  115. router.ServeHTTP(recorder, req)
  116. return recorder
  117. }
  118. func TestPromptAdminConfigRequiresVersionMapsConflictAndNeverEchoesToken(t *testing.T) {
  119. const canary = "prompt-admin-token-canary"
  120. t.Run("missing expected version", func(t *testing.T) {
  121. router := promptAdminRouter(&fakePromptAdminService{})
  122. response := promptAdminRequest(t, router, http.MethodPut, "/admin/prompt-audit/config", map[string]any{})
  123. require.Equal(t, http.StatusBadRequest, response.Code)
  124. require.Contains(t, response.Body.String(), "prompt_audit_invalid_config_request")
  125. })
  126. t.Run("CAS conflict", func(t *testing.T) {
  127. service := &fakePromptAdminService{save: func(context.Context, UpdateConfigRequest, int64) (PublicConfig, error) {
  128. return PublicConfig{}, infraerrors.Conflict(ErrorCodeConfigConflict, "配置已被更新")
  129. }}
  130. response := promptAdminRequest(t, promptAdminRouter(service), http.MethodPut, "/admin/prompt-audit/config", validHandlerUpdateRequest(canary))
  131. require.Equal(t, http.StatusConflict, response.Code)
  132. require.Contains(t, response.Body.String(), ErrorCodeConfigConflict)
  133. require.NotContains(t, response.Body.String(), canary)
  134. })
  135. t.Run("success public DTO", func(t *testing.T) {
  136. service := &fakePromptAdminService{save: func(_ context.Context, req UpdateConfigRequest, actorID int64) (PublicConfig, error) {
  137. require.Equal(t, int64(42), actorID)
  138. require.Equal(t, canary, req.Endpoints[0].Token)
  139. return PublicConfig{ConfigVersion: 8, Endpoints: []PublicEndpoint{{ID: "guard-1", HasToken: true, TokenStatus: "configured"}}}, nil
  140. }}
  141. response := promptAdminRequest(t, promptAdminRouter(service), http.MethodPut, "/admin/prompt-audit/config", validHandlerUpdateRequest(canary))
  142. require.Equal(t, http.StatusOK, response.Code)
  143. body := response.Body.String()
  144. require.NotContains(t, body, canary)
  145. require.NotContains(t, body, "token_ciphertext")
  146. require.NotContains(t, body, `"token":`)
  147. require.Contains(t, body, `"has_token":true`)
  148. })
  149. }
  150. func TestPromptAdminGetConfigReturnsSecretFreeUnavailableError(t *testing.T) {
  151. const canary = "persisted-config-secret-canary"
  152. repository := &switchableSettingRepository{loadErr: errors.New("failed to load token " + canary)}
  153. manager := NewConfigManager(nil, repository, nil, prefixEncryptor{}, testTotpKeyConfig())
  154. require.Error(t, manager.Reload(context.Background()))
  155. service := &PromptService{config: manager}
  156. response := promptAdminRequest(t, promptAdminRouter(service), http.MethodGet, "/admin/prompt-audit/config", nil)
  157. require.Equal(t, http.StatusServiceUnavailable, response.Code)
  158. require.Contains(t, response.Body.String(), ErrorCodeConfigUnavailable)
  159. require.NotContains(t, response.Body.String(), canary)
  160. require.NotContains(t, response.Body.String(), `"config_version"`)
  161. require.NotContains(t, response.Body.String(), `"token"`)
  162. }
  163. func TestPromptAdminProbeSupportsTemporaryOrSavedTokenWithoutEcho(t *testing.T) {
  164. const canary = "probe-token-canary"
  165. for _, tc := range []struct {
  166. name string
  167. token string
  168. tokenApplied bool
  169. }{
  170. {name: "temporary token", token: canary, tokenApplied: true},
  171. {name: "saved token", token: "", tokenApplied: true},
  172. } {
  173. t.Run(tc.name, func(t *testing.T) {
  174. service := &fakePromptAdminService{probe: func(_ context.Context, req ProbeRequest) ProbeResult {
  175. require.Equal(t, tc.token, req.Endpoint.Token)
  176. return ProbeResult{OK: true, Status: "healthy", Message: "ok", TokenApplied: tc.tokenApplied}
  177. }}
  178. endpoint := validHandlerUpdateRequest(tc.token).Endpoints[0]
  179. response := promptAdminRequest(t, promptAdminRouter(service), http.MethodPost, "/admin/prompt-audit/endpoints/probe", ProbeRequest{Endpoint: endpoint})
  180. require.Equal(t, http.StatusOK, response.Code)
  181. require.NotContains(t, response.Body.String(), canary)
  182. require.NotContains(t, response.Body.String(), `"token":`)
  183. require.Contains(t, response.Body.String(), `"token_applied":true`)
  184. })
  185. }
  186. }
  187. func TestPromptAdminRejectsInvalidEventIDsTimesAndPagination(t *testing.T) {
  188. router := promptAdminRouter(&fakePromptAdminService{})
  189. for _, tc := range []struct {
  190. method string
  191. path string
  192. body any
  193. reason string
  194. }{
  195. {http.MethodGet, "/admin/prompt-audit/events/not-a-number", nil, "prompt_audit_invalid_event_id"},
  196. {http.MethodDelete, "/admin/prompt-audit/events/-1", nil, "prompt_audit_invalid_event_id"},
  197. {http.MethodGet, "/admin/prompt-audit/events?group_id=bad", nil, "prompt_audit_invalid_filter_id"},
  198. {http.MethodGet, "/admin/prompt-audit/events?start_at=not-time", nil, "prompt_audit_invalid_time"},
  199. {http.MethodGet, "/admin/prompt-audit/events?page=0", nil, "prompt_audit_invalid_pagination"},
  200. {http.MethodPost, "/admin/prompt-audit/events/batch-delete", map[string]any{"ids": []int64{1, -2}}, "prompt_audit_invalid_event_id"},
  201. } {
  202. response := promptAdminRequest(t, router, tc.method, tc.path, tc.body)
  203. require.Equalf(t, http.StatusBadRequest, response.Code, "%s %s", tc.method, tc.path)
  204. require.Contains(t, response.Body.String(), tc.reason)
  205. }
  206. }
  207. func validHandlerUpdateRequest(token string) UpdateConfigRequest {
  208. return UpdateConfigRequest{
  209. ExpectedConfigVersion: 7,
  210. Strategy: "priority",
  211. WorkerCount: 1,
  212. QueueCapacity: 10,
  213. Scanners: []string{"pii"},
  214. AllGroups: true,
  215. Endpoints: []UpdateEndpoint{{
  216. ID: "guard-1", Name: "Guard One", Protocol: "openai_compatible",
  217. BaseURL: "http://127.0.0.1:18080", Model: DefaultGuardModel, Token: token,
  218. TimeoutMS: 1000, InputLimit: 1024, Enabled: true,
  219. }},
  220. }
  221. }
  222. func TestPromptAdminDeleteConfirmationErrorsStayGeneric(t *testing.T) {
  223. service := &fakePromptAdminService{deleteFilter: func(context.Context, DeleteByFilterRequest, int64) (*DeleteResult, error) {
  224. return nil, errors.New("sensitive-token-or-filter-detail")
  225. }}
  226. response := promptAdminRequest(t, promptAdminRouter(service), http.MethodPost, "/admin/prompt-audit/events/delete-by-filter", DeleteByFilterRequest{
  227. SnapshotMaxID: 3, FilterHash: strings.Repeat("a", 64), ConfirmationToken: "secret-confirmation", Confirm: true,
  228. })
  229. require.Equal(t, http.StatusBadRequest, response.Code)
  230. require.Contains(t, response.Body.String(), "prompt_audit_delete_confirmation_invalid")
  231. require.NotContains(t, response.Body.String(), "sensitive-token")
  232. require.NotContains(t, response.Body.String(), "secret-confirmation")
  233. }