From c40a7707c1edb94b5c8bca7a9f42b4deafe3a48f Mon Sep 17 00:00:00 2001
From: GoEdgeLab <goedgecloud@gmail.com>
Date: Mon, 24 Oct 2022 10:21:13 +0800
Subject: [PATCH] =?UTF-8?q?=E9=98=B2=E7=9B=97=E9=93=BE=E5=8A=9F=E8=83=BD?=
 =?UTF-8?q?=E5=A2=9E=E5=8A=A0=E7=A6=81=E6=AD=A2=E7=9A=84=E6=9D=A5=E6=BA=90?=
 =?UTF-8?q?=E5=9F=9F=E5=90=8D?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 pkg/serverconfigs/http_websocket_config.go |  2 +-
 pkg/serverconfigs/referers_config.go       | 13 +++++-
 pkg/serverconfigs/referers_config_test.go  | 48 ++++++++++++++++++++++
 3 files changed, 61 insertions(+), 2 deletions(-)
 create mode 100644 pkg/serverconfigs/referers_config_test.go

diff --git a/pkg/serverconfigs/http_websocket_config.go b/pkg/serverconfigs/http_websocket_config.go
index b090aa1..6a4ae1b 100644
--- a/pkg/serverconfigs/http_websocket_config.go
+++ b/pkg/serverconfigs/http_websocket_config.go
@@ -6,7 +6,7 @@ import (
 	"time"
 )
 
-// websocket设置
+// HTTPWebsocketConfig websocket设置
 type HTTPWebsocketConfig struct {
 	Id   int64 `yaml:"id" json:"id"`     // ID
 	IsOn bool  `yaml:"isOn" json:"isOn"` // 是否开启
diff --git a/pkg/serverconfigs/referers_config.go b/pkg/serverconfigs/referers_config.go
index 08f53ae..e801271 100644
--- a/pkg/serverconfigs/referers_config.go
+++ b/pkg/serverconfigs/referers_config.go
@@ -11,6 +11,7 @@ type ReferersConfig struct {
 	AllowEmpty      bool     `yaml:"allowEmpty" json:"allowEmpty"`           // 来源域名允许为空
 	AllowSameDomain bool     `yaml:"allowSameDomain" json:"allowSameDomain"` // 允许来源域名和当前访问的域名一致，相当于在站内访问
 	AllowDomains    []string `yaml:"allowDomains" json:"allowDomains"`       // 允许的来源域名列表
+	DenyDomains     []string `yaml:"denyDomains" json:"denyDomains"`         // 禁止的来源域名列表
 }
 
 func (this *ReferersConfig) Init() error {
@@ -30,8 +31,18 @@ func (this *ReferersConfig) MatchDomain(requestDomain string, refererDomain stri
 	}
 
 	if len(this.AllowDomains) == 0 {
+		if len(this.DenyDomains) > 0 {
+			return !configutils.MatchDomains(this.DenyDomains, refererDomain)
+		}
 		return false
 	}
 
-	return configutils.MatchDomains(this.AllowDomains, refererDomain)
+	if configutils.MatchDomains(this.AllowDomains, refererDomain) {
+		if len(this.DenyDomains) > 0 && configutils.MatchDomains(this.DenyDomains, refererDomain) {
+			return false
+		}
+		return true
+	}
+
+	return false
 }
diff --git a/pkg/serverconfigs/referers_config_test.go b/pkg/serverconfigs/referers_config_test.go
new file mode 100644
index 0000000..2a667dc
--- /dev/null
+++ b/pkg/serverconfigs/referers_config_test.go
@@ -0,0 +1,48 @@
+// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn .
+
+package serverconfigs_test
+
+import (
+	"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs"
+	"github.com/iwind/TeaGo/assert"
+	"testing"
+)
+
+func TestReferersConfig_MatchDomain(t *testing.T) {
+	var a = assert.NewAssertion(t)
+
+	{
+		var config = &serverconfigs.ReferersConfig{
+			IsOn:         true,
+			AllowDomains: []string{},
+			DenyDomains:  []string{"a.com", "b.com"},
+		}
+		a.IsTrue(config.MatchDomain("example.com", "c.com"))
+		a.IsTrue(config.MatchDomain("example.com", "d.com"))
+		a.IsFalse(config.MatchDomain("example.com", "a.com"))
+	}
+
+	{
+		var config = &serverconfigs.ReferersConfig{
+			IsOn:         true,
+			AllowDomains: []string{"c.com", "e.com"},
+			DenyDomains:  []string{"a.com", "b.com", "e.com"},
+		}
+		a.IsTrue(config.MatchDomain("example.com", "c.com"))
+		a.IsFalse(config.MatchDomain("example.com", "d.com"))
+		a.IsFalse(config.MatchDomain("example.com", "e.com"))
+		a.IsFalse(config.MatchDomain("example.com", "a.com"))
+	}
+
+	{
+		var config = &serverconfigs.ReferersConfig{
+			IsOn:         true,
+			AllowDomains: []string{"c.com", "e.com"},
+			DenyDomains:  []string{},
+		}
+		a.IsTrue(config.MatchDomain("example.com", "c.com"))
+		a.IsFalse(config.MatchDomain("example.com", "d.com"))
+		a.IsTrue(config.MatchDomain("example.com", "e.com"))
+		a.IsFalse(config.MatchDomain("example.com", "a.com"))
+	}
+}