From 2793e0de89f28d65ce5b7037e8eb9fc858a57cd9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=88=98=E7=A5=A5=E8=B6=85?= <iwind.liu@gmail.com>
Date: Tue, 19 Oct 2021 11:38:46 +0800
Subject: [PATCH] =?UTF-8?q?=E5=A2=9E=E5=8A=A0=E9=98=B2=E7=9B=97=E9=93=BE?=
 =?UTF-8?q?=E8=A7=84=E5=88=99=E5=8F=82=E6=95=B0?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../waf/checkpoints/request_referer_block.go  | 66 +++++++++++++++++++
 internal/waf/checkpoints/utils.go             |  7 ++
 internal/waf/rule.go                          |  2 +-
 3 files changed, 74 insertions(+), 1 deletion(-)
 create mode 100644 internal/waf/checkpoints/request_referer_block.go

diff --git a/internal/waf/checkpoints/request_referer_block.go b/internal/waf/checkpoints/request_referer_block.go
new file mode 100644
index 0000000..b03f1b9
--- /dev/null
+++ b/internal/waf/checkpoints/request_referer_block.go
@@ -0,0 +1,66 @@
+// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package checkpoints
+
+import (
+	"github.com/TeaOSLab/EdgeCommon/pkg/configutils"
+	"github.com/TeaOSLab/EdgeNode/internal/waf/requests"
+	"github.com/iwind/TeaGo/maps"
+	"github.com/iwind/TeaGo/types"
+	"net/url"
+)
+
+// RequestRefererBlockCheckpoint 防盗链
+type RequestRefererBlockCheckpoint struct {
+	Checkpoint
+}
+
+// RequestValue 计算checkpoint值
+// 选项：allowEmpty, allowSameDomain, allowDomains
+func (this *RequestRefererBlockCheckpoint) RequestValue(req requests.Request, param string, options maps.Map) (value interface{}, sysErr error, userErr error) {
+	var referer = req.WAFRaw().Referer()
+
+	if len(referer) == 0 {
+		if options.GetBool("allowEmpty") {
+			value = 1
+			return
+		}
+		value = 0
+		return
+	}
+
+	u, err := url.Parse(referer)
+	if err != nil {
+		value = 0
+		return
+	}
+	var host = u.Host
+
+	if options.GetBool("allowSameDomain") && host == req.WAFRaw().Host {
+		value = 1
+		return
+	}
+
+	var domains = options.GetSlice("allowDomains")
+	var domainStrings = []string{}
+	for _, domain := range domains {
+		domainStrings = append(domainStrings, types.String(domain))
+	}
+
+	if len(domainStrings) == 0 {
+		value = 0
+		return
+	}
+
+	if configutils.MatchDomains(domainStrings, host) {
+		value = 1
+	} else {
+		value = 0
+	}
+
+	return
+}
+
+func (this *RequestRefererBlockCheckpoint) ResponseValue(req requests.Request, resp *requests.Response, param string, options maps.Map) (value interface{}, sysErr error, userErr error) {
+	return
+}
diff --git a/internal/waf/checkpoints/utils.go b/internal/waf/checkpoints/utils.go
index 9054af0..945955f 100644
--- a/internal/waf/checkpoints/utils.go
+++ b/internal/waf/checkpoints/utils.go
@@ -198,6 +198,13 @@ var AllCheckpoints = []*CheckpointDefinition{
 		HasParams:   true,
 		Instance:    new(CC2Checkpoint),
 	},
+	{
+		Name:        "防盗链",
+		Prefix:      "refererBlock",
+		Description: "阻止一些域名访问引用本站资源",
+		HasParams:   true,
+		Instance:    new(RequestRefererBlockCheckpoint),
+	},
 	{
 		Name:        "通用响应Header长度限制",
 		Prefix:      "responseGeneralHeaderLength",
diff --git a/internal/waf/rule.go b/internal/waf/rule.go
index a9ad080..47a4c6b 100644
--- a/internal/waf/rule.go
+++ b/internal/waf/rule.go
@@ -23,7 +23,7 @@ import (
 
 var singleParamRegexp = regexp.MustCompile("^\\${[\\w.-]+}$")
 
-// rule
+// Rule
 type Rule struct {
 	Description       string                 `yaml:"description" json:"description"`
 	Param             string                 `yaml:"param" json:"param"` // such as ${arg.name} or ${args}, can be composite as ${arg.firstName}${arg.lastName}