From 6f18814303eb71a813c51d99338c7031d41ba7f6 Mon Sep 17 00:00:00 2001 From: rinfx <893383980@qq.com> Date: Thu, 24 Oct 2024 11:24:43 +0800 Subject: [PATCH 1/3] bugfix: plugin will block GET request --- .../extensions/ai-security-guard/main.go | 14 ++++++++----- .../wasm-go/pkg/wrapper/response_wrapper.go | 21 +++++++++++++++++++ 2 files changed, 30 insertions(+), 5 deletions(-) diff --git a/plugins/wasm-go/extensions/ai-security-guard/main.go b/plugins/wasm-go/extensions/ai-security-guard/main.go index ca59f7f6a1..4a38db3666 100644 --- a/plugins/wasm-go/extensions/ai-security-guard/main.go +++ b/plugins/wasm-go/extensions/ai-security-guard/main.go @@ -187,10 +187,6 @@ func onHttpRequestHeaders(ctx wrapper.HttpContext, config AISecurityConfig, log log.Debugf("request checking is disabled") ctx.DontReadRequestBody() } - if !config.checkResponse { - log.Debugf("response checking is disabled") - ctx.DontReadResponseBody() - } return types.ActionContinue } @@ -199,7 +195,7 @@ func onHttpRequestBody(ctx wrapper.HttpContext, config AISecurityConfig, body [] content := gjson.GetBytes(body, config.requestContentJsonPath).Raw model := gjson.GetBytes(body, "model").Raw ctx.SetContext("requestModel", model) - log.Debugf("Raw response content is: %s", content) + log.Debugf("Raw request content is: %s", content) if len(content) > 0 { timestamp := time.Now().UTC().Format("2006-01-02T15:04:05Z") randomID, _ := generateHexID(16) @@ -326,6 +322,14 @@ func onHttpResponseHeaders(ctx wrapper.HttpContext, config AISecurityConfig, log log.Warnf("failed to get response headers: %v", err) return types.ActionContinue } + if !config.checkResponse { + log.Debugf("response checking is disabled") + ctx.DontReadResponseBody() + return types.ActionContinue + } + if !wrapper.HasResponseBody() { + return types.ActionContinue + } hdsMap := convertHeaders(headers) ctx.SetContext("headers", hdsMap) return types.HeaderStopIteration diff --git a/plugins/wasm-go/pkg/wrapper/response_wrapper.go b/plugins/wasm-go/pkg/wrapper/response_wrapper.go index 3cd91daf3b..4242b953af 100644 --- a/plugins/wasm-go/pkg/wrapper/response_wrapper.go +++ b/plugins/wasm-go/pkg/wrapper/response_wrapper.go @@ -15,6 +15,9 @@ package wrapper import ( + "strconv" + "strings" + "github.com/higress-group/proxy-wasm-go-sdk/proxywasm" ) @@ -26,3 +29,21 @@ func IsResponseFromUpstream() bool { return false } } + +func HasResponseBody() bool { + contentTypeStr, _ := proxywasm.GetHttpResponseHeader("content-type") + contentLengthStr, _ := proxywasm.GetHttpResponseHeader("content-length") + transferEncodingStr, _ := proxywasm.GetHttpResponseHeader("transfer-encoding") + proxywasm.LogDebugf("check has response body: contentType:%s, contentLengthStr:%s, transferEncodingStr:%s", + contentTypeStr, contentLengthStr, transferEncodingStr) + if contentTypeStr != "" { + return true + } + if contentLengthStr != "" { + contentLength, err := strconv.Atoi(contentLengthStr) + if err == nil && contentLength > 0 { + return true + } + } + return strings.Contains(transferEncodingStr, "chunked") +} From 3fbe662ec4881e3a8fbbdc5e5c83f55891490309 Mon Sep 17 00:00:00 2001 From: rinfx <893383980@qq.com> Date: Thu, 24 Oct 2024 11:29:14 +0800 Subject: [PATCH 2/3] update --- .../extensions/ai-security-guard/main.go | 5 +---- .../wasm-go/pkg/wrapper/response_wrapper.go | 21 ------------------- 2 files changed, 1 insertion(+), 25 deletions(-) diff --git a/plugins/wasm-go/extensions/ai-security-guard/main.go b/plugins/wasm-go/extensions/ai-security-guard/main.go index 4a38db3666..fdc8903a22 100644 --- a/plugins/wasm-go/extensions/ai-security-guard/main.go +++ b/plugins/wasm-go/extensions/ai-security-guard/main.go @@ -327,9 +327,6 @@ func onHttpResponseHeaders(ctx wrapper.HttpContext, config AISecurityConfig, log ctx.DontReadResponseBody() return types.ActionContinue } - if !wrapper.HasResponseBody() { - return types.ActionContinue - } hdsMap := convertHeaders(headers) ctx.SetContext("headers", hdsMap) return types.HeaderStopIteration @@ -403,7 +400,7 @@ func onHttpResponseBody(ctx wrapper.HttpContext, config AISecurityConfig, body [ var jsonData []byte if config.protocolOriginal { jsonData = []byte(denyMessage) - } else if strings.Contains(strings.Join(hdsMap["content-type"], ";"), "event-stream") { + } else if isStreamingResponse { randomID := generateRandomID() jsonData = []byte(fmt.Sprintf(OpenAIStreamResponseFormat, randomID, model, denyMessage, randomID, model)) } else { diff --git a/plugins/wasm-go/pkg/wrapper/response_wrapper.go b/plugins/wasm-go/pkg/wrapper/response_wrapper.go index 4242b953af..3cd91daf3b 100644 --- a/plugins/wasm-go/pkg/wrapper/response_wrapper.go +++ b/plugins/wasm-go/pkg/wrapper/response_wrapper.go @@ -15,9 +15,6 @@ package wrapper import ( - "strconv" - "strings" - "github.com/higress-group/proxy-wasm-go-sdk/proxywasm" ) @@ -29,21 +26,3 @@ func IsResponseFromUpstream() bool { return false } } - -func HasResponseBody() bool { - contentTypeStr, _ := proxywasm.GetHttpResponseHeader("content-type") - contentLengthStr, _ := proxywasm.GetHttpResponseHeader("content-length") - transferEncodingStr, _ := proxywasm.GetHttpResponseHeader("transfer-encoding") - proxywasm.LogDebugf("check has response body: contentType:%s, contentLengthStr:%s, transferEncodingStr:%s", - contentTypeStr, contentLengthStr, transferEncodingStr) - if contentTypeStr != "" { - return true - } - if contentLengthStr != "" { - contentLength, err := strconv.Atoi(contentLengthStr) - if err == nil && contentLength > 0 { - return true - } - } - return strings.Contains(transferEncodingStr, "chunked") -} From 01b6046667c64dd4da888be2689c2a76218ba0af Mon Sep 17 00:00:00 2001 From: rinfx <893383980@qq.com> Date: Thu, 24 Oct 2024 15:52:50 +0800 Subject: [PATCH 3/3] update --- plugins/wasm-go/extensions/ai-security-guard/main.go | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/plugins/wasm-go/extensions/ai-security-guard/main.go b/plugins/wasm-go/extensions/ai-security-guard/main.go index fdc8903a22..5b61589616 100644 --- a/plugins/wasm-go/extensions/ai-security-guard/main.go +++ b/plugins/wasm-go/extensions/ai-security-guard/main.go @@ -317,16 +317,16 @@ func reconvertHeaders(hs map[string][]string) [][2]string { } 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 - } if !config.checkResponse { log.Debugf("response checking is disabled") ctx.DontReadResponseBody() return types.ActionContinue } + 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