From 496346fe95eda4a00f07862c765b82e3c2c54ae8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=BE=84=E6=BD=AD?= Date: Mon, 26 Aug 2024 13:33:44 +0800 Subject: [PATCH] fix ai proxy check request body logic (#1252) --- plugins/wasm-go/extensions/ai-proxy/main.go | 8 ++++---- plugins/wasm-go/pkg/wrapper/request_wrapper.go | 9 +++++++-- 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/plugins/wasm-go/extensions/ai-proxy/main.go b/plugins/wasm-go/extensions/ai-proxy/main.go index 807b48bce..56e9f5808 100644 --- a/plugins/wasm-go/extensions/ai-proxy/main.go +++ b/plugins/wasm-go/extensions/ai-proxy/main.go @@ -74,16 +74,16 @@ func onHttpRequestHeader(ctx wrapper.HttpContext, pluginConfig config.PluginConf if handler, ok := activeProvider.(provider.RequestHeadersHandler); ok { // Disable the route re-calculation since the plugin may modify some headers related to the chosen route. ctx.DisableReroute() - - _, err := handler.OnRequestHeaders(ctx, apiName, log) + hasRequestBody := wrapper.HasRequestBody() + action, err := handler.OnRequestHeaders(ctx, apiName, log) if err == nil { - if wrapper.HasRequestBody() { + if hasRequestBody { ctx.SetRequestBodyBufferLimit(defaultMaxBodyBytes) // Always return types.HeaderStopIteration to support fallback routing, // as long as onHttpRequestBody can be called. return types.HeaderStopIteration } - return types.ActionContinue + return action } _ = util.SendResponse(500, "ai-proxy.proc_req_headers_failed", util.MimeTypeTextPlain, fmt.Sprintf("failed to process request headers: %v", err)) return types.ActionContinue diff --git a/plugins/wasm-go/pkg/wrapper/request_wrapper.go b/plugins/wasm-go/pkg/wrapper/request_wrapper.go index 9a33f04d0..bfbba0cdb 100644 --- a/plugins/wasm-go/pkg/wrapper/request_wrapper.go +++ b/plugins/wasm-go/pkg/wrapper/request_wrapper.go @@ -84,14 +84,19 @@ func IsBinaryResponseBody() bool { } func HasRequestBody() bool { + contentTypeStr, _ := proxywasm.GetHttpRequestHeader("content-type") contentLengthStr, _ := proxywasm.GetHttpRequestHeader("content-length") + transferEncodingStr, _ := proxywasm.GetHttpRequestHeader("transfer-encoding") + proxywasm.LogDebugf("check has request 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 } } - - transferEncodingStr, _ := proxywasm.GetHttpRequestHeader("transfer-encoding") return strings.Contains(transferEncodingStr, "chunked") }