实现重写规则管理

pkg/serverconfigs/http_rewrite_ref.go

@@ -0,0 +1,7 @@
+package serverconfigs
+
+// 重写规则的引用
+type HTTPRewriteRef struct {
+	IsOn          bool  `yaml:"isOn" json:"isOn"`                   // 是否启用
+	RewriteRuleId int64 `yaml:"rewriteRuleId" json:"rewriteRuleId"` // 规则ID
+}

pkg/serverconfigs/http_rewrite_rule.go

@@ -0,0 +1,126 @@
+package serverconfigs
+
+import (
+	"fmt"
+	"github.com/TeaOSLab/EdgeCommon/pkg/configutils"
+	"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs/shared"
+	"regexp"
+)
+
+type HTTPRewriteMode = string
+
+const (
+	HTTPRewriteTargetProxy = 1
+	HTTPRewriteTargetURL   = 2
+)
+
+const (
+	HTTPRewriteModeRedirect HTTPRewriteMode = "redirect" // 跳转
+	HTTPRewriteModeProxy    HTTPRewriteMode = "proxy"    // 代理
+)
+
+// 重写规则定义
+//
+// 参考
+// - http://nginx.org/en/docs/http/ngx_http_rewrite_module.html
+// - https://httpd.apache.org/docs/current/mod/mod_rewrite.html
+// - https://httpd.apache.org/docs/2.4/rewrite/flags.html
+type HTTPRewriteRule struct {
+	Id   int64 `yaml:"id" json:"id"`     // ID
+	IsOn bool  `yaml:"isOn" json:"isOn"` // 是否开启
+
+	// 开启的条件
+	// 语法为：cond param operator value 比如：
+	// - cond ${status} gte 200
+	// - cond ${arg.name} eq lily
+	// - cond ${requestPath} regexp .*\.png
+	CondGroups []*shared.HTTPRequestCondGroup `yaml:"condGroups" json:"condGroups"` // 匹配条件 TODO
+
+	// 规则
+	// 语法为：pattern regexp 比如：
+	// - pattern ^/article/(\d+).html
+	Pattern string `yaml:"pattern" json:"pattern"`
+
+	// 模式
+	Mode           HTTPRewriteMode `yaml:"mode" json:"mode"`
+	RedirectStatus int             `yaml:"redirectStatus" json:"redirectStatus"` // 跳转的状态码
+	ProxyHost      string          `yaml:"proxyHost" json:"proxyHost"`           // 代理模式下的Host
+
+	// TODO 实现对其他代理服务的引用
+
+	// 要替换成的URL
+	// 支持反向引用：${0}, ${1}, ...，也支持?P<NAME>语法
+	// - 如果以 proxy:// 开头，表示目标为代理，首先会尝试作为代理ID请求，如果找不到，会尝试作为代理Host请求
+	Replace string `yaml:"replace" json:"replace"`
+
+	IsBreak bool `yaml:"isBreak" json:"isBreak"` // 终止向下解析
+
+	reg *regexp.Regexp
+}
+
+// 校验
+func (this *HTTPRewriteRule) Init() error {
+	reg, err := regexp.Compile(this.Pattern)
+	if err != nil {
+		return err
+	}
+	this.reg = reg
+
+	// 校验条件
+	for _, condGroup := range this.CondGroups {
+		err := condGroup.Init()
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+// 对某个请求执行规则
+func (this *HTTPRewriteRule) Match(requestPath string, formatter func(source string) string) (replace string, varMapping map[string]string, matched bool) {
+	if this.reg == nil {
+		return "", nil, false
+	}
+
+	matches := this.reg.FindStringSubmatch(requestPath)
+	if len(matches) == 0 {
+		return "", nil, false
+	}
+
+	// 判断条件
+	if len(this.CondGroups) > 0 {
+		for _, cond := range this.CondGroups {
+			if !cond.Match(formatter) {
+				return "", nil, false
+			}
+		}
+	}
+
+	varMapping = map[string]string{}
+	subNames := this.reg.SubexpNames()
+	for index, match := range matches {
+		varMapping[fmt.Sprintf("%d", index)] = match
+		subName := subNames[index]
+		if len(subName) > 0 {
+			varMapping[subName] = match
+		}
+	}
+
+	replace = configutils.ParseVariables(this.Replace, func(varName string) string {
+		v, ok := varMapping[varName]
+		if ok {
+			return v
+		}
+		return "${" + varName + "}"
+	})
+
+	replace = formatter(replace)
+
+	return replace, varMapping, true
+}
+
+// 判断是否是外部URL
+func (this *HTTPRewriteRule) IsExternalURL(url string) bool {
+	return shared.RegexpExternalURL.MatchString(url)
+}

pkg/serverconfigs/http_rewrite_rule_test.go

@@ -0,0 +1,180 @@
+package serverconfigs
+
+import (
+	"github.com/iwind/TeaGo/assert"
+	"github.com/iwind/TeaGo/utils/string"
+	"sync"
+	"testing"
+	"time"
+)
+
+func TestHTTPRewriteRule(t *testing.T) {
+	a := assert.NewAssertion(t).Quiet()
+
+	rule := HTTPRewriteRule{
+		Pattern: "/(hello)/(world)",
+		Replace: "/${1}/${2}",
+	}
+	a.IsNil(rule.Init())
+
+	{
+		_, _, b := rule.Match("/hello/worl", func(source string) string {
+			return source
+		})
+		a.IsFalse(b)
+		a.Log("proxy:", rule.TargetProxy())
+		a.Log("url:", rule.TargetURL())
+	}
+
+	{
+		_, _, b := rule.Match("/hello/world", func(source string) string {
+			return source
+		})
+		a.IsTrue(b)
+		a.Log("proxy:", rule.TargetProxy())
+		a.Log("url:", rule.TargetURL())
+	}
+
+	{
+		r := HTTPRewriteRule{}
+		r.Replace = "http://127.0.0.1${0}"
+		r.Pattern = ".*"
+		err := r.Init()
+		if err != nil {
+			t.Fatal(err)
+		}
+		u, _, b := r.Match("/hello", func(source string) string {
+			return source
+		})
+		a.Log(b)
+		a.Log(u)
+	}
+}
+
+func TestRewriteRule_NamedMatch(t *testing.T) {
+	r := &HTTPRewriteRule{}
+	r.Replace = "http://127.0.0.1/${1}/${last}/${ni}"
+	r.Pattern = "/(\\w+)/(?P<last>\\w+)/(?P<ni>\\w+)"
+	err := r.Init()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	before := time.Now()
+	count := 100
+	for i := 0; i < count; i++ {
+		s, _, b := r.Match("/hello/world/ni", func(source string) string {
+			return source
+		})
+		if i == 0 {
+			if b {
+				t.Log("matched:", s)
+			} else {
+				t.Log("not matched")
+			}
+		}
+	}
+	t.Log(float64(count) / (time.Since(before).Seconds()))
+}
+
+func TestRewriteRule_NamedMatchConcurrent(t *testing.T) {
+	r := &HTTPRewriteRule{}
+	r.Replace = "http://127.0.0.1/${1}/${last}/${ni}"
+	r.Pattern = "/(\\w+)/(?P<last>\\w+)/(?P<ni>\\w+)"
+	err := r.Init()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	threads := 1000
+	count := 1000
+	wg := sync.WaitGroup{}
+	wg.Add(threads * count)
+	fails := 0
+	var locker sync.Mutex
+	for i := 0; i < threads; i++ {
+		go func() {
+			for j := 0; j < count; j++ {
+				func() {
+					defer wg.Done()
+
+					var randomString = stringutil.Rand(16)
+
+					replace, _, b := r.Match("/hello/world/"+randomString, func(source string) string {
+						return source
+					})
+					if !b {
+						locker.Lock()
+						fails++
+						locker.Unlock()
+					} else if replace != "http://127.0.0.1/hello/world/"+randomString {
+						locker.Lock()
+						fails++
+						locker.Unlock()
+					}
+				}()
+			}
+		}()
+	}
+	wg.Wait()
+	if fails > 0 {
+		t.Log("fail")
+	} else {
+		t.Log("success")
+	}
+}
+
+func TestRewriteRule_CaseInsensitive(t *testing.T) {
+	a := assert.NewAssertion(t)
+
+	r := &HTTPRewriteRule{}
+	r.Replace = "http://127.0.0.1${0}"
+	r.Pattern = "(?i)/index.php"
+	err := r.Init()
+	if err != nil {
+		t.Fatal(err)
+	}
+	_, _, ok := r.Match("/index.php", func(source string) string {
+		return source
+	})
+	a.IsTrue(ok)
+
+	_, _, ok = r.Match("/INDEX.php", func(source string) string {
+		return source
+	})
+	a.IsTrue(ok)
+}
+
+func TestRewriteRule_Slashes(t *testing.T) {
+	a := assert.NewAssertion(t)
+
+	r := &HTTPRewriteRule{}
+	r.Replace = "http://127.0.0.1/${0}"
+	r.Pattern = "(?i)/index.php"
+	err := r.Init()
+	if err != nil {
+		t.Fatal(err)
+	}
+	replace, _, ok := r.Match("/index.php", func(source string) string {
+		return source
+	})
+	a.IsTrue(ok)
+	t.Log(replace)
+}
+
+func TestRewriteRuleProxy(t *testing.T) {
+	a := assert.NewAssertion(t).Quiet()
+
+	rule := &HTTPRewriteRule{
+		Pattern: "/(hello)/(world)",
+		Replace: "proxy://lb001/${1}/${2}",
+	}
+	a.IsNil(rule.Init())
+
+	replace, _, b := rule.Match("/hello/world", func(source string) string {
+		return source
+	})
+	a.IsTrue(b)
+	a.IsTrue(rule.TargetProxy() == "lb001")
+	a.IsTrue(replace == "/hello/world")
+}

pkg/serverconfigs/http_web_config.go

@@ -21,6 +21,8 @@ type HTTPWebConfig struct {
 	FirewallRef        *HTTPFirewallRef           `yaml:"firewallRef" json:"firewallRef"`               // 防火墙设置
 	WebsocketRef       *HTTPWebsocketRef          `yaml:"websocketRef" json:"websocketRef"`             // Websocket应用配置
 	Websocket          *HTTPWebsocketConfig       `yaml:"websocket" json:"websocket"`                   // Websocket配置
+	RewriteRefs        []*HTTPRewriteRef          `yaml:"rewriteRefs" json:"rewriteRefs"`               // 重写规则配置
+	RewriteRules       []*HTTPRewriteRule         `yaml:"rewriteRules" json:"rewriteRules"`             // 重写规则
 
 	RequestHeaderPolicyRef  *shared.HTTPHeaderPolicyRef `yaml:"requestHeaderPolicyRef" json:"requestHeaderPolicyRef"`   // 请求Header
 	RequestHeaderPolicy     *shared.HTTPHeaderPolicy    `yaml:"requestHeaderPolicy" json:"requestHeaderPolicy"`         // 请求Header策略
@@ -184,6 +186,14 @@ func (this *HTTPWebConfig) Init() error {
 		}
 	}
 
+	// rewrite rules
+	for _, rewriteRule := range this.RewriteRules {
+		err := rewriteRule.Init()
+		if err != nil {
+			return err
+		}
+	}
+
 	return nil
 }