feat(ai-security-guard): replace denyMessage with structured DenyResponseBody (#3642)

Co-authored-by: rinfx <yucheng.lxr@alibaba-inc.com>

plugins/wasm-go/extensions/ai-security-guard/lvwang/common/text/openai.go

@@ -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)
 		}

plugins/wasm-go/extensions/ai-security-guard/lvwang/multi_modal_guard/image/openai.go

@@ -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")

plugins/wasm-go/extensions/ai-security-guard/lvwang/multi_modal_guard/image/qwen.go

@@ -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")

plugins/wasm-go/extensions/ai-security-guard/lvwang/multi_modal_guard/mcp/mcp.go

@@ -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)
 	}

plugins/wasm-go/extensions/ai-security-guard/lvwang/multi_modal_guard/text/openai.go

@@ -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)
 		}

plugins/wasm-go/extensions/ai-security-guard/lvwang/text_moderation_plus/text/openai.go

@@ -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)
 		}