solve aliyun lvwang content length limit problem (#1569)

plugins/wasm-go/extensions/ai-security-guard/go.mod

@@ -6,7 +6,7 @@ replace github.com/alibaba/higress/plugins/wasm-go => ../..
 
 require (
 	github.com/alibaba/higress/plugins/wasm-go v1.3.6-0.20240522012622-fc6a6aad8906
-	github.com/higress-group/proxy-wasm-go-sdk v0.0.0-20240711023527-ba358c48772f
+	github.com/higress-group/proxy-wasm-go-sdk v1.0.0
 	github.com/tidwall/gjson v1.17.3
 )
 

plugins/wasm-go/extensions/ai-security-guard/go.sum

@@ -3,8 +3,8 @@ github.com/google/uuid v1.3.0 h1:t6JiXgmwXMjEs8VusXIJk2BXHsn+wx8BZdTaoZ5fu7I=
 github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/higress-group/nottinygc v0.0.0-20231101025119-e93c4c2f8520 h1:IHDghbGQ2DTIXHBHxWfqCYQW1fKjyJ/I7W1pMyUDeEA=
 github.com/higress-group/nottinygc v0.0.0-20231101025119-e93c4c2f8520/go.mod h1:Nz8ORLaFiLWotg6GeKlJMhv8cci8mM43uEnLA5t8iew=
-github.com/higress-group/proxy-wasm-go-sdk v0.0.0-20240711023527-ba358c48772f h1:ZIiIBRvIw62gA5MJhuwp1+2wWbqL9IGElQ499rUsYYg=
+github.com/higress-group/proxy-wasm-go-sdk v1.0.0 h1:BZRNf4R7jr9hwRivg/E29nkVaKEak5MWjBDhWjuHijU=
-github.com/higress-group/proxy-wasm-go-sdk v0.0.0-20240711023527-ba358c48772f/go.mod h1:hNFjhrLUIq+kJ9bOcs8QtiplSQ61GZXtd2xHKx4BYRo=
+github.com/higress-group/proxy-wasm-go-sdk v1.0.0/go.mod h1:iiSyFbo+rAtbtGt/bsefv8GU57h9CCLYGJA74/tF5/0=
 github.com/magefile/mage v1.14.0 h1:6QDX3g6z1YvJ4olPhT1wksUcSa/V0a1B+pJb73fBjyo=
 github.com/magefile/mage v1.14.0/go.mod h1:z5UZb/iS3GoOSn0JgWuiw7dxlurVYTu+/jHXqQg881A=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=

plugins/wasm-go/extensions/ai-security-guard/main.go

@@ -55,6 +55,7 @@ const (
 	DefaultDenyMessage               = "很抱歉，我无法回答您的问题"
 
 	AliyunUserAgent = "CIPFrom/AIGateway"
+	LengthLimit     = 1800
 )
 
 type Response struct {
@@ -260,73 +261,92 @@ func onHttpRequestBody(ctx wrapper.HttpContext, config AISecurityConfig, body []
 		log.Info("request content is empty. skip")
 		return types.ActionContinue
 	}
-	timestamp := time.Now().UTC().Format("2006-01-02T15:04:05Z")
+	contentIndex := 0
-	randomID, _ := generateHexID(16)
+	sessionID, _ := generateHexID(20)
-	params := map[string]string{
+	var singleCall func()
-		"Format":            "JSON",
+	callback := func(statusCode int, responseHeaders http.Header, responseBody []byte) {
-		"Version":           "2022-03-02",
+		log.Info(string(responseBody))
-		"SignatureMethod":   "Hmac-SHA1",
+		if statusCode != 200 || gjson.GetBytes(responseBody, "Code").Int() != 200 {
-		"SignatureNonce":    randomID,
+			proxywasm.ResumeHttpRequest()
-		"SignatureVersion":  "1.0",
+			return
-		"Action":            "TextModerationPlus",
+		}
-		"AccessKeyId":       config.ak,
+		var response Response
-		"Timestamp":         timestamp,
+		err := json.Unmarshal(responseBody, &response)
-		"Service":           config.requestCheckService,
+		if err != nil {
-		"ServiceParameters": fmt.Sprintf(`{"content": "%s"}`, marshalStr(content, log)),
+			log.Error("failed to unmarshal aliyun content security response at request phase")
-	}
+			proxywasm.ResumeHttpRequest()
-	if config.token != "" {
+			return
-		params["SecurityToken"] = config.token
+		}
-	}
+		if riskLevelToInt(response.Data.RiskLevel) < riskLevelToInt(config.riskLevelBar) {
-	signature := getSign(params, config.sk+"&")
+			if contentIndex >= len(content) {
-	reqParams := url.Values{}
-	for k, v := range params {
-		reqParams.Add(k, v)
-	}
-	reqParams.Add("Signature", signature)
-	err := config.client.Post(fmt.Sprintf("/?%s", reqParams.Encode()), [][2]string{{"User-Agent", AliyunUserAgent}}, nil,
-		func(statusCode int, responseHeaders http.Header, responseBody []byte) {
-			log.Info(string(responseBody))
-			if statusCode != 200 || gjson.GetBytes(responseBody, "Code").Int() != 200 {
 				proxywasm.ResumeHttpRequest()
-				return
-			}
-			var response Response
-			err := json.Unmarshal(responseBody, &response)
-			if err != nil {
-				log.Error("failed to unmarshal aliyun content security response at request phase")
-				proxywasm.ResumeHttpRequest()
-				return
-			}
-			if riskLevelToInt(response.Data.RiskLevel) < riskLevelToInt(config.riskLevelBar) {
-				proxywasm.ResumeHttpRequest()
-				return
-			}
-			denyMessage := DefaultDenyMessage
-			if config.denyMessage != "" {
-				denyMessage = config.denyMessage
-			} else if response.Data.Advice != nil && response.Data.Advice[0].Answer != "" {
-				denyMessage = response.Data.Advice[0].Answer
-			}
-			marshalledDenyMessage := marshalStr(denyMessage, log)
-			if config.protocolOriginal {
-				proxywasm.SendHttpResponse(uint32(config.denyCode), [][2]string{{"content-type", "application/json"}}, []byte(marshalledDenyMessage), -1)
-			} else if gjson.GetBytes(body, "stream").Bool() {
-				randomID := generateRandomID()
-				jsonData := []byte(fmt.Sprintf(OpenAIStreamResponseFormat, randomID, model, marshalledDenyMessage, randomID, model))
-				proxywasm.SendHttpResponse(uint32(config.denyCode), [][2]string{{"content-type", "text/event-stream;charset=UTF-8"}}, jsonData, -1)
 			} else {
-				randomID := generateRandomID()
+				singleCall()
-				jsonData := []byte(fmt.Sprintf(OpenAIResponseFormat, randomID, model, marshalledDenyMessage))
-				proxywasm.SendHttpResponse(uint32(config.denyCode), [][2]string{{"content-type", "application/json"}}, jsonData, -1)
 			}
-			ctx.DontReadResponseBody()
+			return
-			config.incrementCounter("ai_sec_request_deny", 1)
+		}
-		},
+		denyMessage := DefaultDenyMessage
-	)
+		if config.denyMessage != "" {
-	if err != nil {
+			denyMessage = config.denyMessage
-		log.Errorf("failed call the safe check service: %v", err)
+		} else if response.Data.Advice != nil && response.Data.Advice[0].Answer != "" {
-		return types.ActionContinue
+			denyMessage = response.Data.Advice[0].Answer
+		}
+		marshalledDenyMessage := marshalStr(denyMessage, log)
+		if config.protocolOriginal {
+			proxywasm.SendHttpResponse(uint32(config.denyCode), [][2]string{{"content-type", "application/json"}}, []byte(marshalledDenyMessage), -1)
+		} else if gjson.GetBytes(body, "stream").Bool() {
+			randomID := generateRandomID()
+			jsonData := []byte(fmt.Sprintf(OpenAIStreamResponseFormat, randomID, model, marshalledDenyMessage, randomID, model))
+			proxywasm.SendHttpResponse(uint32(config.denyCode), [][2]string{{"content-type", "text/event-stream;charset=UTF-8"}}, jsonData, -1)
+		} else {
+			randomID := generateRandomID()
+			jsonData := []byte(fmt.Sprintf(OpenAIResponseFormat, randomID, model, marshalledDenyMessage))
+			proxywasm.SendHttpResponse(uint32(config.denyCode), [][2]string{{"content-type", "application/json"}}, jsonData, -1)
+		}
+		ctx.DontReadResponseBody()
+		config.incrementCounter("ai_sec_request_deny", 1)
+		proxywasm.ResumeHttpRequest()
 	}
+	singleCall = func() {
+		timestamp := time.Now().UTC().Format("2006-01-02T15:04:05Z")
+		randomID, _ := generateHexID(16)
+		var nextContentIndex int
+		if contentIndex+LengthLimit >= len(content) {
+			nextContentIndex = len(content)
+		} else {
+			nextContentIndex = contentIndex + LengthLimit
+		}
+		contentPiece := content[contentIndex:nextContentIndex]
+		contentIndex = nextContentIndex
+		log.Debugf("current content piece: %s", contentPiece)
+		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":           config.requestCheckService,
+			"ServiceParameters": fmt.Sprintf(`{"sessionId": "%s","content": "%s"}`, sessionID, marshalStr(contentPiece, log)),
+		}
+		if config.token != "" {
+			params["SecurityToken"] = config.token
+		}
+		signature := getSign(params, config.sk+"&")
+		reqParams := url.Values{}
+		for k, v := range params {
+			reqParams.Add(k, v)
+		}
+		reqParams.Add("Signature", signature)
+		err := config.client.Post(fmt.Sprintf("/?%s", reqParams.Encode()), [][2]string{{"User-Agent", AliyunUserAgent}}, nil, callback)
+		if err != nil {
+			log.Errorf("failed call the safe check service: %v", err)
+			proxywasm.ResumeHttpRequest()
+		}
+	}
+	singleCall()
 	return types.ActionPause
 }
 
@@ -385,73 +405,94 @@ func onHttpResponseBody(ctx wrapper.HttpContext, config AISecurityConfig, body [
 		log.Info("response content is empty. skip")
 		return types.ActionContinue
 	}
-	timestamp := time.Now().UTC().Format("2006-01-02T15:04:05Z")
+	contentIndex := 0
-	randomID, _ := generateHexID(16)
+	sessionID, _ := generateHexID(20)
-	params := map[string]string{
+	var singleCall func()
-		"Format":            "JSON",
+	callback := func(statusCode int, responseHeaders http.Header, responseBody []byte) {
-		"Version":           "2022-03-02",
+		log.Info(string(responseBody))
-		"SignatureMethod":   "Hmac-SHA1",
+		if statusCode != 200 || gjson.GetBytes(responseBody, "Code").Int() != 200 {
-		"SignatureNonce":    randomID,
+			proxywasm.ResumeHttpResponse()
-		"SignatureVersion":  "1.0",
+			return
-		"Action":            "TextModerationPlus",
+		}
-		"AccessKeyId":       config.ak,
+		var response Response
-		"Timestamp":         timestamp,
+		err := json.Unmarshal(responseBody, &response)
-		"Service":           config.responseCheckService,
+		if err != nil {
-		"ServiceParameters": fmt.Sprintf(`{"content": "%s"}`, marshalStr(content, log)),
+			log.Error("failed to unmarshal aliyun content security response at response phase")
-	}
+			proxywasm.ResumeHttpResponse()
-	if config.token != "" {
+			return
-		params["SecurityToken"] = config.token
+		}
-	}
+		if riskLevelToInt(response.Data.RiskLevel) < riskLevelToInt(config.riskLevelBar) {
-	signature := getSign(params, config.sk+"&")
+			if contentIndex >= len(content) {
-	reqParams := url.Values{}
+				proxywasm.ResumeHttpResponse()
-	for k, v := range params {
-		reqParams.Add(k, v)
-	}
-	reqParams.Add("Signature", signature)
-	err := config.client.Post(fmt.Sprintf("/?%s", reqParams.Encode()), [][2]string{{"User-Agent", AliyunUserAgent}}, nil,
-		func(statusCode int, responseHeaders http.Header, responseBody []byte) {
-			defer proxywasm.ResumeHttpResponse()
-			log.Info(string(responseBody))
-			if statusCode != 200 || gjson.GetBytes(responseBody, "Code").Int() != 200 {
-				return
-			}
-			var response Response
-			err := json.Unmarshal(responseBody, &response)
-			if err != nil {
-				log.Error("failed to unmarshal aliyun content security response at response phase")
-				return
-			}
-			if riskLevelToInt(response.Data.RiskLevel) < riskLevelToInt(config.riskLevelBar) {
-				return
-			}
-			denyMessage := DefaultDenyMessage
-			if config.denyMessage != "" {
-				denyMessage = config.denyMessage
-			} else if response.Data.Advice != nil && response.Data.Advice[0].Answer != "" {
-				denyMessage = response.Data.Advice[0].Answer
-			}
-			marshalledDenyMessage := marshalStr(denyMessage, log)
-			var jsonData []byte
-			if config.protocolOriginal {
-				jsonData = []byte(marshalledDenyMessage)
-			} else if isStreamingResponse {
-				randomID := generateRandomID()
-				jsonData = []byte(fmt.Sprintf(OpenAIStreamResponseFormat, randomID, model, marshalledDenyMessage, randomID, model))
 			} else {
-				randomID := generateRandomID()
+				singleCall()
-				jsonData = []byte(fmt.Sprintf(OpenAIResponseFormat, randomID, model, marshalledDenyMessage))
 			}
-			delete(hdsMap, "content-length")
+			return
-			hdsMap[":status"] = []string{fmt.Sprint(config.denyCode)}
+		}
-			proxywasm.ReplaceHttpResponseHeaders(reconvertHeaders(hdsMap))
+		denyMessage := DefaultDenyMessage
-			proxywasm.ReplaceHttpResponseBody(jsonData)
+		if config.denyMessage != "" {
-			config.incrementCounter("ai_sec_response_deny", 1)
+			denyMessage = config.denyMessage
-		},
+		} else if response.Data.Advice != nil && response.Data.Advice[0].Answer != "" {
-	)
+			denyMessage = response.Data.Advice[0].Answer
-	if err != nil {
+		}
-		log.Errorf("failed call the safe check service: %v", err)
+		marshalledDenyMessage := marshalStr(denyMessage, log)
-		return types.ActionContinue
+		var jsonData []byte
+		if config.protocolOriginal {
+			jsonData = []byte(marshalledDenyMessage)
+		} else if isStreamingResponse {
+			randomID := generateRandomID()
+			jsonData = []byte(fmt.Sprintf(OpenAIStreamResponseFormat, randomID, model, marshalledDenyMessage, randomID, model))
+		} else {
+			randomID := generateRandomID()
+			jsonData = []byte(fmt.Sprintf(OpenAIResponseFormat, randomID, model, marshalledDenyMessage))
+		}
+		delete(hdsMap, "content-length")
+		hdsMap[":status"] = []string{fmt.Sprint(config.denyCode)}
+		proxywasm.ReplaceHttpResponseHeaders(reconvertHeaders(hdsMap))
+		proxywasm.ReplaceHttpResponseBody(jsonData)
+		config.incrementCounter("ai_sec_response_deny", 1)
+		proxywasm.ResumeHttpResponse()
 	}
+	singleCall = func() {
+		timestamp := time.Now().UTC().Format("2006-01-02T15:04:05Z")
+		randomID, _ := generateHexID(16)
+		var nextContentIndex int
+		if contentIndex+LengthLimit >= len(content) {
+			nextContentIndex = len(content)
+		} else {
+			nextContentIndex = contentIndex + LengthLimit
+		}
+		contentPiece := content[contentIndex:nextContentIndex]
+		contentIndex = nextContentIndex
+		log.Debugf("current content piece: %s", contentPiece)
+		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":           config.responseCheckService,
+			"ServiceParameters": fmt.Sprintf(`{"sessionId": "%s","content": "%s"}`, sessionID, marshalStr(contentPiece, log)),
+		}
+		if config.token != "" {
+			params["SecurityToken"] = config.token
+		}
+		signature := getSign(params, config.sk+"&")
+		reqParams := url.Values{}
+		for k, v := range params {
+			reqParams.Add(k, v)
+		}
+		reqParams.Add("Signature", signature)
+		err := config.client.Post(fmt.Sprintf("/?%s", reqParams.Encode()), [][2]string{{"User-Agent", AliyunUserAgent}}, nil, callback)
+		if err != nil {
+			log.Errorf("failed call the safe check service: %v", err)
+			proxywasm.ResumeHttpResponse()
+		}
+	}
+	singleCall()
 	return types.ActionPause
 }