WAF增加多个动作

internal/waf/rule_set.go

@@ -1,9 +1,13 @@
 package waf
 
 import (
+	"github.com/TeaOSLab/EdgeNode/internal/remotelogs"
 	"github.com/TeaOSLab/EdgeNode/internal/waf/requests"
+	"github.com/iwind/TeaGo/lists"
+	"github.com/iwind/TeaGo/logs"
 	"github.com/iwind/TeaGo/maps"
 	"github.com/iwind/TeaGo/utils/string"
+	"net/http"
 )
 
 type RuleConnector = string
@@ -14,16 +18,17 @@ const (
 )
 
 type RuleSet struct {
-	Id          string        `yaml:"id" json:"id"`
-	Code        string        `yaml:"code" json:"code"`
-	IsOn        bool          `yaml:"isOn" json:"isOn"`
-	Name        string        `yaml:"name" json:"name"`
-	Description string        `yaml:"description" json:"description"`
-	Rules       []*Rule       `yaml:"rules" json:"rules"`
-	Connector   RuleConnector `yaml:"connector" json:"connector"` // rules connector
+	Id          string          `yaml:"id" json:"id"`
+	Code        string          `yaml:"code" json:"code"`
+	IsOn        bool            `yaml:"isOn" json:"isOn"`
+	Name        string          `yaml:"name" json:"name"`
+	Description string          `yaml:"description" json:"description"`
+	Rules       []*Rule         `yaml:"rules" json:"rules"`
+	Connector   RuleConnector   `yaml:"connector" json:"connector"` // rules connector
+	Actions     []*ActionConfig `yaml:"actions" json:"actions"`
 
-	Action        ActionString `yaml:"action" json:"action"`
-	ActionOptions maps.Map     `yaml:"actionOptions" json:"actionOptions"` // TODO TO BE IMPLEMENTED
+	actionCodes     []string
+	actionInstances []ActionInterface
 
 	hasRules bool
 }
@@ -35,7 +40,7 @@ func NewRuleSet() *RuleSet {
 	}
 }
 
-func (this *RuleSet) Init() error {
+func (this *RuleSet) Init(waf *WAF) error {
 	this.hasRules = len(this.Rules) > 0
 	if this.hasRules {
 		for _, rule := range this.Rules {
@@ -45,6 +50,31 @@ func (this *RuleSet) Init() error {
 			}
 		}
 	}
+
+	// action codes
+	var actionCodes = []string{}
+	for _, action := range this.Actions {
+		if !lists.ContainsString(actionCodes, action.Code) {
+			actionCodes = append(actionCodes, action.Code)
+		}
+	}
+	this.actionCodes = actionCodes
+
+	// action instances
+	this.actionInstances = []ActionInterface{}
+	for _, action := range this.Actions {
+		instance := FindActionInstance(action.Code, action.Options)
+		if instance == nil {
+			remotelogs.Error("WAF_RULE_SET", "can not find instance for action '"+action.Code+"'")
+		} else {
+			this.actionInstances = append(this.actionInstances, instance)
+		}
+		err := instance.Init(waf)
+		if err != nil {
+			remotelogs.Error("WAF_RULE_SET", "init action '"+action.Code+"' failed: "+err.Error())
+		}
+	}
+
 	return nil
 }
 
@@ -52,7 +82,75 @@ func (this *RuleSet) AddRule(rule ...*Rule) {
 	this.Rules = append(this.Rules, rule...)
 }
 
-func (this *RuleSet) MatchRequest(req *requests.Request) (b bool, err error) {
+// AddAction 添加动作
+func (this *RuleSet) AddAction(code string, options maps.Map) {
+	if options == nil {
+		options = maps.Map{}
+	}
+	this.Actions = append(this.Actions, &ActionConfig{
+		Code:    code,
+		Options: options,
+	})
+}
+
+// HasSpecialActions 除了Allow之外是否还有别的动作
+func (this *RuleSet) HasSpecialActions() bool {
+	for _, action := range this.Actions {
+		if action.Code != ActionAllow {
+			return true
+		}
+	}
+	return false
+}
+
+// HasAttackActions 检查是否含有攻击防御动作
+func (this *RuleSet) HasAttackActions() bool {
+	for _, action := range this.actionInstances {
+		if action.IsAttack() {
+			return true
+		}
+	}
+	return false
+}
+
+func (this *RuleSet) ActionCodes() []string {
+	return this.actionCodes
+}
+
+func (this *RuleSet) PerformActions(waf *WAF, group *RuleGroup, req requests.Request, writer http.ResponseWriter) bool {
+	// 先执行allow
+	for _, instance := range this.actionInstances {
+		if !instance.WillChange() {
+			if waf.onActionCallback != nil {
+				goNext := waf.onActionCallback(instance)
+				if !goNext {
+					return false
+				}
+			}
+			logs.Printf("perform1: %#v", instance) // TODO
+			instance.Perform(waf, group, this, req, writer)
+		}
+	}
+
+	// 再执行block|verify
+	for _, instance := range this.actionInstances {
+		// 只执行第一个可能改变请求的动作，其余的都会被忽略
+		if instance.WillChange() {
+			if waf.onActionCallback != nil {
+				goNext := waf.onActionCallback(instance)
+				if !goNext {
+					return false
+				}
+			}
+			logs.Printf("perform2: %#v", instance) // TODO
+			return instance.Perform(waf, group, this, req, writer)
+		}
+	}
+
+	return true
+}
+
+func (this *RuleSet) MatchRequest(req requests.Request) (b bool, err error) {
 	if !this.hasRules {
 		return false, nil
 	}
@@ -93,7 +191,7 @@ func (this *RuleSet) MatchRequest(req *requests.Request) (b bool, err error) {
 	return
 }
 
-func (this *RuleSet) MatchResponse(req *requests.Request, resp *requests.Response) (b bool, err error) {
+func (this *RuleSet) MatchResponse(req requests.Request, resp *requests.Response) (b bool, err error) {
 	if !this.hasRules {
 		return false, nil
 	}