prompt_audit_route_coverage_test.go 5.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144
  1. package routes
  2. import (
  3. "net/http"
  4. "net/http/httptest"
  5. "os"
  6. "path/filepath"
  7. "regexp"
  8. "sort"
  9. "strings"
  10. "testing"
  11. "github.com/Wei-Shaw/sub2api/internal/handler"
  12. "github.com/Wei-Shaw/sub2api/internal/securityaudit"
  13. servermiddleware "github.com/Wei-Shaw/sub2api/internal/server/middleware"
  14. "github.com/gin-gonic/gin"
  15. "github.com/stretchr/testify/require"
  16. )
  17. func TestEveryGatewayPOSTRouteIsClassifiedForPromptAuditCoverage(t *testing.T) {
  18. routeSource, err := os.ReadFile("gateway.go")
  19. require.NoError(t, err)
  20. pattern := regexp.MustCompile(`(?:gateway|gemini|r|codexDirect|antigravityV1|antigravityV1Beta)\.POST\("([^"]+)"`)
  21. matches := pattern.FindAllStringSubmatch(string(routeSource), -1)
  22. actual := map[string]struct{}{}
  23. for _, match := range matches {
  24. actual[match[1]] = struct{}{}
  25. }
  26. audited := map[string][]string{
  27. "/messages": {"gateway_handler.go", "openai_gateway_handler.go"},
  28. "/responses": {"gateway_handler_responses.go", "openai_gateway_handler.go"},
  29. "/responses/*subpath": {"gateway_handler_responses.go", "openai_gateway_handler.go"},
  30. "/chat/completions": {"gateway_handler_chat_completions.go", "openai_chat_completions.go"},
  31. "/embeddings": {"openai_embeddings.go"},
  32. "/alpha/search": {"openai_alpha_search.go"},
  33. "/live": {"openai_live.go"},
  34. "/realtime/calls": {"openai_live.go"},
  35. "/images/generations": {"openai_images.go", "grok_media.go"},
  36. "/images/edits": {"openai_images.go", "grok_media.go"},
  37. "/images/generations/async": {"image_task_handler.go"},
  38. "/images/edits/async": {"image_task_handler.go"},
  39. "/images/batches": {"batch_image_handler.go"},
  40. "/videos": {"grok_media.go"},
  41. "/videos/generations": {"grok_media.go"},
  42. "/videos/edits": {"grok_media.go"},
  43. "/videos/extensions": {"grok_media.go"},
  44. "/models/*modelAction": {"gemini_v1beta_handler.go"},
  45. "/tts": {"grok_audio.go"},
  46. "/web_search": {"gateway_web_search.go"},
  47. "/x_search": {"gateway_web_search.go"},
  48. }
  49. excluded := map[string]string{
  50. "/messages/count_tokens": "tokenization only; it does not execute a model request",
  51. "/images/batches/:id/cancel": "control-plane cancellation with no user prompt",
  52. "/stt": "speech transcription is not a text-generation prompt",
  53. "/custom-voices": "voice profile management has no model prompt",
  54. }
  55. unclassified := make([]string, 0)
  56. for route := range actual {
  57. if _, ok := audited[route]; ok {
  58. continue
  59. }
  60. if _, ok := excluded[route]; ok {
  61. continue
  62. }
  63. unclassified = append(unclassified, route)
  64. }
  65. sort.Strings(unclassified)
  66. require.Empty(t, unclassified, "new gateway POST routes must be audited or explicitly classified with a no-prompt reason")
  67. for route, files := range audited {
  68. _, exists := actual[route]
  69. require.Truef(t, exists, "stale prompt-audit route manifest entry %s", route)
  70. for _, filename := range files {
  71. source, readErr := os.ReadFile(filepath.Join("..", "..", "handler", filename))
  72. require.NoError(t, readErr)
  73. require.Containsf(t, string(source), "checkSecurityAudit", "%s route handler %s bypasses Coordinator", route, filename)
  74. }
  75. }
  76. for route, reason := range excluded {
  77. require.NotEmpty(t, strings.TrimSpace(reason))
  78. _, exists := actual[route]
  79. require.Truef(t, exists, "stale excluded route %s", route)
  80. }
  81. }
  82. func TestResponsesWebSocketHasFirstAndSubsequentTurnPromptGates(t *testing.T) {
  83. routeSource, err := os.ReadFile("gateway.go")
  84. require.NoError(t, err)
  85. require.GreaterOrEqual(t, strings.Count(string(routeSource), `.GET("/responses"`), 2)
  86. handlerSource, err := os.ReadFile(filepath.Join("..", "..", "handler", "openai_gateway_handler.go"))
  87. require.NoError(t, err)
  88. require.Contains(t, string(handlerSource), `checkSecurityAuditStage`)
  89. require.Contains(t, string(handlerSource), `"first_turn"`)
  90. require.Contains(t, string(handlerSource), `"subsequent_turn"`)
  91. wsStart := strings.Index(string(handlerSource), `func (h *OpenAIGatewayHandler) ResponsesWebSocket`)
  92. require.NotEqual(t, -1, wsStart)
  93. wsSource := string(handlerSource)[wsStart:]
  94. require.Less(t,
  95. strings.Index(wsSource, `"first_turn"`),
  96. strings.Index(wsSource, `TryAcquireUserSlotForAPIKey`),
  97. "the first response.create gate must precede per-request user/account slots",
  98. )
  99. }
  100. func TestPromptAuditAdminRoutesRejectUnauthenticatedAndNonAdminRequests(t *testing.T) {
  101. gin.SetMode(gin.TestMode)
  102. router := gin.New()
  103. handlers := &handler.Handlers{Admin: &handler.AdminHandlers{
  104. PromptAudit: securityaudit.NewPromptAdminHandler(nil),
  105. }}
  106. adminAuth := servermiddleware.AdminAuthMiddleware(func(c *gin.Context) {
  107. if c.GetHeader("Authorization") == "" {
  108. servermiddleware.AbortWithError(c, http.StatusUnauthorized, "UNAUTHORIZED", "Authorization required")
  109. return
  110. }
  111. servermiddleware.AbortWithError(c, http.StatusForbidden, "FORBIDDEN", "Admin access required")
  112. })
  113. auditLog := servermiddleware.AuditLogMiddleware(func(c *gin.Context) { c.Next() })
  114. stepUp := servermiddleware.StepUpAuthMiddleware(func(c *gin.Context) { c.Next() })
  115. RegisterAdminRoutes(router.Group("/api/v1"), handlers, adminAuth, auditLog, stepUp, nil, nil, nil)
  116. for _, tc := range []struct {
  117. name string
  118. auth string
  119. wantStatus int
  120. }{
  121. {name: "unauthenticated", wantStatus: http.StatusUnauthorized},
  122. {name: "non-admin", auth: "Bearer user-token", wantStatus: http.StatusForbidden},
  123. } {
  124. t.Run(tc.name, func(t *testing.T) {
  125. recorder := httptest.NewRecorder()
  126. request := httptest.NewRequest(http.MethodGet, "/api/v1/admin/prompt-audit/config", nil)
  127. if tc.auth != "" {
  128. request.Header.Set("Authorization", tc.auth)
  129. }
  130. router.ServeHTTP(recorder, request)
  131. require.Equal(t, tc.wantStatus, recorder.Code)
  132. })
  133. }
  134. }