mirror of
https://github.com/alibaba/higress.git
synced 2026-05-08 04:17:27 +08:00
feat(ai-security-guard): replace denyMessage with structured DenyResponseBody (#3642)
Co-authored-by: rinfx <yucheng.lxr@alibaba-inc.com>
This commit is contained in:
@@ -41,15 +41,43 @@ description: 阿里云内容安全检测
|
||||
| `consumerResponseCheckService` | map | optional | - | 为不同消费者指定特定的响应检测服务 |
|
||||
| `consumerRiskLevel` | map | optional | - | 为不同消费者指定各维度的拦截风险等级 |
|
||||
|
||||
补充说明一下 `denyMessage`,对非法请求的处理逻辑为:
|
||||
- 如果配置了 `denyMessage`,返回内容为 `denyMessage` 配置内容,格式为openai格式的流式/非流式响应
|
||||
- 如果没有配置 `denyMessage`,优先返回阿里云内容安全的建议回答,格式为openai格式的流式/非流式响应
|
||||
- 如果阿里云内容安全未返回建议的回答,返回内容为内置的兜底回答,内容为`"很抱歉,我无法回答您的问题"`,格式为openai格式的流式/非流式响应
|
||||
### 拒绝响应结构
|
||||
|
||||
如果用户使用了非openai格式的协议,此时对非法请求的处理逻辑为:
|
||||
- 如果配置了 `denyMessage`,返回用户配置的 `denyMessage` 内容,非流式响应
|
||||
- 如果没有配置 `denyMessage`,优先返回阿里云内容安全的建议回答,非流式响应
|
||||
- 如果阿里云内容安全未返回建议回答,返回内置的兜底回答,内容为`"很抱歉,我无法回答您的问题"`,非流式响应
|
||||
内容被拦截时,插件(`MultiModalGuard` action)统一返回以下结构化 JSON 对象,各协议的承载位置如下:
|
||||
|
||||
```json
|
||||
{
|
||||
"blockedDetails": [
|
||||
{
|
||||
"Type": "contentModeration",
|
||||
"Level": "high",
|
||||
"Suggestion": "block"
|
||||
}
|
||||
],
|
||||
"requestId": "AAAAAA-BBBB-CCCC-DDDD-EEEEEEE****",
|
||||
"guardCode": 200
|
||||
}
|
||||
```
|
||||
|
||||
字段说明:
|
||||
|
||||
| 字段 | 类型 | 说明 |
|
||||
| --- | --- | --- |
|
||||
| `blockedDetails` | array | 命中拦截的维度明细;若安全服务未返回明细,则根据顶层风险信号自动合成 |
|
||||
| `blockedDetails[].Type` | string | 风险类型:`contentModeration` / `promptAttack` / `sensitiveData` / `maliciousUrl` / `modelHallucination` |
|
||||
| `blockedDetails[].Level` | string | 风险等级:`high` / `medium` / `low` 等 |
|
||||
| `blockedDetails[].Suggestion` | string | 安全服务建议操作,通常为 `block` |
|
||||
| `requestId` | string | 安全服务的请求 ID,用于追踪 |
|
||||
| `guardCode` | int | 安全服务返回的业务码(非 HTTP 状态码,成功检测时为 `200`) |
|
||||
|
||||
各协议承载位置:
|
||||
|
||||
- **`text_generation`(OpenAI 非流式)**:上述结构体序列化为 JSON 字符串后放入 `choices[0].message.content`
|
||||
- **`text_generation`(OpenAI 流式 SSE)**:同上,放入首个 chunk 的 `delta.content`
|
||||
- **`text_generation`(`protocol=original`)**:上述结构体直接作为 JSON 响应 body 返回
|
||||
- **`image_generation`**:上述结构体直接作为 JSON 响应 body 返回(HTTP 403)
|
||||
- **`mcp`(JSON-RPC)**:上述结构体序列化为 JSON 字符串后放入 `error.message`
|
||||
- **`mcp`(SSE)**:同上,通过 SSE 事件返回
|
||||
|
||||
补充说明一下内容合规检测、提示词攻击检测、敏感内容检测三种风险的四个等级:
|
||||
|
||||
|
||||
@@ -41,6 +41,43 @@ Plugin Priority: `300`
|
||||
| `consumerResponseCheckService` | map | optional | - | Specify specific response detection services for different consumers |
|
||||
| `consumerRiskLevel` | map | optional | - | Specify interception risk levels for different consumers in different dimensions |
|
||||
|
||||
### Deny Response Body
|
||||
|
||||
When content is blocked, the plugin (`MultiModalGuard` action) returns the following structured JSON object. The location in the response depends on the protocol:
|
||||
|
||||
```json
|
||||
{
|
||||
"blockedDetails": [
|
||||
{
|
||||
"Type": "contentModeration",
|
||||
"Level": "high",
|
||||
"Suggestion": "block"
|
||||
}
|
||||
],
|
||||
"requestId": "AAAAAA-BBBB-CCCC-DDDD-EEEEEEE****",
|
||||
"guardCode": 200
|
||||
}
|
||||
```
|
||||
|
||||
Field descriptions:
|
||||
|
||||
| Field | Type | Description |
|
||||
| --- | --- | --- |
|
||||
| `blockedDetails` | array | Details of the triggered blocking dimensions. Synthesised from top-level risk signals when the security service returns no detail entries. |
|
||||
| `blockedDetails[].Type` | string | Risk type: `contentModeration` / `promptAttack` / `sensitiveData` / `maliciousUrl` / `modelHallucination` |
|
||||
| `blockedDetails[].Level` | string | Risk level: `high` / `medium` / `low` etc. |
|
||||
| `blockedDetails[].Suggestion` | string | Action recommended by the security service, usually `block` |
|
||||
| `requestId` | string | Request ID from the security service, for tracing |
|
||||
| `guardCode` | int | Business code returned by the security service (not an HTTP status code; `200` indicates a successful check that detected a risk) |
|
||||
|
||||
How the body is embedded per protocol:
|
||||
|
||||
- **`text_generation` (OpenAI non-streaming)**: serialised as a JSON string and placed in `choices[0].message.content`
|
||||
- **`text_generation` (OpenAI streaming SSE)**: same, placed in `delta.content` of the first chunk
|
||||
- **`text_generation` (`protocol=original`)**: returned directly as the JSON response body
|
||||
- **`image_generation`**: returned directly as the JSON response body (HTTP 403)
|
||||
- **`mcp` (JSON-RPC)**: serialised as a JSON string and placed in `error.message`
|
||||
- **`mcp` (SSE)**: same, returned via SSE event
|
||||
|
||||
## Examples of configuration
|
||||
### Check if the input is legal
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
package config
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"regexp"
|
||||
@@ -584,3 +585,68 @@ func IsRiskLevelAcceptable(action string, data Data, config AISecurityConfig, co
|
||||
return LevelToInt(data.RiskLevel) < LevelToInt(config.GetRiskLevelBar(consumer))
|
||||
}
|
||||
}
|
||||
|
||||
type DenyResponseBody struct {
|
||||
BlockedDetails []Detail `json:"blockedDetails"`
|
||||
RequestId string `json:"requestId"`
|
||||
// GuardCode is the business code returned by the security service (typically 200 when the check
|
||||
// succeeded and a risk was detected). It is NOT an HTTP status code.
|
||||
GuardCode int `json:"guardCode"`
|
||||
}
|
||||
|
||||
func BuildDenyResponseBody(response Response, config AISecurityConfig, consumer string) ([]byte, error) {
|
||||
body := DenyResponseBody{
|
||||
BlockedDetails: GetUnacceptableDetail(response.Data, config, consumer),
|
||||
RequestId: response.RequestId,
|
||||
GuardCode: response.Code,
|
||||
}
|
||||
return json.Marshal(body)
|
||||
}
|
||||
|
||||
func GetUnacceptableDetail(data Data, config AISecurityConfig, consumer string) []Detail {
|
||||
result := []Detail{}
|
||||
for _, detail := range data.Detail {
|
||||
switch detail.Type {
|
||||
case ContentModerationType:
|
||||
if LevelToInt(detail.Level) >= LevelToInt(config.GetContentModerationLevelBar(consumer)) {
|
||||
result = append(result, detail)
|
||||
}
|
||||
case PromptAttackType:
|
||||
if LevelToInt(detail.Level) >= LevelToInt(config.GetPromptAttackLevelBar(consumer)) {
|
||||
result = append(result, detail)
|
||||
}
|
||||
case SensitiveDataType:
|
||||
if LevelToInt(detail.Level) >= LevelToInt(config.GetSensitiveDataLevelBar(consumer)) {
|
||||
result = append(result, detail)
|
||||
}
|
||||
case MaliciousUrlDataType:
|
||||
if LevelToInt(detail.Level) >= LevelToInt(config.GetMaliciousUrlLevelBar(consumer)) {
|
||||
result = append(result, detail)
|
||||
}
|
||||
case ModelHallucinationDataType:
|
||||
if LevelToInt(detail.Level) >= LevelToInt(config.GetModelHallucinationLevelBar(consumer)) {
|
||||
result = append(result, detail)
|
||||
}
|
||||
}
|
||||
}
|
||||
// Fallback: when the security service returns a top-level risk signal but no Detail entries,
|
||||
// synthesise detail items from RiskLevel/AttackLevel so blockedDetails is never empty on a
|
||||
// real block event.
|
||||
if len(result) == 0 {
|
||||
if LevelToInt(data.RiskLevel) >= LevelToInt(config.GetContentModerationLevelBar(consumer)) {
|
||||
result = append(result, Detail{
|
||||
Type: ContentModerationType,
|
||||
Level: data.RiskLevel,
|
||||
Suggestion: "block",
|
||||
})
|
||||
}
|
||||
if LevelToInt(data.AttackLevel) >= LevelToInt(config.GetPromptAttackLevelBar(consumer)) {
|
||||
result = append(result, Detail{
|
||||
Type: PromptAttackType,
|
||||
Level: data.AttackLevel,
|
||||
Suggestion: "block",
|
||||
})
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
@@ -65,13 +65,19 @@ func HandleTextGenerationStreamingResponseBody(ctx wrapper.HttpContext, config c
|
||||
return
|
||||
}
|
||||
if !cfg.IsRiskLevelAcceptable(config.Action, response.Data, config, consumer) {
|
||||
denyMessage := cfg.DefaultDenyMessage
|
||||
if response.Data.Advice != nil && response.Data.Advice[0].Answer != "" {
|
||||
denyMessage = "\n" + response.Data.Advice[0].Answer
|
||||
} else if config.DenyMessage != "" {
|
||||
denyMessage = config.DenyMessage
|
||||
denyBody, err := cfg.BuildDenyResponseBody(response, config, consumer)
|
||||
if err != nil {
|
||||
log.Errorf("failed to build deny response body: %v", err)
|
||||
endStream := ctx.GetContext("end_of_stream_received").(bool) && ctx.BufferQueueSize() == 0
|
||||
proxywasm.InjectEncodedDataToFilterChain(bytes.Join(bufferQueue, []byte("")), endStream)
|
||||
bufferQueue = [][]byte{}
|
||||
if !endStream {
|
||||
ctx.SetContext("during_call", false)
|
||||
singleCall()
|
||||
}
|
||||
return
|
||||
}
|
||||
marshalledDenyMessage := wrapper.MarshalStr(denyMessage)
|
||||
marshalledDenyMessage := wrapper.MarshalStr(string(denyBody))
|
||||
randomID := utils.GenerateRandomChatID()
|
||||
jsonData := []byte(fmt.Sprintf(cfg.OpenAIStreamResponseFormat, randomID, marshalledDenyMessage, randomID))
|
||||
proxywasm.InjectEncodedDataToFilterChain(jsonData, true)
|
||||
@@ -199,21 +205,22 @@ func HandleTextGenerationResponseBody(ctx wrapper.HttpContext, config cfg.AISecu
|
||||
}
|
||||
return
|
||||
}
|
||||
denyMessage := cfg.DefaultDenyMessage
|
||||
if config.DenyMessage != "" {
|
||||
denyMessage = config.DenyMessage
|
||||
} else if response.Data.Advice != nil && response.Data.Advice[0].Answer != "" {
|
||||
denyMessage = response.Data.Advice[0].Answer
|
||||
denyBody, err := cfg.BuildDenyResponseBody(response, config, consumer)
|
||||
if err != nil {
|
||||
log.Errorf("failed to build deny response body: %v", err)
|
||||
proxywasm.ResumeHttpResponse()
|
||||
return
|
||||
}
|
||||
marshalledDenyMessage := wrapper.MarshalStr(denyMessage)
|
||||
if config.ProtocolOriginal {
|
||||
proxywasm.SendHttpResponse(uint32(config.DenyCode), [][2]string{{"content-type", "application/json"}}, []byte(marshalledDenyMessage), -1)
|
||||
proxywasm.SendHttpResponse(uint32(config.DenyCode), [][2]string{{"content-type", "application/json"}}, denyBody, -1)
|
||||
} else if isStreamingResponse {
|
||||
randomID := utils.GenerateRandomChatID()
|
||||
marshalledDenyMessage := wrapper.MarshalStr(string(denyBody))
|
||||
jsonData := []byte(fmt.Sprintf(cfg.OpenAIStreamResponseFormat, randomID, marshalledDenyMessage, randomID))
|
||||
proxywasm.SendHttpResponse(uint32(config.DenyCode), [][2]string{{"content-type", "text/event-stream;charset=UTF-8"}}, jsonData, -1)
|
||||
} else {
|
||||
randomID := utils.GenerateRandomChatID()
|
||||
marshalledDenyMessage := wrapper.MarshalStr(string(denyBody))
|
||||
jsonData := []byte(fmt.Sprintf(cfg.OpenAIResponseFormat, randomID, marshalledDenyMessage))
|
||||
proxywasm.SendHttpResponse(uint32(config.DenyCode), [][2]string{{"content-type", "application/json"}}, jsonData, -1)
|
||||
}
|
||||
|
||||
@@ -85,14 +85,13 @@ func HandleOpenAIImageGenerationRequestBody(ctx wrapper.HttpContext, config cfg.
|
||||
}
|
||||
return
|
||||
}
|
||||
denyMessage := cfg.DefaultDenyMessage
|
||||
if config.DenyMessage != "" {
|
||||
denyMessage = config.DenyMessage
|
||||
} else if response.Data.Advice != nil && response.Data.Advice[0].Answer != "" {
|
||||
denyMessage = response.Data.Advice[0].Answer
|
||||
denyBody, err := cfg.BuildDenyResponseBody(response, config, consumer)
|
||||
if err != nil {
|
||||
log.Errorf("failed to build deny response body: %v", err)
|
||||
proxywasm.ResumeHttpRequest()
|
||||
return
|
||||
}
|
||||
marshalledDenyMessage := wrapper.MarshalStr(denyMessage)
|
||||
proxywasm.SendHttpResponse(403, [][2]string{{"content-type", "application/json"}}, []byte(marshalledDenyMessage), -1)
|
||||
proxywasm.SendHttpResponse(403, [][2]string{{"content-type", "application/json"}}, denyBody, -1)
|
||||
ctx.DontReadResponseBody()
|
||||
config.IncrementCounter("ai_sec_request_deny", 1)
|
||||
endTime := time.Now().UnixMilli()
|
||||
@@ -157,14 +156,13 @@ func HandleOpenAIImageGenerationRequestBody(ctx wrapper.HttpContext, config cfg.
|
||||
return
|
||||
}
|
||||
|
||||
denyMessage := cfg.DefaultDenyMessage
|
||||
if config.DenyMessage != "" {
|
||||
denyMessage = config.DenyMessage
|
||||
} else if response.Data.Advice != nil && response.Data.Advice[0].Answer != "" {
|
||||
denyMessage = response.Data.Advice[0].Answer
|
||||
denyBody, err := cfg.BuildDenyResponseBody(response, config, consumer)
|
||||
if err != nil {
|
||||
log.Errorf("failed to build deny response body: %v", err)
|
||||
proxywasm.ResumeHttpRequest()
|
||||
return
|
||||
}
|
||||
marshalledDenyMessage := wrapper.MarshalStr(denyMessage)
|
||||
proxywasm.SendHttpResponse(403, [][2]string{{"content-type", "application/json"}}, []byte(marshalledDenyMessage), -1)
|
||||
proxywasm.SendHttpResponse(403, [][2]string{{"content-type", "application/json"}}, denyBody, -1)
|
||||
ctx.DontReadResponseBody()
|
||||
config.IncrementCounter("ai_sec_request_deny", 1)
|
||||
ctx.SetUserAttribute("safecheck_request_rt", endTime-startTime)
|
||||
@@ -244,7 +242,13 @@ func HandleOpenAIImageGenerationResponseBody(ctx wrapper.HttpContext, config cfg
|
||||
}
|
||||
return
|
||||
}
|
||||
proxywasm.SendHttpResponse(403, [][2]string{{"content-type", "application/json"}}, []byte("illegal image"), -1)
|
||||
denyBody, err := cfg.BuildDenyResponseBody(response, config, consumer)
|
||||
if err != nil {
|
||||
log.Errorf("failed to build deny response body: %v", err)
|
||||
proxywasm.ResumeHttpResponse()
|
||||
return
|
||||
}
|
||||
proxywasm.SendHttpResponse(403, [][2]string{{"content-type", "application/json"}}, denyBody, -1)
|
||||
config.IncrementCounter("ai_sec_request_deny", 1)
|
||||
ctx.SetUserAttribute("safecheck_request_rt", endTime-startTime)
|
||||
ctx.SetUserAttribute("safecheck_status", "reqeust deny")
|
||||
|
||||
@@ -243,14 +243,13 @@ func HandleQwenImageGenerationRequestBody(ctx wrapper.HttpContext, config cfg.AI
|
||||
}
|
||||
return
|
||||
}
|
||||
denyMessage := cfg.DefaultDenyMessage
|
||||
if config.DenyMessage != "" {
|
||||
denyMessage = config.DenyMessage
|
||||
} else if response.Data.Advice != nil && response.Data.Advice[0].Answer != "" {
|
||||
denyMessage = response.Data.Advice[0].Answer
|
||||
denyBody, err := cfg.BuildDenyResponseBody(response, config, consumer)
|
||||
if err != nil {
|
||||
log.Errorf("failed to build deny response body: %v", err)
|
||||
proxywasm.ResumeHttpRequest()
|
||||
return
|
||||
}
|
||||
marshalledDenyMessage := wrapper.MarshalStr(denyMessage)
|
||||
proxywasm.SendHttpResponse(403, [][2]string{{"content-type", "application/json"}}, []byte(marshalledDenyMessage), -1)
|
||||
proxywasm.SendHttpResponse(403, [][2]string{{"content-type", "application/json"}}, denyBody, -1)
|
||||
ctx.DontReadResponseBody()
|
||||
config.IncrementCounter("ai_sec_request_deny", 1)
|
||||
endTime := time.Now().UnixMilli()
|
||||
@@ -315,14 +314,13 @@ func HandleQwenImageGenerationRequestBody(ctx wrapper.HttpContext, config cfg.AI
|
||||
return
|
||||
}
|
||||
|
||||
denyMessage := cfg.DefaultDenyMessage
|
||||
if config.DenyMessage != "" {
|
||||
denyMessage = config.DenyMessage
|
||||
} else if response.Data.Advice != nil && response.Data.Advice[0].Answer != "" {
|
||||
denyMessage = response.Data.Advice[0].Answer
|
||||
denyBody, err := cfg.BuildDenyResponseBody(response, config, consumer)
|
||||
if err != nil {
|
||||
log.Errorf("failed to build deny response body: %v", err)
|
||||
proxywasm.ResumeHttpRequest()
|
||||
return
|
||||
}
|
||||
marshalledDenyMessage := wrapper.MarshalStr(denyMessage)
|
||||
proxywasm.SendHttpResponse(403, [][2]string{{"content-type", "application/json"}}, []byte(marshalledDenyMessage), -1)
|
||||
proxywasm.SendHttpResponse(403, [][2]string{{"content-type", "application/json"}}, denyBody, -1)
|
||||
ctx.DontReadResponseBody()
|
||||
config.IncrementCounter("ai_sec_request_deny", 1)
|
||||
ctx.SetUserAttribute("safecheck_request_rt", endTime-startTime)
|
||||
@@ -402,14 +400,13 @@ func HandleQwenImageGenerationResponseBody(ctx wrapper.HttpContext, config cfg.A
|
||||
}
|
||||
return
|
||||
}
|
||||
denyMessage := cfg.DefaultDenyMessage
|
||||
if config.DenyMessage != "" {
|
||||
denyMessage = config.DenyMessage
|
||||
} else if response.Data.Advice != nil && response.Data.Advice[0].Answer != "" {
|
||||
denyMessage = response.Data.Advice[0].Answer
|
||||
denyBody, err := cfg.BuildDenyResponseBody(response, config, consumer)
|
||||
if err != nil {
|
||||
log.Errorf("failed to build deny response body: %v", err)
|
||||
proxywasm.ResumeHttpResponse()
|
||||
return
|
||||
}
|
||||
marshalledDenyMessage := wrapper.MarshalStr(denyMessage)
|
||||
proxywasm.SendHttpResponse(403, [][2]string{{"content-type", "application/json"}}, []byte(marshalledDenyMessage), -1)
|
||||
proxywasm.SendHttpResponse(403, [][2]string{{"content-type", "application/json"}}, denyBody, -1)
|
||||
config.IncrementCounter("ai_sec_request_deny", 1)
|
||||
ctx.SetUserAttribute("safecheck_request_rt", endTime-startTime)
|
||||
ctx.SetUserAttribute("safecheck_status", "reqeust deny")
|
||||
|
||||
@@ -2,6 +2,7 @@ package mcp
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
@@ -18,9 +19,9 @@ import (
|
||||
|
||||
const (
|
||||
MethodToolCall = "tools/call"
|
||||
DenyResponse = `{"jsonrpc":"2.0","id":0,"error":{"code":403,"message":"blocked by security guard"}}`
|
||||
DenyResponse = `{"jsonrpc":"2.0","id":0,"error":{"code":403,"message":"%s"}}`
|
||||
DenySSEResponse = `event: message
|
||||
data: {"jsonrpc":"2.0","id":0,"error":{"code":403,"message":"blocked by security guard"}}
|
||||
data: {"jsonrpc":"2.0","id":0,"error":{"code":403,"message":"%s"}}
|
||||
|
||||
`
|
||||
)
|
||||
@@ -78,7 +79,15 @@ func HandleMcpRequestBody(ctx wrapper.HttpContext, config cfg.AISecurityConfig,
|
||||
ctx.SetUserAttribute("safecheck_riskWords", response.Data.Result[0].RiskWords)
|
||||
}
|
||||
ctx.WriteUserAttributeToLogWithKey(wrapper.AILogKey)
|
||||
proxywasm.SendHttpResponse(uint32(config.DenyCode), [][2]string{{"content-type", "application/json"}}, []byte(DenyResponse), -1)
|
||||
denyBody, err := cfg.BuildDenyResponseBody(response, config, consumer)
|
||||
if err != nil {
|
||||
log.Errorf("failed to build deny response body: %v", err)
|
||||
proxywasm.ResumeHttpRequest()
|
||||
return
|
||||
}
|
||||
marshalledDenyMessage := wrapper.MarshalStr(string(denyBody))
|
||||
denyResponse := fmt.Sprintf(DenyResponse, marshalledDenyMessage)
|
||||
proxywasm.SendHttpResponse(uint32(config.DenyCode), [][2]string{{"content-type", "application/json"}}, []byte(denyResponse), -1)
|
||||
}
|
||||
singleCall = func() {
|
||||
var nextContentIndex int
|
||||
@@ -124,7 +133,15 @@ func HandleMcpStreamingResponseBody(ctx wrapper.HttpContext, config cfg.AISecuri
|
||||
return
|
||||
}
|
||||
if !cfg.IsRiskLevelAcceptable(config.Action, response.Data, config, consumer) {
|
||||
proxywasm.InjectEncodedDataToFilterChain([]byte(DenySSEResponse), true)
|
||||
denyBody, err := cfg.BuildDenyResponseBody(response, config, consumer)
|
||||
if err != nil {
|
||||
log.Errorf("failed to build deny response body: %v", err)
|
||||
proxywasm.InjectEncodedDataToFilterChain(frontBuffer, false)
|
||||
return
|
||||
}
|
||||
marshalledDenyMessage := wrapper.MarshalStr(string(denyBody))
|
||||
denySSEResponse := fmt.Sprintf(DenySSEResponse, marshalledDenyMessage)
|
||||
proxywasm.InjectEncodedDataToFilterChain([]byte(denySSEResponse), true)
|
||||
} else {
|
||||
proxywasm.InjectEncodedDataToFilterChain(frontBuffer, false)
|
||||
}
|
||||
@@ -212,8 +229,16 @@ func HandleMcpResponseBody(ctx wrapper.HttpContext, config cfg.AISecurityConfig,
|
||||
ctx.SetUserAttribute("safecheck_riskWords", response.Data.Result[0].RiskWords)
|
||||
}
|
||||
ctx.WriteUserAttributeToLogWithKey(wrapper.AILogKey)
|
||||
denyBody, err := cfg.BuildDenyResponseBody(response, config, consumer)
|
||||
if err != nil {
|
||||
log.Errorf("failed to build deny response body: %v", err)
|
||||
proxywasm.ResumeHttpResponse()
|
||||
return
|
||||
}
|
||||
marshalledDenyMessage := wrapper.MarshalStr(string(denyBody))
|
||||
denyResponseBody := fmt.Sprintf(DenyResponse, marshalledDenyMessage)
|
||||
proxywasm.RemoveHttpResponseHeader("content-length")
|
||||
proxywasm.ReplaceHttpResponseBody([]byte(DenyResponse))
|
||||
proxywasm.ReplaceHttpResponseBody([]byte(denyResponseBody))
|
||||
proxywasm.ResumeHttpResponse()
|
||||
// proxywasm.SendHttpResponse(uint32(config.DenyCode), [][2]string{{"content-type", "application/json"}}, []byte(DenyResponse), -1)
|
||||
}
|
||||
|
||||
@@ -96,21 +96,22 @@ func HandleTextGenerationRequestBody(ctx wrapper.HttpContext, config cfg.AISecur
|
||||
}
|
||||
return
|
||||
}
|
||||
denyMessage := cfg.DefaultDenyMessage
|
||||
if config.DenyMessage != "" {
|
||||
denyMessage = config.DenyMessage
|
||||
} else if response.Data.Advice != nil && response.Data.Advice[0].Answer != "" {
|
||||
denyMessage = response.Data.Advice[0].Answer
|
||||
denyBody, err := cfg.BuildDenyResponseBody(response, config, consumer)
|
||||
if err != nil {
|
||||
log.Errorf("failed to build deny response body: %v", err)
|
||||
proxywasm.ResumeHttpRequest()
|
||||
return
|
||||
}
|
||||
marshalledDenyMessage := wrapper.MarshalStr(denyMessage)
|
||||
if config.ProtocolOriginal {
|
||||
proxywasm.SendHttpResponse(uint32(config.DenyCode), [][2]string{{"content-type", "application/json"}}, []byte(marshalledDenyMessage), -1)
|
||||
proxywasm.SendHttpResponse(uint32(config.DenyCode), [][2]string{{"content-type", "application/json"}}, denyBody, -1)
|
||||
} else if gjson.GetBytes(body, "stream").Bool() {
|
||||
randomID := utils.GenerateRandomChatID()
|
||||
marshalledDenyMessage := wrapper.MarshalStr(string(denyBody))
|
||||
jsonData := []byte(fmt.Sprintf(cfg.OpenAIStreamResponseFormat, randomID, marshalledDenyMessage, randomID))
|
||||
proxywasm.SendHttpResponse(uint32(config.DenyCode), [][2]string{{"content-type", "text/event-stream;charset=UTF-8"}}, jsonData, -1)
|
||||
} else {
|
||||
randomID := utils.GenerateRandomChatID()
|
||||
marshalledDenyMessage := wrapper.MarshalStr(string(denyBody))
|
||||
jsonData := []byte(fmt.Sprintf(cfg.OpenAIResponseFormat, randomID, marshalledDenyMessage))
|
||||
proxywasm.SendHttpResponse(uint32(config.DenyCode), [][2]string{{"content-type", "application/json"}}, jsonData, -1)
|
||||
}
|
||||
@@ -178,21 +179,22 @@ func HandleTextGenerationRequestBody(ctx wrapper.HttpContext, config cfg.AISecur
|
||||
return
|
||||
}
|
||||
|
||||
denyMessage := cfg.DefaultDenyMessage
|
||||
if config.DenyMessage != "" {
|
||||
denyMessage = config.DenyMessage
|
||||
} else if response.Data.Advice != nil && response.Data.Advice[0].Answer != "" {
|
||||
denyMessage = response.Data.Advice[0].Answer
|
||||
denyBody, err := cfg.BuildDenyResponseBody(response, config, consumer)
|
||||
if err != nil {
|
||||
log.Errorf("failed to build deny response body: %v", err)
|
||||
proxywasm.ResumeHttpRequest()
|
||||
return
|
||||
}
|
||||
marshalledDenyMessage := wrapper.MarshalStr(denyMessage)
|
||||
if config.ProtocolOriginal {
|
||||
proxywasm.SendHttpResponse(uint32(config.DenyCode), [][2]string{{"content-type", "application/json"}}, []byte(marshalledDenyMessage), -1)
|
||||
proxywasm.SendHttpResponse(uint32(config.DenyCode), [][2]string{{"content-type", "application/json"}}, denyBody, -1)
|
||||
} else if gjson.GetBytes(body, "stream").Bool() {
|
||||
randomID := utils.GenerateRandomChatID()
|
||||
marshalledDenyMessage := wrapper.MarshalStr(string(denyBody))
|
||||
jsonData := []byte(fmt.Sprintf(cfg.OpenAIStreamResponseFormat, randomID, marshalledDenyMessage, randomID))
|
||||
proxywasm.SendHttpResponse(uint32(config.DenyCode), [][2]string{{"content-type", "text/event-stream;charset=UTF-8"}}, jsonData, -1)
|
||||
} else {
|
||||
randomID := utils.GenerateRandomChatID()
|
||||
marshalledDenyMessage := wrapper.MarshalStr(string(denyBody))
|
||||
jsonData := []byte(fmt.Sprintf(cfg.OpenAIResponseFormat, randomID, marshalledDenyMessage))
|
||||
proxywasm.SendHttpResponse(uint32(config.DenyCode), [][2]string{{"content-type", "application/json"}}, jsonData, -1)
|
||||
}
|
||||
|
||||
@@ -53,21 +53,22 @@ func HandleTextGenerationRequestBody(ctx wrapper.HttpContext, config cfg.AISecur
|
||||
}
|
||||
return
|
||||
}
|
||||
denyMessage := cfg.DefaultDenyMessage
|
||||
if config.DenyMessage != "" {
|
||||
denyMessage = config.DenyMessage
|
||||
} else if response.Data.Advice != nil && response.Data.Advice[0].Answer != "" {
|
||||
denyMessage = response.Data.Advice[0].Answer
|
||||
denyBody, err := cfg.BuildDenyResponseBody(response, config, consumer)
|
||||
if err != nil {
|
||||
log.Errorf("failed to build deny response body: %v", err)
|
||||
proxywasm.ResumeHttpRequest()
|
||||
return
|
||||
}
|
||||
marshalledDenyMessage := wrapper.MarshalStr(denyMessage)
|
||||
if config.ProtocolOriginal {
|
||||
proxywasm.SendHttpResponse(uint32(config.DenyCode), [][2]string{{"content-type", "application/json"}}, []byte(marshalledDenyMessage), -1)
|
||||
proxywasm.SendHttpResponse(uint32(config.DenyCode), [][2]string{{"content-type", "application/json"}}, denyBody, -1)
|
||||
} else if gjson.GetBytes(body, "stream").Bool() {
|
||||
randomID := utils.GenerateRandomChatID()
|
||||
marshalledDenyMessage := wrapper.MarshalStr(string(denyBody))
|
||||
jsonData := []byte(fmt.Sprintf(cfg.OpenAIStreamResponseFormat, randomID, marshalledDenyMessage, randomID))
|
||||
proxywasm.SendHttpResponse(uint32(config.DenyCode), [][2]string{{"content-type", "text/event-stream;charset=UTF-8"}}, jsonData, -1)
|
||||
} else {
|
||||
randomID := utils.GenerateRandomChatID()
|
||||
marshalledDenyMessage := wrapper.MarshalStr(string(denyBody))
|
||||
jsonData := []byte(fmt.Sprintf(cfg.OpenAIResponseFormat, randomID, marshalledDenyMessage))
|
||||
proxywasm.SendHttpResponse(uint32(config.DenyCode), [][2]string{{"content-type", "application/json"}}, jsonData, -1)
|
||||
}
|
||||
|
||||
@@ -156,6 +156,90 @@ var mcpConfig = func() json.RawMessage {
|
||||
return data
|
||||
}()
|
||||
|
||||
// 测试配置:MCP配置(启用请求检查)
|
||||
var mcpRequestConfig = func() json.RawMessage {
|
||||
data, _ := json.Marshal(map[string]interface{}{
|
||||
"serviceName": "security-service",
|
||||
"servicePort": 8080,
|
||||
"serviceHost": "security.example.com",
|
||||
"accessKey": "test-ak",
|
||||
"secretKey": "test-sk",
|
||||
"checkRequest": true,
|
||||
"checkResponse": false,
|
||||
"action": "MultiModalGuard",
|
||||
"apiType": "mcp",
|
||||
"requestContentJsonPath": "params.arguments",
|
||||
"contentModerationLevelBar": "high",
|
||||
"promptAttackLevelBar": "high",
|
||||
"sensitiveDataLevelBar": "S3",
|
||||
"timeout": 2000,
|
||||
})
|
||||
return data
|
||||
}()
|
||||
|
||||
// 测试配置:MultiModalGuard 文本生成
|
||||
var multiModalGuardTextConfig = func() json.RawMessage {
|
||||
data, _ := json.Marshal(map[string]interface{}{
|
||||
"serviceName": "security-service",
|
||||
"servicePort": 8080,
|
||||
"serviceHost": "security.example.com",
|
||||
"accessKey": "test-ak",
|
||||
"secretKey": "test-sk",
|
||||
"checkRequest": true,
|
||||
"checkResponse": true,
|
||||
"action": "MultiModalGuard",
|
||||
"apiType": "text_generation",
|
||||
"contentModerationLevelBar": "high",
|
||||
"promptAttackLevelBar": "high",
|
||||
"sensitiveDataLevelBar": "S3",
|
||||
"timeout": 2000,
|
||||
"bufferLimit": 1000,
|
||||
})
|
||||
return data
|
||||
}()
|
||||
|
||||
// 测试配置:MultiModalGuard OpenAI 图像生成
|
||||
var multiModalGuardImageOpenAIConfig = func() json.RawMessage {
|
||||
data, _ := json.Marshal(map[string]interface{}{
|
||||
"serviceName": "security-service",
|
||||
"servicePort": 8080,
|
||||
"serviceHost": "security.example.com",
|
||||
"accessKey": "test-ak",
|
||||
"secretKey": "test-sk",
|
||||
"checkRequest": true,
|
||||
"checkResponse": true,
|
||||
"action": "MultiModalGuard",
|
||||
"apiType": "image_generation",
|
||||
"providerType": "openai",
|
||||
"contentModerationLevelBar": "high",
|
||||
"promptAttackLevelBar": "high",
|
||||
"sensitiveDataLevelBar": "S3",
|
||||
"timeout": 2000,
|
||||
})
|
||||
return data
|
||||
}()
|
||||
|
||||
// 测试配置:MultiModalGuard Qwen 图像生成
|
||||
var multiModalGuardImageQwenConfig = func() json.RawMessage {
|
||||
data, _ := json.Marshal(map[string]interface{}{
|
||||
"serviceName": "security-service",
|
||||
"servicePort": 8080,
|
||||
"serviceHost": "security.example.com",
|
||||
"accessKey": "test-ak",
|
||||
"secretKey": "test-sk",
|
||||
"checkRequest": true,
|
||||
"checkResponse": true,
|
||||
"action": "MultiModalGuard",
|
||||
"apiType": "image_generation",
|
||||
"providerType": "qwen",
|
||||
"contentModerationLevelBar": "high",
|
||||
"promptAttackLevelBar": "high",
|
||||
"sensitiveDataLevelBar": "S3",
|
||||
"timeout": 2000,
|
||||
})
|
||||
return data
|
||||
}()
|
||||
|
||||
func TestParseConfig(t *testing.T) {
|
||||
test.RunGoTest(t, func(t *testing.T) {
|
||||
// 测试基础配置解析
|
||||
@@ -330,6 +414,51 @@ func TestOnHttpRequestBody(t *testing.T) {
|
||||
// 空内容应该直接通过
|
||||
require.Equal(t, types.ActionContinue, action)
|
||||
})
|
||||
|
||||
// TextModerationPlus(默认 action,含 agent/OpenAI 形态)请求拦截应返回 choices[0].message.content 内的 blockedDetails JSON
|
||||
t.Run("text moderation plus request deny returns blockedDetails in openai completion shape", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(basicConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/v1/chat/completions"},
|
||||
{":method", "POST"},
|
||||
})
|
||||
|
||||
body := `{"messages": [{"role": "user", "content": "trigger deny"}]}`
|
||||
require.Equal(t, types.ActionPause, host.CallOnHttpRequestBody([]byte(body)))
|
||||
|
||||
securityResponse := `{"Code": 200, "Message": "Success", "RequestId": "req-tmp-deny", "Data": {"RiskLevel": "high"}}`
|
||||
host.CallOnHttpCall([][2]string{
|
||||
{":status", "200"},
|
||||
{"content-type", "application/json"},
|
||||
}, []byte(securityResponse))
|
||||
|
||||
local := host.GetLocalResponse()
|
||||
require.NotNil(t, local, "expected SendHttpResponse for request deny")
|
||||
require.Contains(t, string(local.Data), "blockedDetails")
|
||||
require.Contains(t, string(local.Data), "req-tmp-deny")
|
||||
|
||||
type openAIChatCompletion struct {
|
||||
Choices []struct {
|
||||
Message struct {
|
||||
Content string `json:"content"`
|
||||
} `json:"message"`
|
||||
} `json:"choices"`
|
||||
}
|
||||
var outer openAIChatCompletion
|
||||
require.NoError(t, json.Unmarshal(local.Data, &outer))
|
||||
require.Len(t, outer.Choices, 1)
|
||||
|
||||
var deny cfg.DenyResponseBody
|
||||
require.NoError(t, json.Unmarshal([]byte(outer.Choices[0].Message.Content), &deny))
|
||||
require.Equal(t, "req-tmp-deny", deny.RequestId)
|
||||
require.Equal(t, 200, deny.GuardCode)
|
||||
require.NotEmpty(t, deny.BlockedDetails)
|
||||
require.Equal(t, cfg.ContentModerationType, deny.BlockedDetails[0].Type)
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
@@ -649,3 +778,444 @@ func TestUtilityFunctions(t *testing.T) {
|
||||
require.Len(t, id, 38) // "chatcmpl-" + 29 random chars
|
||||
})
|
||||
}
|
||||
|
||||
func TestMultiModalGuardTextGenerationDeny(t *testing.T) {
|
||||
test.RunTest(t, func(t *testing.T) {
|
||||
// MultiModalGuard text_generation request deny → exercises multi_modal_guard/text/openai.go BuildDenyResponseBody path
|
||||
t.Run("multi modal guard text request deny returns blockedDetails", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(multiModalGuardTextConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/v1/chat/completions"},
|
||||
{":method", "POST"},
|
||||
})
|
||||
|
||||
body := `{"messages": [{"role": "user", "content": "trigger deny"}]}`
|
||||
require.Equal(t, types.ActionPause, host.CallOnHttpRequestBody([]byte(body)))
|
||||
|
||||
securityResponse := `{"Code": 200, "Message": "Success", "RequestId": "req-mmg-text-deny", "Data": {"RiskLevel": "high"}}`
|
||||
host.CallOnHttpCall([][2]string{
|
||||
{":status", "200"},
|
||||
{"content-type", "application/json"},
|
||||
}, []byte(securityResponse))
|
||||
|
||||
local := host.GetLocalResponse()
|
||||
require.NotNil(t, local, "expected SendHttpResponse for request deny")
|
||||
require.Contains(t, string(local.Data), "blockedDetails")
|
||||
require.Contains(t, string(local.Data), "req-mmg-text-deny")
|
||||
})
|
||||
|
||||
// MultiModalGuard text_generation response deny → exercises common/text/openai.go HandleTextGenerationResponseBody BuildDenyResponseBody path
|
||||
t.Run("multi modal guard text response deny returns blockedDetails", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(multiModalGuardTextConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/v1/chat/completions"},
|
||||
{":method", "POST"},
|
||||
})
|
||||
|
||||
host.CallOnHttpResponseHeaders([][2]string{
|
||||
{":status", "200"},
|
||||
{"content-type", "application/json"},
|
||||
})
|
||||
|
||||
body := `{"choices": [{"message": {"role": "assistant", "content": "bad response content"}}]}`
|
||||
action := host.CallOnHttpResponseBody([]byte(body))
|
||||
require.Equal(t, types.ActionPause, action)
|
||||
|
||||
securityResponse := `{"Code": 200, "Message": "Success", "RequestId": "req-mmg-resp-deny", "Data": {"RiskLevel": "high"}}`
|
||||
host.CallOnHttpCall([][2]string{
|
||||
{":status", "200"},
|
||||
{"content-type", "application/json"},
|
||||
}, []byte(securityResponse))
|
||||
|
||||
local := host.GetLocalResponse()
|
||||
require.NotNil(t, local, "expected SendHttpResponse for response deny")
|
||||
require.Contains(t, string(local.Data), "blockedDetails")
|
||||
require.Contains(t, string(local.Data), "req-mmg-resp-deny")
|
||||
})
|
||||
|
||||
// MultiModalGuard text_generation request pass
|
||||
t.Run("multi modal guard text request pass", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(multiModalGuardTextConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/v1/chat/completions"},
|
||||
{":method", "POST"},
|
||||
})
|
||||
|
||||
body := `{"messages": [{"role": "user", "content": "Hello"}]}`
|
||||
require.Equal(t, types.ActionPause, host.CallOnHttpRequestBody([]byte(body)))
|
||||
|
||||
securityResponse := `{"Code": 200, "Message": "Success", "RequestId": "req-mmg-pass", "Data": {"RiskLevel": "low"}}`
|
||||
host.CallOnHttpCall([][2]string{
|
||||
{":status", "200"},
|
||||
{"content-type", "application/json"},
|
||||
}, []byte(securityResponse))
|
||||
|
||||
action := host.GetHttpStreamAction()
|
||||
require.Equal(t, types.ActionContinue, action)
|
||||
host.CompleteHttp()
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
func TestMultiModalGuardImageGenerationDeny(t *testing.T) {
|
||||
test.RunTest(t, func(t *testing.T) {
|
||||
// OpenAI image generation request deny → exercises multi_modal_guard/image/openai.go BuildDenyResponseBody path
|
||||
t.Run("openai image request deny returns blockedDetails", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(multiModalGuardImageOpenAIConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/v1/images/generations"},
|
||||
{":method", "POST"},
|
||||
})
|
||||
|
||||
body := `{"prompt": "generate bad image"}`
|
||||
require.Equal(t, types.ActionPause, host.CallOnHttpRequestBody([]byte(body)))
|
||||
|
||||
securityResponse := `{"Code": 200, "Message": "Success", "RequestId": "req-img-openai-deny", "Data": {"RiskLevel": "high"}}`
|
||||
host.CallOnHttpCall([][2]string{
|
||||
{":status", "200"},
|
||||
{"content-type", "application/json"},
|
||||
}, []byte(securityResponse))
|
||||
|
||||
local := host.GetLocalResponse()
|
||||
require.NotNil(t, local, "expected SendHttpResponse for OpenAI image request deny")
|
||||
require.Contains(t, string(local.Data), "blockedDetails")
|
||||
require.Contains(t, string(local.Data), "req-img-openai-deny")
|
||||
})
|
||||
|
||||
// OpenAI image generation request pass
|
||||
t.Run("openai image request pass", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(multiModalGuardImageOpenAIConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/v1/images/generations"},
|
||||
{":method", "POST"},
|
||||
})
|
||||
|
||||
body := `{"prompt": "a cute cat"}`
|
||||
require.Equal(t, types.ActionPause, host.CallOnHttpRequestBody([]byte(body)))
|
||||
|
||||
securityResponse := `{"Code": 200, "Message": "Success", "RequestId": "req-img-pass", "Data": {"RiskLevel": "low"}}`
|
||||
host.CallOnHttpCall([][2]string{
|
||||
{":status", "200"},
|
||||
{"content-type", "application/json"},
|
||||
}, []byte(securityResponse))
|
||||
|
||||
action := host.GetHttpStreamAction()
|
||||
require.Equal(t, types.ActionContinue, action)
|
||||
host.CompleteHttp()
|
||||
})
|
||||
|
||||
// Qwen image generation request deny → exercises multi_modal_guard/image/qwen.go BuildDenyResponseBody path
|
||||
t.Run("qwen image request deny returns blockedDetails", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(multiModalGuardImageQwenConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/v1/images/generations"},
|
||||
{":method", "POST"},
|
||||
})
|
||||
|
||||
body := `{"input": {"prompt": "generate bad image"}}`
|
||||
require.Equal(t, types.ActionPause, host.CallOnHttpRequestBody([]byte(body)))
|
||||
|
||||
securityResponse := `{"Code": 200, "Message": "Success", "RequestId": "req-img-qwen-deny", "Data": {"RiskLevel": "high"}}`
|
||||
host.CallOnHttpCall([][2]string{
|
||||
{":status", "200"},
|
||||
{"content-type", "application/json"},
|
||||
}, []byte(securityResponse))
|
||||
|
||||
local := host.GetLocalResponse()
|
||||
require.NotNil(t, local, "expected SendHttpResponse for Qwen image request deny")
|
||||
require.Contains(t, string(local.Data), "blockedDetails")
|
||||
require.Contains(t, string(local.Data), "req-img-qwen-deny")
|
||||
})
|
||||
|
||||
// Qwen image generation request pass
|
||||
t.Run("qwen image request pass", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(multiModalGuardImageQwenConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/v1/images/generations"},
|
||||
{":method", "POST"},
|
||||
})
|
||||
|
||||
body := `{"input": {"prompt": "a cute cat"}}`
|
||||
require.Equal(t, types.ActionPause, host.CallOnHttpRequestBody([]byte(body)))
|
||||
|
||||
securityResponse := `{"Code": 200, "Message": "Success", "RequestId": "req-qwen-pass", "Data": {"RiskLevel": "low"}}`
|
||||
host.CallOnHttpCall([][2]string{
|
||||
{":status", "200"},
|
||||
{"content-type", "application/json"},
|
||||
}, []byte(securityResponse))
|
||||
|
||||
action := host.GetHttpStreamAction()
|
||||
require.Equal(t, types.ActionContinue, action)
|
||||
host.CompleteHttp()
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
func TestMCPRequestDeny(t *testing.T) {
|
||||
test.RunTest(t, func(t *testing.T) {
|
||||
// MCP request deny → exercises multi_modal_guard/mcp/mcp.go HandleMcpRequestBody BuildDenyResponseBody path
|
||||
t.Run("mcp request deny returns blockedDetails", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(mcpRequestConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/mcp/call"},
|
||||
{":method", "POST"},
|
||||
})
|
||||
|
||||
body := `{"method": "tools/call", "params": {"arguments": "bad request content"}}`
|
||||
require.Equal(t, types.ActionPause, host.CallOnHttpRequestBody([]byte(body)))
|
||||
|
||||
securityResponse := `{"Code": 200, "Message": "Success", "RequestId": "req-mcp-deny", "Data": {"RiskLevel": "high"}}`
|
||||
host.CallOnHttpCall([][2]string{
|
||||
{":status", "200"},
|
||||
{"content-type", "application/json"},
|
||||
}, []byte(securityResponse))
|
||||
|
||||
local := host.GetLocalResponse()
|
||||
require.NotNil(t, local, "expected SendHttpResponse for MCP request deny")
|
||||
require.Contains(t, string(local.Data), "blockedDetails")
|
||||
require.Contains(t, string(local.Data), "req-mcp-deny")
|
||||
})
|
||||
|
||||
// MCP request pass
|
||||
t.Run("mcp request pass", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(mcpRequestConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/mcp/call"},
|
||||
{":method", "POST"},
|
||||
})
|
||||
|
||||
body := `{"method": "tools/call", "params": {"arguments": "safe content"}}`
|
||||
require.Equal(t, types.ActionPause, host.CallOnHttpRequestBody([]byte(body)))
|
||||
|
||||
securityResponse := `{"Code": 200, "Message": "Success", "RequestId": "req-mcp-pass", "Data": {"RiskLevel": "low"}}`
|
||||
host.CallOnHttpCall([][2]string{
|
||||
{":status", "200"},
|
||||
{"content-type", "application/json"},
|
||||
}, []byte(securityResponse))
|
||||
|
||||
action := host.GetHttpStreamAction()
|
||||
require.Equal(t, types.ActionContinue, action)
|
||||
host.CompleteHttp()
|
||||
})
|
||||
|
||||
// MCP request skip non-tool-call method
|
||||
t.Run("mcp request skip non-tool-call", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(mcpRequestConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/mcp/call"},
|
||||
{":method", "POST"},
|
||||
})
|
||||
|
||||
body := `{"method": "resources/list", "params": {}}`
|
||||
action := host.CallOnHttpRequestBody([]byte(body))
|
||||
require.Equal(t, types.ActionContinue, action)
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
func TestTextModerationPlusResponseDeny(t *testing.T) {
|
||||
test.RunTest(t, func(t *testing.T) {
|
||||
// TextModerationPlus response deny → exercises text_moderation_plus/text (via common/text) BuildDenyResponseBody response path
|
||||
t.Run("text moderation plus response deny returns blockedDetails", func(t *testing.T) {
|
||||
host, status := test.NewTestHost(basicConfig)
|
||||
defer host.Reset()
|
||||
require.Equal(t, types.OnPluginStartStatusOK, status)
|
||||
|
||||
host.CallOnHttpRequestHeaders([][2]string{
|
||||
{":authority", "example.com"},
|
||||
{":path", "/v1/chat/completions"},
|
||||
{":method", "POST"},
|
||||
})
|
||||
|
||||
host.CallOnHttpResponseHeaders([][2]string{
|
||||
{":status", "200"},
|
||||
{"content-type", "application/json"},
|
||||
})
|
||||
|
||||
body := `{"choices": [{"message": {"role": "assistant", "content": "bad response"}}]}`
|
||||
action := host.CallOnHttpResponseBody([]byte(body))
|
||||
require.Equal(t, types.ActionPause, action)
|
||||
|
||||
securityResponse := `{"Code": 200, "Message": "Success", "RequestId": "req-tmp-resp-deny", "Data": {"RiskLevel": "high"}}`
|
||||
host.CallOnHttpCall([][2]string{
|
||||
{":status", "200"},
|
||||
{"content-type", "application/json"},
|
||||
}, []byte(securityResponse))
|
||||
|
||||
local := host.GetLocalResponse()
|
||||
require.NotNil(t, local, "expected SendHttpResponse for response deny")
|
||||
require.Contains(t, string(local.Data), "blockedDetails")
|
||||
require.Contains(t, string(local.Data), "req-tmp-resp-deny")
|
||||
|
||||
// Verify OpenAI completion shape wrapper
|
||||
type openAIChatCompletion struct {
|
||||
Choices []struct {
|
||||
Message struct {
|
||||
Content string `json:"content"`
|
||||
} `json:"message"`
|
||||
} `json:"choices"`
|
||||
}
|
||||
var outer openAIChatCompletion
|
||||
require.NoError(t, json.Unmarshal(local.Data, &outer))
|
||||
require.Len(t, outer.Choices, 1)
|
||||
|
||||
var deny cfg.DenyResponseBody
|
||||
require.NoError(t, json.Unmarshal([]byte(outer.Choices[0].Message.Content), &deny))
|
||||
require.Equal(t, "req-tmp-resp-deny", deny.RequestId)
|
||||
require.Equal(t, 200, deny.GuardCode)
|
||||
require.NotEmpty(t, deny.BlockedDetails)
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
func TestBuildDenyResponseBody(t *testing.T) {
|
||||
makeConfig := func(contentBar, promptBar string) cfg.AISecurityConfig {
|
||||
return cfg.AISecurityConfig{
|
||||
ContentModerationLevelBar: contentBar,
|
||||
PromptAttackLevelBar: promptBar,
|
||||
SensitiveDataLevelBar: "S4",
|
||||
MaliciousUrlLevelBar: "max",
|
||||
ModelHallucinationLevelBar: "max",
|
||||
Action: cfg.MultiModalGuard,
|
||||
}
|
||||
}
|
||||
|
||||
t.Run("guardCode equals response.Code", func(t *testing.T) {
|
||||
resp := cfg.Response{
|
||||
Code: 200,
|
||||
RequestId: "req-123",
|
||||
Data: cfg.Data{},
|
||||
}
|
||||
body, err := cfg.BuildDenyResponseBody(resp, makeConfig("high", "high"), "")
|
||||
require.NoError(t, err)
|
||||
|
||||
var result cfg.DenyResponseBody
|
||||
require.NoError(t, json.Unmarshal(body, &result))
|
||||
require.Equal(t, 200, result.GuardCode)
|
||||
require.Equal(t, "req-123", result.RequestId)
|
||||
})
|
||||
|
||||
t.Run("blockedDetails from Data.Detail", func(t *testing.T) {
|
||||
resp := cfg.Response{
|
||||
Code: 200,
|
||||
RequestId: "req-456",
|
||||
Data: cfg.Data{
|
||||
Detail: []cfg.Detail{
|
||||
{Type: cfg.ContentModerationType, Level: "high", Suggestion: "block"},
|
||||
{Type: cfg.PromptAttackType, Level: "low", Suggestion: "block"},
|
||||
},
|
||||
},
|
||||
}
|
||||
config := makeConfig("high", "high")
|
||||
body, err := cfg.BuildDenyResponseBody(resp, config, "")
|
||||
require.NoError(t, err)
|
||||
|
||||
var result cfg.DenyResponseBody
|
||||
require.NoError(t, json.Unmarshal(body, &result))
|
||||
// only the contentModeration entry meets the "high" bar; promptAttack at "low" does not
|
||||
require.Len(t, result.BlockedDetails, 1)
|
||||
require.Equal(t, cfg.ContentModerationType, result.BlockedDetails[0].Type)
|
||||
require.Equal(t, "high", result.BlockedDetails[0].Level)
|
||||
})
|
||||
|
||||
t.Run("blockedDetails fallback from RiskLevel when Detail is empty", func(t *testing.T) {
|
||||
resp := cfg.Response{
|
||||
Code: 200,
|
||||
RequestId: "req-789",
|
||||
Data: cfg.Data{
|
||||
RiskLevel: "high",
|
||||
// Detail deliberately empty
|
||||
},
|
||||
}
|
||||
config := makeConfig("high", "high")
|
||||
body, err := cfg.BuildDenyResponseBody(resp, config, "")
|
||||
require.NoError(t, err)
|
||||
|
||||
var result cfg.DenyResponseBody
|
||||
require.NoError(t, json.Unmarshal(body, &result))
|
||||
require.NotEmpty(t, result.BlockedDetails, "expected fallback detail from RiskLevel")
|
||||
require.Equal(t, cfg.ContentModerationType, result.BlockedDetails[0].Type)
|
||||
require.Equal(t, "high", result.BlockedDetails[0].Level)
|
||||
require.Equal(t, "block", result.BlockedDetails[0].Suggestion)
|
||||
})
|
||||
|
||||
t.Run("blockedDetails fallback from AttackLevel when Detail is empty", func(t *testing.T) {
|
||||
resp := cfg.Response{
|
||||
Code: 200,
|
||||
RequestId: "req-abc",
|
||||
Data: cfg.Data{
|
||||
AttackLevel: "high",
|
||||
// Detail deliberately empty
|
||||
},
|
||||
}
|
||||
config := makeConfig("high", "high")
|
||||
body, err := cfg.BuildDenyResponseBody(resp, config, "")
|
||||
require.NoError(t, err)
|
||||
|
||||
var result cfg.DenyResponseBody
|
||||
require.NoError(t, json.Unmarshal(body, &result))
|
||||
require.NotEmpty(t, result.BlockedDetails, "expected fallback detail from AttackLevel")
|
||||
require.Equal(t, cfg.PromptAttackType, result.BlockedDetails[0].Type)
|
||||
require.Equal(t, "high", result.BlockedDetails[0].Level)
|
||||
require.Equal(t, "block", result.BlockedDetails[0].Suggestion)
|
||||
})
|
||||
|
||||
t.Run("blockedDetails empty when risk levels below threshold", func(t *testing.T) {
|
||||
resp := cfg.Response{
|
||||
Code: 200,
|
||||
RequestId: "req-def",
|
||||
Data: cfg.Data{
|
||||
RiskLevel: "low",
|
||||
AttackLevel: "low",
|
||||
},
|
||||
}
|
||||
// threshold is "high", so "low" must not produce fallback entries
|
||||
config := makeConfig("high", "high")
|
||||
body, err := cfg.BuildDenyResponseBody(resp, config, "")
|
||||
require.NoError(t, err)
|
||||
|
||||
var result cfg.DenyResponseBody
|
||||
require.NoError(t, json.Unmarshal(body, &result))
|
||||
require.Empty(t, result.BlockedDetails)
|
||||
})
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user