From f6ee4ed166e3487b0a06a39dbf04a1533d192aff Mon Sep 17 00:00:00 2001
From: Kent Dong <ch3cho@qq.com>
Date: Wed, 24 Jul 2024 16:18:00 +0800
Subject: [PATCH] fix: Bypass the response body processing in ai-proxy if it is
 returned internally (#1149)

---
 plugins/wasm-go/extensions/ai-proxy/main.go   |  6 ++++
 .../wasm-go/extensions/frontend-gray/go.mod   |  2 +-
 .../wasm-go/extensions/frontend-gray/go.sum   |  1 +
 .../wasm-go/pkg/wrapper/response_wrapper.go   | 28 +++++++++++++++++++
 4 files changed, 36 insertions(+), 1 deletion(-)
 create mode 100644 plugins/wasm-go/pkg/wrapper/response_wrapper.go

diff --git a/plugins/wasm-go/extensions/ai-proxy/main.go b/plugins/wasm-go/extensions/ai-proxy/main.go
index 9be121dcf..67ee77244 100644
--- a/plugins/wasm-go/extensions/ai-proxy/main.go
+++ b/plugins/wasm-go/extensions/ai-proxy/main.go
@@ -111,6 +111,12 @@ func onHttpRequestBody(ctx wrapper.HttpContext, pluginConfig config.PluginConfig
 }
 
 func onHttpResponseHeaders(ctx wrapper.HttpContext, pluginConfig config.PluginConfig, log wrapper.Log) types.Action {
+	if !wrapper.IsResponseFromUpstream() {
+		// Response is not coming from the upstream. Let it pass through.
+		ctx.DontReadResponseBody()
+		return types.ActionContinue
+	}
+
 	activeProvider := pluginConfig.GetProvider()
 
 	if activeProvider == nil {
diff --git a/plugins/wasm-go/extensions/frontend-gray/go.mod b/plugins/wasm-go/extensions/frontend-gray/go.mod
index 37cd97090..835da044f 100644
--- a/plugins/wasm-go/extensions/frontend-gray/go.mod
+++ b/plugins/wasm-go/extensions/frontend-gray/go.mod
@@ -6,7 +6,7 @@ replace github.com/alibaba/higress/plugins/wasm-go => ../..
 
 require (
 	github.com/alibaba/higress/plugins/wasm-go v0.0.0-20240531060402-2807ddfbb79e
-	github.com/higress-group/proxy-wasm-go-sdk v0.0.0-20240327114451-d6b7174a84fc
+	github.com/higress-group/proxy-wasm-go-sdk v0.0.0-20240711023527-ba358c48772f
 	github.com/stretchr/testify v1.8.4
 	github.com/tidwall/gjson v1.17.0
 )
diff --git a/plugins/wasm-go/extensions/frontend-gray/go.sum b/plugins/wasm-go/extensions/frontend-gray/go.sum
index d6289a648..11f68bcbf 100644
--- a/plugins/wasm-go/extensions/frontend-gray/go.sum
+++ b/plugins/wasm-go/extensions/frontend-gray/go.sum
@@ -12,6 +12,7 @@ github.com/higress-group/proxy-wasm-go-sdk v0.0.0-20240226065437-8f7a0b3c9071 h1
 github.com/higress-group/proxy-wasm-go-sdk v0.0.0-20240226065437-8f7a0b3c9071/go.mod h1:hNFjhrLUIq+kJ9bOcs8QtiplSQ61GZXtd2xHKx4BYRo=
 github.com/higress-group/proxy-wasm-go-sdk v0.0.0-20240327114451-d6b7174a84fc h1:t2AT8zb6N/59Y78lyRWedVoVWHNRSCBh0oWCC+bluTQ=
 github.com/higress-group/proxy-wasm-go-sdk v0.0.0-20240327114451-d6b7174a84fc/go.mod h1:hNFjhrLUIq+kJ9bOcs8QtiplSQ61GZXtd2xHKx4BYRo=
+github.com/higress-group/proxy-wasm-go-sdk v0.0.0-20240711023527-ba358c48772f/go.mod h1:hNFjhrLUIq+kJ9bOcs8QtiplSQ61GZXtd2xHKx4BYRo=
 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=
diff --git a/plugins/wasm-go/pkg/wrapper/response_wrapper.go b/plugins/wasm-go/pkg/wrapper/response_wrapper.go
new file mode 100644
index 000000000..3cd91daf3
--- /dev/null
+++ b/plugins/wasm-go/pkg/wrapper/response_wrapper.go
@@ -0,0 +1,28 @@
+// Copyright (c) 2022 Alibaba Group Holding Ltd.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package wrapper
+
+import (
+	"github.com/higress-group/proxy-wasm-go-sdk/proxywasm"
+)
+
+func IsResponseFromUpstream() bool {
+	if codeDetails, err := proxywasm.GetProperty([]string{"response", "code_details"}); err == nil {
+		return string(codeDetails) == "via_upstream"
+	} else {
+		proxywasm.LogErrorf("get response code details failed: %v", err)
+		return false
+	}
+}