jiazhizhong/higress
1
0
Fork 0
mirror of https://github.com/alibaba/higress.git synced 2026-03-16 16:30:47 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: Fix the quotation issue of deny message in ai-security-guard (#1352)

Browse Source
This commit is contained in:
Kent Dong
2024-09-27 18:45:51 +08:00
committed by GitHub
parent 1b119ed371
commit 71aae9ddf6
1 changed files with 15 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
plugins/wasm-go/extensions/ai-security-guard/main.go
View File

@@ -7,6 +7,7 @@ import (
"crypto/sha1"
"encoding/base64"
"encoding/hex"
"encoding/json"
"errors"
"fmt"
mrand "math/rand"
@@ -194,7 +195,7 @@ func onHttpRequestHeaders(ctx wrapper.HttpContext, config AISecurityConfig, log
}
func onHttpRequestBody(ctx wrapper.HttpContext, config AISecurityConfig, body []byte, log wrapper.Log) types.Action {
proxywasm.LogDebugf("checking request body...")
log.Debugf("checking request body...")
content := gjson.GetBytes(body, config.requestContentJsonPath).String()
model := gjson.GetBytes(body, "model").Raw
ctx.SetContext("requestModel", model)
@@ -231,12 +232,14 @@ func onHttpRequestBody(ctx wrapper.HttpContext, config AISecurityConfig, body []
respAdvice := respData.Get("Advice")
respResult := respData.Get("Result")
var denyMessage string
messageNeedSerialization := true
if config.protocolOriginal {
// not openai
if config.denyMessage != "" {
denyMessage = config.denyMessage
} else if respAdvice.Exists() {
denyMessage = respAdvice.Array()[0].Get("Answer").Raw
messageNeedSerialization = false
} else {
denyMessage = DefaultDenyMessage
}
@@ -244,12 +247,20 @@ func onHttpRequestBody(ctx wrapper.HttpContext, config AISecurityConfig, body []
// openai
if respAdvice.Exists() {
denyMessage = respAdvice.Array()[0].Get("Answer").Raw
messageNeedSerialization = false
} else if config.denyMessage != "" {
denyMessage = config.denyMessage
} else {
denyMessage = DefaultDenyMessage
}
}
if messageNeedSerialization {
if data, err := json.Marshal(denyMessage); err == nil {
denyMessage = string(data)
} else {
denyMessage = fmt.Sprintf("\"%s\"", DefaultDenyMessage)
}
}
if respResult.Array()[0].Get("Label").String() != "nonLabel" {
proxywasm.SetProperty([]string{TracingPrefix, "ai_sec_risklabel"}, []byte(respResult.Array()[0].Get("Label").String()))
proxywasm.SetProperty([]string{TracingPrefix, "ai_sec_deny_phase"}, []byte("request"))
@@ -280,7 +291,7 @@ func onHttpRequestBody(ctx wrapper.HttpContext, config AISecurityConfig, body []
}
return types.ActionPause
} else {
proxywasm.LogDebugf("request content is empty. skip")
log.Debugf("request content is empty. skip")
return types.ActionContinue
}
}
@@ -320,7 +331,7 @@ func onHttpResponseHeaders(ctx wrapper.HttpContext, config AISecurityConfig, log
}
func onHttpResponseBody(ctx wrapper.HttpContext, config AISecurityConfig, body []byte, log wrapper.Log) types.Action {
proxywasm.LogDebugf("checking response body...")
log.Debugf("checking response body...")
hdsMap := ctx.GetContext("headers").(map[string][]string)
isStreamingResponse := strings.Contains(strings.Join(hdsMap["content-type"], ";"), "event-stream")
model := ctx.GetStringContext("requestModel", "unknown")
@@ -411,7 +422,7 @@ func onHttpResponseBody(ctx wrapper.HttpContext, config AISecurityConfig, body [
}
return types.ActionPause
} else {
proxywasm.LogDebugf("request content is empty. skip")
log.Debugf("request content is empty. skip")
return types.ActionContinue
}
}