mirror of
https://github.com/alibaba/higress.git
synced 2026-06-09 20:57:32 +08:00
add plugin: ai-security-guard (#1034)
This commit is contained in:
rinfx
GitHub
272
plugins/wasm-go/extensions/ai-security-guard/main.go
Normal file
272
plugins/wasm-go/extensions/ai-security-guard/main.go
Normal file
@@ -0,0 +1,272 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"crypto/hmac"
|
||||
"crypto/rand"
|
||||
"crypto/sha1"
|
||||
"encoding/base64"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/alibaba/higress/plugins/wasm-go/pkg/wrapper"
|
||||
"github.com/higress-group/proxy-wasm-go-sdk/proxywasm"
|
||||
"github.com/higress-group/proxy-wasm-go-sdk/proxywasm/types"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
func main() {
|
||||
wrapper.SetCtx(
|
||||
"ai-security-guard",
|
||||
wrapper.ParseConfigBy(parseConfig),
|
||||
wrapper.ProcessRequestHeadersBy(onHttpRequestHeaders),
|
||||
wrapper.ProcessRequestBodyBy(onHttpRequestBody),
|
||||
wrapper.ProcessResponseHeadersBy(onHttpResponseHeaders),
|
||||
wrapper.ProcessResponseBodyBy(onHttpResponseBody),
|
||||
)
|
||||
}
|
||||
|
||||
type AISecurityConfig struct {
|
||||
client wrapper.HttpClient
|
||||
ak string
|
||||
sk string
|
||||
}
|
||||
|
||||
type StandardResponse struct {
|
||||
Code int `json:"Code"`
|
||||
Phase string `json:"BlockPhase"`
|
||||
Message string `json:"Message"`
|
||||
}
|
||||
|
||||
func urlEncoding(rawStr string) string {
|
||||
encodedStr := url.PathEscape(rawStr)
|
||||
encodedStr = strings.ReplaceAll(encodedStr, "+", "%20")
|
||||
encodedStr = strings.ReplaceAll(encodedStr, ":", "%3A")
|
||||
encodedStr = strings.ReplaceAll(encodedStr, "=", "%3D")
|
||||
encodedStr = strings.ReplaceAll(encodedStr, "&", "%26")
|
||||
return encodedStr
|
||||
}
|
||||
|
||||
func hmacSha1(message, secret string) string {
|
||||
key := []byte(secret)
|
||||
h := hmac.New(sha1.New, key)
|
||||
h.Write([]byte(message))
|
||||
hash := h.Sum(nil)
|
||||
return base64.StdEncoding.EncodeToString(hash)
|
||||
}
|
||||
|
||||
func getSign(params map[string]string, secret string) string {
|
||||
paramArray := []string{}
|
||||
for k, v := range params {
|
||||
paramArray = append(paramArray, urlEncoding(k)+"="+urlEncoding(v))
|
||||
}
|
||||
sort.Slice(paramArray, func(i, j int) bool {
|
||||
return paramArray[i] <= paramArray[j]
|
||||
})
|
||||
canonicalStr := strings.Join(paramArray, "&")
|
||||
signStr := "POST&%2F&" + urlEncoding(canonicalStr)
|
||||
fmt.Println(signStr)
|
||||
return hmacSha1(signStr, secret)
|
||||
}
|
||||
|
||||
func generateHexID(length int) (string, error) {
|
||||
bytes := make([]byte, length/2)
|
||||
if _, err := rand.Read(bytes); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return hex.EncodeToString(bytes), nil
|
||||
}
|
||||
|
||||
func parseConfig(json gjson.Result, config *AISecurityConfig, log wrapper.Log) error {
|
||||
serviceName := json.Get("serviceName").String()
|
||||
servicePort := json.Get("servicePort").Int()
|
||||
domain := json.Get("domain").String()
|
||||
config.ak = json.Get("ak").String()
|
||||
config.sk = json.Get("sk").String()
|
||||
if serviceName == "" || servicePort == 0 || domain == "" {
|
||||
return errors.New("invalid service config")
|
||||
}
|
||||
config.client = wrapper.NewClusterClient(wrapper.DnsCluster{
|
||||
ServiceName: serviceName,
|
||||
Port: servicePort,
|
||||
Domain: domain,
|
||||
})
|
||||
return nil
|
||||
}
|
||||
|
||||
func onHttpRequestHeaders(ctx wrapper.HttpContext, config AISecurityConfig, log wrapper.Log) types.Action {
|
||||
return types.ActionContinue
|
||||
}
|
||||
|
||||
func onHttpRequestBody(ctx wrapper.HttpContext, config AISecurityConfig, body []byte, log wrapper.Log) types.Action {
|
||||
messages := gjson.GetBytes(body, "messages").Array()
|
||||
if len(messages) > 0 {
|
||||
role := messages[len(messages)-1].Get("role").String()
|
||||
content := messages[len(messages)-1].Get("content").String()
|
||||
if role != "user" {
|
||||
return types.ActionContinue
|
||||
}
|
||||
timestamp := time.Now().UTC().Format("2006-01-02T15:04:05Z")
|
||||
randomID, _ := generateHexID(16)
|
||||
params := map[string]string{
|
||||
"Format": "JSON",
|
||||
"Version": "2022-03-02",
|
||||
"SignatureMethod": "Hmac-SHA1",
|
||||
"SignatureNonce": randomID,
|
||||
"SignatureVersion": "1.0",
|
||||
"Action": "TextModerationPlus",
|
||||
"AccessKeyId": config.ak,
|
||||
"Timestamp": timestamp,
|
||||
"Service": "llm_query_moderation",
|
||||
"ServiceParameters": `{"content": "` + content + `"}`,
|
||||
}
|
||||
signature := getSign(params, config.sk+"&")
|
||||
reqParams := url.Values{}
|
||||
for k, v := range params {
|
||||
reqParams.Add(k, v)
|
||||
}
|
||||
reqParams.Add("Signature", signature)
|
||||
config.client.Post(fmt.Sprintf("/?%s", reqParams.Encode()), nil, nil,
|
||||
func(statusCode int, responseHeaders http.Header, responseBody []byte) {
|
||||
respData := gjson.GetBytes(responseBody, "Data")
|
||||
if respData.Exists() {
|
||||
respAdvice := respData.Get("Advice")
|
||||
respResult := respData.Get("Result")
|
||||
if respAdvice.Exists() {
|
||||
sr := StandardResponse{
|
||||
Code: 403,
|
||||
Phase: "Request",
|
||||
Message: respAdvice.Array()[0].Get("Answer").String(),
|
||||
}
|
||||
jsonData, _ := json.MarshalIndent(sr, "", " ")
|
||||
proxywasm.SetProperty([]string{"risklabel"}, []byte(respResult.Array()[0].Get("Label").String()))
|
||||
proxywasm.SendHttpResponse(403, [][2]string{{"content-type", "application/json"}}, jsonData, -1)
|
||||
} else if respResult.Array()[0].Get("Label").String() != "nonLabel" {
|
||||
sr := StandardResponse{
|
||||
Code: 403,
|
||||
Phase: "Request",
|
||||
Message: "risk detected",
|
||||
}
|
||||
jsonData, _ := json.MarshalIndent(sr, "", " ")
|
||||
proxywasm.SetProperty([]string{"risklabel"}, []byte(respResult.Array()[0].Get("Label").String()))
|
||||
proxywasm.SendHttpResponse(403, [][2]string{{"content-type", "application/json"}}, jsonData, -1)
|
||||
} else {
|
||||
proxywasm.ResumeHttpRequest()
|
||||
}
|
||||
} else {
|
||||
proxywasm.ResumeHttpRequest()
|
||||
}
|
||||
},
|
||||
)
|
||||
return types.ActionPause
|
||||
} else {
|
||||
return types.ActionContinue
|
||||
}
|
||||
}
|
||||
|
||||
func convertHeaders(hs [][2]string) map[string][]string {
|
||||
ret := make(map[string][]string)
|
||||
for _, h := range hs {
|
||||
k, v := strings.ToLower(h[0]), h[1]
|
||||
ret[k] = append(ret[k], v)
|
||||
}
|
||||
return ret
|
||||
}
|
||||
|
||||
// headers: map[string][]string -> [][2]string
|
||||
func reconvertHeaders(hs map[string][]string) [][2]string {
|
||||
var ret [][2]string
|
||||
for k, vs := range hs {
|
||||
for _, v := range vs {
|
||||
ret = append(ret, [2]string{k, v})
|
||||
}
|
||||
}
|
||||
sort.SliceStable(ret, func(i, j int) bool {
|
||||
return ret[i][0] < ret[j][0]
|
||||
})
|
||||
return ret
|
||||
}
|
||||
|
||||
func onHttpResponseHeaders(ctx wrapper.HttpContext, config AISecurityConfig, log wrapper.Log) types.Action {
|
||||
headers, err := proxywasm.GetHttpResponseHeaders()
|
||||
if err != nil {
|
||||
log.Warnf("failed to get response headers: %v", err)
|
||||
return types.ActionContinue
|
||||
}
|
||||
hdsMap := convertHeaders(headers)
|
||||
ctx.SetContext("headers", hdsMap)
|
||||
return types.HeaderStopIteration
|
||||
}
|
||||
|
||||
func onHttpResponseBody(ctx wrapper.HttpContext, config AISecurityConfig, body []byte, log wrapper.Log) types.Action {
|
||||
messages := gjson.GetBytes(body, "choices").Array()
|
||||
if len(messages) > 0 {
|
||||
content := messages[0].Get("message").Get("content").String()
|
||||
timestamp := time.Now().UTC().Format("2006-01-02T15:04:05Z")
|
||||
randomID, _ := generateHexID(16)
|
||||
params := map[string]string{
|
||||
"Format": "JSON",
|
||||
"Version": "2022-03-02",
|
||||
"SignatureMethod": "Hmac-SHA1",
|
||||
"SignatureNonce": randomID,
|
||||
"SignatureVersion": "1.0",
|
||||
"Action": "TextModerationPlus",
|
||||
"AccessKeyId": config.ak,
|
||||
"Timestamp": timestamp,
|
||||
"Service": "llm_response_moderation",
|
||||
"ServiceParameters": `{"content": "` + content + `"}`,
|
||||
}
|
||||
signature := getSign(params, config.sk+"&")
|
||||
reqParams := url.Values{}
|
||||
for k, v := range params {
|
||||
reqParams.Add(k, v)
|
||||
}
|
||||
reqParams.Add("Signature", signature)
|
||||
config.client.Post(fmt.Sprintf("/?%s", reqParams.Encode()), nil, nil,
|
||||
func(statusCode int, responseHeaders http.Header, responseBody []byte) {
|
||||
defer proxywasm.ResumeHttpResponse()
|
||||
respData := gjson.GetBytes(responseBody, "Data")
|
||||
if respData.Exists() {
|
||||
respAdvice := respData.Get("Advice")
|
||||
respResult := respData.Get("Result")
|
||||
if respAdvice.Exists() {
|
||||
sr := StandardResponse{
|
||||
Code: 403,
|
||||
Phase: "Response",
|
||||
Message: respAdvice.Array()[0].Get("Answer").String(),
|
||||
}
|
||||
jsonData, _ := json.MarshalIndent(sr, "", " ")
|
||||
hdsMap := ctx.GetContext("headers").(map[string][]string)
|
||||
delete(hdsMap, "content-length")
|
||||
hdsMap[":status"] = []string{"403"}
|
||||
proxywasm.ReplaceHttpResponseHeaders(reconvertHeaders(hdsMap))
|
||||
proxywasm.ReplaceHttpResponseBody(jsonData)
|
||||
proxywasm.SetProperty([]string{"risklabel"}, []byte(respResult.Array()[0].Get("Label").String()))
|
||||
} else if respResult.Array()[0].Get("Label").String() != "nonLabel" {
|
||||
sr := StandardResponse{
|
||||
Code: 403,
|
||||
Phase: "Response",
|
||||
Message: "risk detected",
|
||||
}
|
||||
jsonData, _ := json.MarshalIndent(sr, "", " ")
|
||||
hdsMap := ctx.GetContext("headers").(map[string][]string)
|
||||
delete(hdsMap, "content-length")
|
||||
hdsMap[":status"] = []string{"403"}
|
||||
proxywasm.ReplaceHttpResponseHeaders(reconvertHeaders(hdsMap))
|
||||
proxywasm.ReplaceHttpResponseBody(jsonData)
|
||||
proxywasm.SetProperty([]string{"risklabel"}, []byte(respResult.Array()[0].Get("Label").String()))
|
||||
}
|
||||
}
|
||||
},
|
||||
)
|
||||
return types.ActionPause
|
||||
} else {
|
||||
return types.ActionContinue
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user