security_audit_errors.go 6.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187
  1. package handler
  2. import (
  3. "context"
  4. "encoding/json"
  5. "net/http"
  6. "strings"
  7. "time"
  8. "github.com/Wei-Shaw/sub2api/internal/pkg/googleapi"
  9. "github.com/Wei-Shaw/sub2api/internal/securityaudit"
  10. "github.com/Wei-Shaw/sub2api/internal/service"
  11. coderws "github.com/coder/websocket"
  12. "github.com/gin-gonic/gin"
  13. )
  14. func (h *OpenAIGatewayHandler) openAISecurityAuditError(c *gin.Context, decision *securityaudit.Decision) {
  15. if decision == nil {
  16. return
  17. }
  18. if decision.Legacy != nil && decision.Legacy.Blocked {
  19. h.errorResponse(c, securityAuditStatus(decision), securityAuditErrorCode(decision), securityAuditMessage(decision))
  20. return
  21. }
  22. errType := "api_error"
  23. if decision.Kind == securityaudit.DecisionBlock {
  24. errType = "permission_error"
  25. }
  26. c.JSON(securityAuditStatus(decision), gin.H{"error": gin.H{
  27. "type": errType, "code": securityAuditErrorCode(decision), "message": securityAuditMessage(decision),
  28. }})
  29. }
  30. func (h *GatewayHandler) openAISecurityAuditError(c *gin.Context, decision *securityaudit.Decision) {
  31. if decision == nil {
  32. return
  33. }
  34. if decision.Legacy != nil && decision.Legacy.Blocked {
  35. h.chatCompletionsErrorResponse(c, securityAuditStatus(decision), securityAuditErrorCode(decision), securityAuditMessage(decision))
  36. return
  37. }
  38. errType := "api_error"
  39. if decision.Kind == securityaudit.DecisionBlock {
  40. errType = "permission_error"
  41. }
  42. c.JSON(securityAuditStatus(decision), gin.H{"error": gin.H{
  43. "type": errType, "code": securityAuditErrorCode(decision), "message": securityAuditMessage(decision),
  44. }})
  45. }
  46. func (h *GatewayHandler) responsesSecurityAuditError(c *gin.Context, decision *securityaudit.Decision) {
  47. if decision == nil {
  48. return
  49. }
  50. if decision.Legacy != nil && decision.Legacy.Blocked {
  51. h.responsesErrorResponse(c, securityAuditStatus(decision), securityAuditErrorCode(decision), securityAuditMessage(decision))
  52. return
  53. }
  54. c.JSON(securityAuditStatus(decision), gin.H{"error": gin.H{
  55. "type": "api_error", "code": securityAuditErrorCode(decision), "message": securityAuditMessage(decision),
  56. }})
  57. }
  58. func (h *GatewayHandler) anthropicSecurityAuditError(c *gin.Context, decision *securityaudit.Decision) {
  59. if decision == nil {
  60. return
  61. }
  62. if decision.Legacy != nil && decision.Legacy.Blocked {
  63. h.errorResponse(c, securityAuditStatus(decision), securityAuditErrorCode(decision), securityAuditMessage(decision))
  64. return
  65. }
  66. errType := "api_error"
  67. if decision.Kind == securityaudit.DecisionBlock {
  68. errType = "permission_error"
  69. }
  70. c.JSON(securityAuditStatus(decision), gin.H{"type": "error", "error": gin.H{
  71. "type": errType, "code": securityAuditErrorCode(decision), "message": securityAuditMessage(decision),
  72. }})
  73. }
  74. func (h *OpenAIGatewayHandler) anthropicSecurityAuditError(c *gin.Context, decision *securityaudit.Decision) {
  75. if decision == nil {
  76. return
  77. }
  78. if decision.Legacy != nil && decision.Legacy.Blocked {
  79. h.anthropicErrorResponse(c, securityAuditStatus(decision), securityAuditErrorCode(decision), securityAuditMessage(decision))
  80. return
  81. }
  82. errType := "api_error"
  83. if decision.Kind == securityaudit.DecisionBlock {
  84. errType = "permission_error"
  85. }
  86. c.JSON(securityAuditStatus(decision), gin.H{"type": "error", "error": gin.H{
  87. "type": errType, "code": securityAuditErrorCode(decision), "message": securityAuditMessage(decision),
  88. }})
  89. }
  90. func googleSecurityAuditError(c *gin.Context, decision *securityaudit.Decision) {
  91. if decision == nil {
  92. return
  93. }
  94. if decision.Legacy != nil && decision.Legacy.Blocked {
  95. googleError(c, securityAuditStatus(decision), securityAuditMessage(decision))
  96. return
  97. }
  98. status := securityAuditStatus(decision)
  99. googleStatus := googleapi.HTTPStatusToGoogleStatus(status)
  100. if status == http.StatusServiceUnavailable {
  101. googleStatus = "UNAVAILABLE"
  102. }
  103. requestID := ""
  104. if c != nil && c.Request != nil {
  105. requestID = contentModerationRequestID(c.Request.Context())
  106. }
  107. c.JSON(status, gin.H{"error": gin.H{
  108. "code": status, "message": securityAuditMessage(decision), "status": googleStatus,
  109. "details": []gin.H{{
  110. "@type": "type.googleapis.com/google.rpc.ErrorInfo",
  111. "reason": securityAuditErrorCode(decision), "domain": "sub2api.securityaudit",
  112. "metadata": gin.H{"request_id": requestID},
  113. }},
  114. }})
  115. }
  116. func writeSecurityAuditWSError(ctx context.Context, conn *coderws.Conn, decision *securityaudit.Decision) {
  117. if conn == nil || decision == nil {
  118. return
  119. }
  120. if decision.Legacy != nil && decision.Legacy.Blocked {
  121. legacy := decision.Legacy
  122. writeContentModerationWSError(ctx, conn, (legacyContentModerationDecision{legacy}).toService())
  123. return
  124. }
  125. if ctx == nil {
  126. ctx = context.Background()
  127. }
  128. payload, err := json.Marshal(gin.H{
  129. "event_id": "evt_prompt_guard_rejected", "type": "error",
  130. "error": gin.H{"type": "invalid_request_error", "code": securityAuditErrorCode(decision), "message": securityAuditMessage(decision)},
  131. })
  132. if err != nil {
  133. return
  134. }
  135. writeCtx, cancel := context.WithTimeout(ctx, 2*time.Second)
  136. defer cancel()
  137. _ = conn.Write(writeCtx, coderws.MessageText, payload)
  138. }
  139. type legacyContentModerationDecision struct{ value *securityaudit.LegacyDecision }
  140. func (d legacyContentModerationDecision) toService() *service.ContentModerationDecision {
  141. if d.value == nil {
  142. return nil
  143. }
  144. return &service.ContentModerationDecision{Allowed: d.value.Allowed, Blocked: d.value.Blocked, Flagged: d.value.Flagged, Message: d.value.Message, StatusCode: d.value.StatusCode, Action: d.value.Action}
  145. }
  146. func securityAuditWSCloseStatus(decision *securityaudit.Decision) coderws.StatusCode {
  147. if decision == nil {
  148. return coderws.StatusInternalError
  149. }
  150. if decision.Legacy != nil && decision.Legacy.Blocked {
  151. return coderws.StatusPolicyViolation
  152. }
  153. if decision.Kind == securityaudit.DecisionBlock {
  154. return coderws.StatusCode(4403)
  155. }
  156. return coderws.StatusTryAgainLater
  157. }
  158. func securityAuditWSCloseReason(decision *securityaudit.Decision) string {
  159. if decision == nil {
  160. return securityaudit.ErrorCodeUnavailable
  161. }
  162. if decision.Legacy != nil && decision.Legacy.Blocked {
  163. message := strings.TrimSpace(decision.Legacy.Message)
  164. if message != "" {
  165. return message
  166. }
  167. return "content_policy_violation"
  168. }
  169. code := securityAuditErrorCode(decision)
  170. if code == "" {
  171. return securityaudit.ErrorCodeUnavailable
  172. }
  173. return code
  174. }