实现UA名单功能

pkg/serverconfigs/http_web_config.go

@@ -29,6 +29,7 @@ type HTTPWebConfig struct {
 	RewriteRules       []*HTTPRewriteRule                  `yaml:"rewriteRules" json:"rewriteRules"`             // 重写规则
 	FastcgiRef         *HTTPFastcgiRef                     `yaml:"fastcgiRef" json:"fastcgiRef"`                 // Fastcgi引用
 	FastcgiList        []*HTTPFastcgiConfig                `yaml:"fastcgiList" json:"fastcgiList"`               // Fastcgi配置
+	UserAgent          *UserAgentConfig                    `yaml:"userAgent" json:"userAgent"`                   // UserAgent配置
 
 	RequestHeaderPolicyRef  *shared.HTTPHeaderPolicyRef `yaml:"requestHeaderPolicyRef" json:"requestHeaderPolicyRef"`   // 请求Header
 	RequestHeaderPolicy     *shared.HTTPHeaderPolicy    `yaml:"requestHeaderPolicy" json:"requestHeaderPolicy"`         // 请求Header策略
@@ -283,6 +284,14 @@ func (this *HTTPWebConfig) Init() error {
 		}
 	}
 
+	// user agent
+	if this.UserAgent != nil {
+		err := this.UserAgent.Init()
+		if err != nil {
+			return err
+		}
+	}
+
 	return nil
 }
 

pkg/serverconfigs/user_agent_config.go

@@ -0,0 +1,117 @@
+// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn .
+
+package serverconfigs
+
+import (
+	"net/http"
+	"regexp"
+	"strings"
+)
+
+type UserAgentAction = string
+
+const (
+	UserAgentActionAllow UserAgentAction = "allow"
+	UserAgentActionDeny  UserAgentAction = "deny"
+)
+
+type UserAgentKeyword struct {
+	keyword string
+	reg     *regexp.Regexp
+}
+
+func NewUserAgentKeyword(keyword string) *UserAgentKeyword {
+	var reg *regexp.Regexp
+	if strings.Contains(keyword, "*") {
+		var pieces = strings.Split(keyword, "*")
+		for index, piece := range pieces {
+			pieces[index] = regexp.QuoteMeta(piece)
+		}
+		var newKeyword = strings.Join(pieces, ".*")
+		r, _ := regexp.Compile("(?i)" + newKeyword)
+		reg = r
+	} else {
+		r, _ := regexp.Compile("(?i)" + regexp.QuoteMeta(keyword))
+		reg = r
+	}
+
+	return &UserAgentKeyword{
+		keyword: keyword,
+		reg:     reg,
+	}
+}
+
+type UserAgentFilter struct {
+	Keywords []string        `yaml:"keywords" json:"keywords"` // 关键词
+	Action   UserAgentAction `yaml:"action" json:"action"`     // 动作
+
+	keywords []*UserAgentKeyword
+}
+
+func (this *UserAgentFilter) Init() error {
+	this.keywords = []*UserAgentKeyword{}
+	for _, keyword := range this.Keywords {
+		this.keywords = append(this.keywords, NewUserAgentKeyword(keyword))
+	}
+
+	return nil
+}
+
+func (this *UserAgentFilter) Match(userAgent string) bool {
+	if len(this.Keywords) == 0 {
+		return false
+	}
+
+	for _, keyword := range this.keywords {
+		if len(keyword.keyword) == 0 {
+			// 空白
+			if len(userAgent) == 0 {
+				return true
+			}
+		} else if keyword.reg != nil {
+			if keyword.reg.MatchString(userAgent) {
+				return true
+			}
+		} else if strings.Contains(userAgent, keyword.keyword) {
+			return true
+		}
+	}
+
+	return false
+}
+
+type UserAgentConfig struct {
+	IsPrior bool               `yaml:"isPrior" json:"isPrior"`
+	IsOn    bool               `yaml:"isOn" json:"isOn"`
+	Filters []*UserAgentFilter `yaml:"filters" json:"filters"`
+}
+
+func NewUserAgentConfig() *UserAgentConfig {
+	return &UserAgentConfig{}
+}
+
+func (this *UserAgentConfig) Init() error {
+	for _, filter := range this.Filters {
+		err := filter.Init()
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func (this *UserAgentConfig) AllowRequest(req *http.Request) bool {
+	if len(this.Filters) == 0 {
+		return true
+	}
+
+	var userAgent = req.UserAgent()
+
+	for _, filter := range this.Filters {
+		if filter.Match(userAgent) {
+			return filter.Action == UserAgentActionAllow
+		}
+	}
+
+	return true
+}

pkg/serverconfigs/user_agent_config_test.go

@@ -0,0 +1,71 @@
+// 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"
+	"net/http"
+	"testing"
+)
+
+func TestUserAgentConfig_AllowRequest(t *testing.T) {
+	var a = assert.NewAssertion(t)
+
+	var config = serverconfigs.NewUserAgentConfig()
+	config.Filters = append(config.Filters, &serverconfigs.UserAgentFilter{
+		Keywords: []string{"Chrome", "Google*Bot", "Opera*a*c", "|(*"},
+		Action:   serverconfigs.UserAgentActionAllow,
+	})
+	config.Filters = append(config.Filters, &serverconfigs.UserAgentFilter{
+		Keywords: []string{"Google*a*c"},
+		Action:   serverconfigs.UserAgentActionDeny,
+	})
+	config.Filters = append(config.Filters, &serverconfigs.UserAgentFilter{
+		Keywords: []string{""},
+		Action:   serverconfigs.UserAgentActionDeny,
+	})
+	config.Filters = append(config.Filters, &serverconfigs.UserAgentFilter{
+		Keywords: []string{"mozilla", "firefox"},
+		Action:   serverconfigs.UserAgentActionDeny,
+	})
+	err := config.Init()
+	if err != nil {
+		t.Fatal(err)
+	}
+	{
+		req, err := http.NewRequest(http.MethodGet, "/", nil)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		{
+			req.Header.Set("User-Agent", "")
+			a.IsFalse(config.AllowRequest(req))
+		}
+		{
+			req.Header.Set("User-Agent", "chrome")
+			a.IsTrue(config.AllowRequest(req))
+		}
+		{
+			req.Header.Set("User-Agent", "mozilla")
+			a.IsFalse(config.AllowRequest(req))
+		}
+		{
+			req.Header.Set("User-Agent", "Firefox")
+			a.IsFalse(config.AllowRequest(req))
+		}
+		{
+			req.Header.Set("User-Agent", "Google Bot")
+			a.IsTrue(config.AllowRequest(req))
+		}
+		{
+			req.Header.Set("User-Agent", "opera abc")
+			a.IsTrue(config.AllowRequest(req))
+		}
+		{
+			req.Header.Set("User-Agent", "google abc")
+			a.IsFalse(config.AllowRequest(req))
+		}
+	}
+}