From 45d08955c48d18452eb959a81987013f42eed3eb Mon Sep 17 00:00:00 2001
From: GoEdgeLab <goedgecloud@gmail.com>
Date: Sun, 24 Dec 2023 10:03:24 +0800
Subject: [PATCH] =?UTF-8?q?WAF=E5=8F=82=E6=95=B0=E4=B8=AD=E5=A2=9E?=
 =?UTF-8?q?=E5=8A=A0=E2=80=9C=E8=AF=B7=E6=B1=82=E6=9D=A5=E6=BA=90=E2=80=9D?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/waf/checkpoints/request_referer.go   |  2 +-
 .../waf/checkpoints/request_referer_origin.go | 44 +++++++++++++++++++
 .../request_referer_origin_test.go            | 38 ++++++++++++++++
 internal/waf/checkpoints/utils.go             | 10 ++++-
 4 files changed, 92 insertions(+), 2 deletions(-)
 create mode 100644 internal/waf/checkpoints/request_referer_origin.go
 create mode 100644 internal/waf/checkpoints/request_referer_origin_test.go

diff --git a/internal/waf/checkpoints/request_referer.go b/internal/waf/checkpoints/request_referer.go
index 93b2c21..9ad6790 100644
--- a/internal/waf/checkpoints/request_referer.go
+++ b/internal/waf/checkpoints/request_referer.go
@@ -23,5 +23,5 @@ func (this *RequestRefererCheckpoint) ResponseValue(req requests.Request, resp *
 }
 
 func (this *RequestRefererCheckpoint) CacheLife() utils.CacheLife {
-	return utils.CacheShortLife
+	return utils.CacheMiddleLife
 }
diff --git a/internal/waf/checkpoints/request_referer_origin.go b/internal/waf/checkpoints/request_referer_origin.go
new file mode 100644
index 0000000..8008298
--- /dev/null
+++ b/internal/waf/checkpoints/request_referer_origin.go
@@ -0,0 +1,44 @@
+package checkpoints
+
+import (
+	"github.com/TeaOSLab/EdgeNode/internal/waf/requests"
+	"github.com/TeaOSLab/EdgeNode/internal/waf/utils"
+	"github.com/iwind/TeaGo/maps"
+)
+
+type RequestRefererOriginCheckpoint struct {
+	Checkpoint
+}
+
+func (this *RequestRefererOriginCheckpoint) RequestValue(req requests.Request, param string, options maps.Map, ruleId int64) (value any, hasRequestBody bool, sysErr error, userErr error) {
+	var s []string
+
+	var referer = req.WAFRaw().Referer()
+	if len(referer) > 0 {
+		s = append(s, referer)
+	}
+
+	var origin = req.WAFRaw().Header.Get("Origin")
+	if len(origin) > 0 {
+		s = append(s, origin)
+	}
+
+	if len(s) > 0 {
+		value = s
+	} else {
+		value = ""
+	}
+
+	return
+}
+
+func (this *RequestRefererOriginCheckpoint) ResponseValue(req requests.Request, resp *requests.Response, param string, options maps.Map, ruleId int64) (value any, hasRequestBody bool, sysErr error, userErr error) {
+	if this.IsRequest() {
+		return this.RequestValue(req, param, options, ruleId)
+	}
+	return
+}
+
+func (this *RequestRefererOriginCheckpoint) CacheLife() utils.CacheLife {
+	return utils.CacheMiddleLife
+}
diff --git a/internal/waf/checkpoints/request_referer_origin_test.go b/internal/waf/checkpoints/request_referer_origin_test.go
new file mode 100644
index 0000000..d99f0ed
--- /dev/null
+++ b/internal/waf/checkpoints/request_referer_origin_test.go
@@ -0,0 +1,38 @@
+package checkpoints_test
+
+import (
+	"github.com/TeaOSLab/EdgeNode/internal/waf/checkpoints"
+	"github.com/TeaOSLab/EdgeNode/internal/waf/requests"
+	"net/http"
+	"testing"
+)
+
+func TestRequestRefererOriginCheckpoint_RequestValue(t *testing.T) {
+	rawReq, err := http.NewRequest(http.MethodGet, "https://example.com", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	var req = requests.NewTestRequest(rawReq)
+
+	var checkpoint = &checkpoints.RequestRefererOriginCheckpoint{}
+
+	{
+		t.Log(checkpoint.RequestValue(req, "", nil, 0))
+	}
+
+	{
+		rawReq.Header.Set("Referer", "https://example.com/hello.yaml")
+		t.Log(checkpoint.RequestValue(req, "", nil, 0))
+	}
+
+	{
+		rawReq.Header.Set("Origin", "https://example.com/world.yaml")
+		t.Log(checkpoint.RequestValue(req, "", nil, 0))
+	}
+
+	{
+		rawReq.Header.Del("Referer")
+		rawReq.Header.Set("Origin", "https://example.com/world.yaml")
+		t.Log(checkpoint.RequestValue(req, "", nil, 0))
+	}
+}
diff --git a/internal/waf/checkpoints/utils.go b/internal/waf/checkpoints/utils.go
index 62a4b6a..f65301c 100644
--- a/internal/waf/checkpoints/utils.go
+++ b/internal/waf/checkpoints/utils.go
@@ -163,7 +163,15 @@ var AllCheckpoints = []*CheckpointDefinition{
 		Priority:    100,
 	},
 	{
-		Name:        "请求来源URL",
+		Name:        "请求来源",
+		Prefix:      "refererOrigin",
+		Description: "请求报头中的Referer或Origin值",
+		HasParams:   false,
+		Instance:    new(RequestRefererOriginCheckpoint),
+		Priority:    100,
+	},
+	{
+		Name:        "请求来源Referer",
 		Prefix:      "referer",
 		Description: "请求Header中的Referer值",
 		HasParams:   false,