From 42c9c3d824ae0dba7b82697612cdeda36be6020e Mon Sep 17 00:00:00 2001
From: rinfx <893383980@qq.com>
Date: Wed, 15 May 2024 20:34:47 +0800
Subject: [PATCH] waf skip body when protocol is grpc, websocket or sse (#943)

---
 plugins/wasm-go/extensions/waf/go.sum         |  3 +-
 .../extensions/waf/wasmplugin/plugin.go       | 17 ++++++++++
 .../extensions/waf/wasmplugin/utils.go        | 31 +++++++++++++++++++
 3 files changed, 49 insertions(+), 2 deletions(-)

diff --git a/plugins/wasm-go/extensions/waf/go.sum b/plugins/wasm-go/extensions/waf/go.sum
index 82e6729c6..4d318e425 100644
--- a/plugins/wasm-go/extensions/waf/go.sum
+++ b/plugins/wasm-go/extensions/waf/go.sum
@@ -10,8 +10,7 @@ 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-20240226065437-8f7a0b3c9071 h1:STb5rOHRZOzoiAa+gTz2LFqO1nYj7U/1eIVUJJadU4A=
-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-20240318034951-d5306e367c43 h1:dCw7F/9ciw4NZN7w68wQRaygZ2zGOWMTIEoRvP1tlWs=
 github.com/higress-group/proxy-wasm-go-sdk v0.0.0-20240318034951-d5306e367c43/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=
diff --git a/plugins/wasm-go/extensions/waf/wasmplugin/plugin.go b/plugins/wasm-go/extensions/waf/wasmplugin/plugin.go
index b74dbf0e8..f4beb38b0 100644
--- a/plugins/wasm-go/extensions/waf/wasmplugin/plugin.go
+++ b/plugins/wasm-go/extensions/waf/wasmplugin/plugin.go
@@ -69,6 +69,15 @@ func parseConfig(json gjson.Result, config *WafConfig, log wrapper.Log) error {
 }
 
 func onHttpRequestHeaders(ctx wrapper.HttpContext, config WafConfig, log wrapper.Log) types.Action {
+	ctx.SetContext("skipwaf", false)
+
+	if ignoreBody() {
+		ctx.DontReadRequestBody()
+		ctx.DontReadResponseBody()
+		ctx.SetContext("skipwaf", true)
+		return types.ActionContinue
+	}
+
 	ctx.SetContext("interruptionHandled", false)
 	ctx.SetContext("processedRequestBody", false)
 	ctx.SetContext("processedResponseBody", false)
@@ -192,6 +201,10 @@ func onHttpRequestBody(ctx wrapper.HttpContext, config WafConfig, body []byte, l
 }
 
 func onHttpResponseHeaders(ctx wrapper.HttpContext, config WafConfig, log wrapper.Log) types.Action {
+	if ctx.GetContext("skipwaf").(bool) {
+		return types.ActionContinue
+	}
+
 	if ctx.GetContext("interruptionHandled").(bool) {
 		return types.ActionContinue
 	}
@@ -306,6 +319,10 @@ func onHttpResponseBody(ctx wrapper.HttpContext, config WafConfig, body []byte,
 }
 
 func onHttpStreamDone(ctx wrapper.HttpContext, config WafConfig, log wrapper.Log) {
+	if ctx.GetContext("skipwaf").(bool) {
+		return
+	}
+
 	tx := ctx.GetContext("tx").(ctypes.Transaction)
 
 	if !tx.IsRuleEngineOff() {
diff --git a/plugins/wasm-go/extensions/waf/wasmplugin/utils.go b/plugins/wasm-go/extensions/waf/wasmplugin/utils.go
index d8066def5..d34977b69 100644
--- a/plugins/wasm-go/extensions/waf/wasmplugin/utils.go
+++ b/plugins/wasm-go/extensions/waf/wasmplugin/utils.go
@@ -138,3 +138,34 @@ func logError(error ctypes.MatchedRule) {
 		proxywasm.LogDebug(msg)
 	}
 }
+
+func isWebSocketRequest() bool {
+	if value, err := proxywasm.GetHttpRequestHeader("Upgrade"); err == nil {
+		if value == "websocket" {
+			return true
+		}
+	}
+	return false
+}
+
+func isSSERequest() bool {
+	if value, err := proxywasm.GetHttpRequestHeader("Accept"); err == nil {
+		if value == "text/event-stream" {
+			return true
+		}
+	}
+	return false
+}
+
+func isGrpcRequest() bool {
+	if value, err := proxywasm.GetHttpRequestHeader("Content-Type"); err == nil {
+		if value == "application/grpc" {
+			return true
+		}
+	}
+	return false
+}
+
+func ignoreBody() bool {
+	return isWebSocketRequest() || isSSERequest() || isGrpcRequest()
+}