WAF策略中增加验证码相关定制设置

internal/nodes/http_request_waf.go

@@ -194,7 +194,7 @@ func (this *HTTPRequest) checkWAFRequest(firewallPolicy *firewallconfigs.HTTPFir
 	}
 
 	// 规则测试
-	w := sharedWAFManager.FindWAF(firewallPolicy.Id)
+	w := waf.SharedWAFManager.FindWAF(firewallPolicy.Id)
 	if w == nil {
 		return
 	}
@@ -261,7 +261,7 @@ func (this *HTTPRequest) checkWAFResponse(firewallPolicy *firewallconfigs.HTTPFi
 		return
 	}
 
-	w := sharedWAFManager.FindWAF(firewallPolicy.Id)
+	w := waf.SharedWAFManager.FindWAF(firewallPolicy.Id)
 	if w == nil {
 		return
 	}

internal/nodes/node.go

@@ -21,6 +21,7 @@ import (
 	"github.com/TeaOSLab/EdgeNode/internal/stats"
 	"github.com/TeaOSLab/EdgeNode/internal/trackers"
 	"github.com/TeaOSLab/EdgeNode/internal/utils"
+	"github.com/TeaOSLab/EdgeNode/internal/waf"
 	"github.com/andybalholm/brotli"
 	"github.com/iwind/TeaGo/Tea"
 	"github.com/iwind/TeaGo/lists"
@@ -865,7 +866,7 @@ func (this *Node) onReload(config *nodeconfigs.NodeConfig) {
 	}
 
 	// WAF策略
-	sharedWAFManager.UpdatePolicies(config.FindAllFirewallPolicies())
+	waf.SharedWAFManager.UpdatePolicies(config.FindAllFirewallPolicies())
 	iplibrary.SharedActionManager.UpdateActions(config.FirewallActions)
 
 	// 统计指标

internal/waf/action_allow.go

@@ -6,6 +6,7 @@ import (
 )
 
 type AllowAction struct {
+	BaseAction
 }
 
 func (this *AllowAction) Init(waf *WAF) error {

internal/waf/action_base.go

@@ -7,6 +7,17 @@ import (
 )
 
 type BaseAction struct {
+	currentActionId int64
+}
+
+// ActionId 读取ActionId
+func (this *BaseAction) ActionId() int64 {
+	return this.currentActionId
+}
+
+// SetActionId 设置Id
+func (this *BaseAction) SetActionId(actionId int64) {
+	this.currentActionId = actionId
 }
 
 // CloseConn 关闭连接

internal/waf/action_block.go

@@ -20,6 +20,8 @@ var urlPrefixReg = regexp.MustCompile("^(?i)(http|https)://")
 var httpClient = utils.SharedHttpClient(5 * time.Second)
 
 type BlockAction struct {
+	BaseAction
+
 	StatusCode int    `yaml:"statusCode" json:"statusCode"`
 	Body       string `yaml:"body" json:"body"` // supports HTML
 	URL        string `yaml:"url" json:"url"`

internal/waf/action_captcha.go

@@ -18,16 +18,71 @@ const (
 )
 
 type CaptchaAction struct {
+	BaseAction
+
 	Life              int32 `yaml:"life" json:"life"`
 	MaxFails          int   `yaml:"maxFails" json:"maxFails"`                   // 最大失败次数
 	FailBlockTimeout  int   `yaml:"failBlockTimeout" json:"failBlockTimeout"`   // 失败拦截时间
+	FailBlockScopeAll bool  `yaml:"failBlockScopeAll" json:"failBlockScopeAll"` // 是否全局有效
 
-	Language       string `yaml:"language" json:"language"`             // 语言，zh-CN, en-US ...
+	CountLetters int8 `yaml:"countLetters" json:"countLetters"`
+
+	UIIsOn          bool   `yaml:"uiIsOn" json:"uiIsOn"`                   // 是否使用自定义UI
+	UITitle         string `yaml:"uiTitle" json:"uiTitle"`                 // 消息标题
+	UIPrompt        string `yaml:"uiPrompt" json:"uiPrompt"`               // 消息提示
+	UIButtonTitle   string `yaml:"uiButtonTitle" json:"uiButtonTitle"`     // 按钮标题
+	UIShowRequestId bool   `yaml:"uiShowRequestId" json:"uiShowRequestId"` // 是否显示请求ID
+	UICss           string `yaml:"uiCss" json:"uiCss"`                     // CSS样式
+	UIFooter        string `yaml:"uiFooter" json:"uiFooter"`               // 页脚
+	UIBody          string `yaml:"uiBody" json:"uiBody"`                   // 内容轮廓
+
+	Lang           string `yaml:"lang" json:"lang"`                     // 语言，zh-CN, en-US ...
 	AddToWhiteList bool   `yaml:"addToWhiteList" json:"addToWhiteList"` // 是否加入到白名单
 	Scope          string `yaml:"scope" json:"scope"`
 }
 
 func (this *CaptchaAction) Init(waf *WAF) error {
+	if waf.DefaultCaptchaAction != nil {
+		if this.Life <= 0 {
+			this.Life = waf.DefaultCaptchaAction.Life
+		}
+		if this.MaxFails <= 0 {
+			this.MaxFails = waf.DefaultCaptchaAction.MaxFails
+		}
+		if this.FailBlockTimeout <= 0 {
+			this.FailBlockTimeout = waf.DefaultCaptchaAction.FailBlockTimeout
+		}
+		this.FailBlockScopeAll = waf.DefaultCaptchaAction.FailBlockScopeAll
+
+		if this.CountLetters <= 0 {
+			this.CountLetters = waf.DefaultCaptchaAction.CountLetters
+		}
+
+		this.UIIsOn = waf.DefaultCaptchaAction.UIIsOn
+		if len(this.UITitle) == 0 {
+			this.UITitle = waf.DefaultCaptchaAction.UITitle
+		}
+		if len(this.UIPrompt) == 0 {
+			this.UIPrompt = waf.DefaultCaptchaAction.UIPrompt
+		}
+		if len(this.UIButtonTitle) == 0 {
+			this.UIButtonTitle = waf.DefaultCaptchaAction.UIButtonTitle
+		}
+		this.UIShowRequestId = waf.DefaultCaptchaAction.UIShowRequestId
+		if len(this.UICss) == 0 {
+			this.UICss = waf.DefaultCaptchaAction.UICss
+		}
+		if len(this.UIFooter) == 0 {
+			this.UIFooter = waf.DefaultCaptchaAction.UIFooter
+		}
+		if len(this.UIBody) == 0 {
+			this.UIBody = waf.DefaultCaptchaAction.UIBody
+		}
+		if len(this.Lang) == 0 {
+			this.Lang = waf.DefaultCaptchaAction.Lang
+		}
+	}
+
 	return nil
 }
 
@@ -49,7 +104,7 @@ func (this *CaptchaAction) Perform(waf *WAF, group *RuleGroup, set *RuleSet, req
 		return true
 	}
 
-	refURL := request.WAFRaw().URL.String()
+	var refURL = request.WAFRaw().URL.String()
 
 	// 覆盖配置
 	if strings.HasPrefix(refURL, CaptchaPath) {
@@ -63,10 +118,8 @@ func (this *CaptchaAction) Perform(waf *WAF, group *RuleGroup, set *RuleSet, req
 	}
 
 	var captchaConfig = maps.Map{
-		"action":           this,
+		"actionId":  this.ActionId(),
 		"timestamp": time.Now().Unix(),
-		"maxFails":         this.MaxFails,
-		"failBlockTimeout": this.FailBlockTimeout,
 		"url":       refURL,
 		"policyId":  waf.Id,
 		"groupId":   group.Id,

internal/waf/action_go_group.go

@@ -8,6 +8,8 @@ import (
 )
 
 type GoGroupAction struct {
+	BaseAction
+
 	GroupId string `yaml:"groupId" json:"groupId"`
 }
 

internal/waf/action_go_set.go

@@ -8,6 +8,8 @@ import (
 )
 
 type GoSetAction struct {
+	BaseAction
+
 	GroupId string `yaml:"groupId" json:"groupId"`
 	SetId   string `yaml:"setId" json:"setId"`
 }

internal/waf/action_interface.go

@@ -11,6 +11,12 @@ type ActionInterface interface {
 	// Init 初始化
 	Init(waf *WAF) error
 
+	// ActionId 读取ActionId
+	ActionId() int64
+
+	// SetActionId 设置ID
+	SetActionId(id int64)
+
 	// Code 代号
 	Code() string
 
@@ -20,6 +26,6 @@ type ActionInterface interface {
 	// WillChange determine if the action will change the request
 	WillChange() bool
 
-	// Perform perform the action
+	// Perform the action
 	Perform(waf *WAF, group *RuleGroup, set *RuleSet, request requests.Request, writer http.ResponseWriter) (allow bool)
 }

internal/waf/action_log.go

@@ -6,6 +6,7 @@ import (
 )
 
 type LogAction struct {
+	BaseAction
 }
 
 func (this *LogAction) Init(waf *WAF) error {

internal/waf/action_notify.go

@@ -50,6 +50,7 @@ func init() {
 }
 
 type NotifyAction struct {
+	BaseAction
 }
 
 func (this *NotifyAction) Init(waf *WAF) error {
@@ -69,7 +70,7 @@ func (this *NotifyAction) WillChange() bool {
 	return false
 }
 
-// Perform perform the action
+// Perform the action
 func (this *NotifyAction) Perform(waf *WAF, group *RuleGroup, set *RuleSet, request requests.Request, writer http.ResponseWriter) (allow bool) {
 	select {
 	case notifyChan <- &notifyTask{

internal/waf/action_tag.go

@@ -6,6 +6,8 @@ import (
 )
 
 type TagAction struct {
+	BaseAction
+
 	Tags []string `yaml:"tags" json:"tags"`
 }
 

internal/waf/action_utils.go

@@ -5,15 +5,19 @@ import (
 	"github.com/TeaOSLab/EdgeNode/internal/remotelogs"
 	"github.com/iwind/TeaGo/maps"
 	"reflect"
+	"sync/atomic"
 )
 
+var seedActionId int64 = 1
+
 func FindActionInstance(action ActionString, options maps.Map) ActionInterface {
 	for _, def := range AllActions {
 		if def.Code == action {
 			if def.Type != nil {
 				// create new instance
-				ptrValue := reflect.New(def.Type)
+				var ptrValue = reflect.New(def.Type)
-				instance := ptrValue.Interface().(ActionInterface)
+				var instance = ptrValue.Interface().(ActionInterface)
+				instance.SetActionId(atomic.AddInt64(&seedActionId, 1))
 
 				if len(options) > 0 {
 					optionsJSON, err := json.Marshal(options)

internal/waf/captcha_validator.go

@@ -6,7 +6,6 @@ import (
 	"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs/firewallconfigs"
 	"github.com/TeaOSLab/EdgeNode/internal/ttlcache"
 	"github.com/TeaOSLab/EdgeNode/internal/utils"
-	"github.com/TeaOSLab/EdgeNode/internal/utils/jsonutils"
 	"github.com/TeaOSLab/EdgeNode/internal/waf/requests"
 	"github.com/dchest/captcha"
 	"github.com/iwind/TeaGo/logs"
@@ -26,8 +25,8 @@ func NewCaptchaValidator() *CaptchaValidator {
 	return &CaptchaValidator{}
 }
 
-func (this *CaptchaValidator) Run(request requests.Request, writer http.ResponseWriter) {
+func (this *CaptchaValidator) Run(req requests.Request, writer http.ResponseWriter) {
-	var info = request.WAFRaw().URL.Query().Get("info")
+	var info = req.WAFRaw().URL.Query().Get("info")
 	if len(info) == 0 {
 		writer.WriteHeader(http.StatusBadRequest)
 		_, _ = writer.Write([]byte("invalid request"))
@@ -41,35 +40,48 @@ func (this *CaptchaValidator) Run(request requests.Request, writer http.Response
 
 	var timestamp = m.GetInt64("timestamp")
 	if timestamp < time.Now().Unix()-600 { // 10分钟之后信息过期
-		http.Redirect(writer, request.WAFRaw(), m.GetString("url"), http.StatusTemporaryRedirect)
+		http.Redirect(writer, req.WAFRaw(), m.GetString("url"), http.StatusTemporaryRedirect)
-		return
-	}
-
-	var actionConfig = &CaptchaAction{}
-	err = jsonutils.MapToObject(m.GetMap("action"), actionConfig)
-	if err != nil {
-		http.Redirect(writer, request.WAFRaw(), m.GetString("url"), http.StatusTemporaryRedirect)
 		return
 	}
 
+	var actionId = m.GetInt64("actionId")
 	var setId = m.GetInt64("setId")
 	var originURL = m.GetString("url")
-	var maxFails = m.GetInt("maxFails")
-	var failBlockTimeout = m.GetInt("failBlockTimeout")
 	var policyId = m.GetInt64("policyId")
 	var groupId = m.GetInt64("groupId")
-	if request.WAFRaw().Method == http.MethodPost && len(request.WAFRaw().FormValue("GOEDGE_WAF_CAPTCHA_ID")) > 0 {
+
-		this.validate(actionConfig, maxFails, failBlockTimeout, policyId, groupId, setId, originURL, request, writer)
+	var waf = SharedWAFManager.FindWAF(policyId)
+	if waf == nil {
+		http.Redirect(writer, req.WAFRaw(), originURL, http.StatusTemporaryRedirect)
+		return
+	}
+	var actionConfig = waf.FindAction(actionId)
+	if actionConfig == nil {
+		http.Redirect(writer, req.WAFRaw(), originURL, http.StatusTemporaryRedirect)
+		return
+	}
+	captchaActionConfig, ok := actionConfig.(*CaptchaAction)
+	if !ok {
+		http.Redirect(writer, req.WAFRaw(), originURL, http.StatusTemporaryRedirect)
+		return
+	}
+
+	if req.WAFRaw().Method == http.MethodPost && len(req.WAFRaw().FormValue("GOEDGE_WAF_CAPTCHA_ID")) > 0 {
+		this.validate(captchaActionConfig, policyId, groupId, setId, originURL, req, writer)
 	} else {
 		// 增加计数
-		this.IncreaseFails(request, maxFails, failBlockTimeout, policyId, groupId, setId)
+		this.IncreaseFails(req, captchaActionConfig, policyId, groupId, setId)
-		this.show(actionConfig, request, writer)
+		this.show(captchaActionConfig, req, writer)
 	}
 }
 
-func (this *CaptchaValidator) show(actionConfig *CaptchaAction, request requests.Request, writer http.ResponseWriter) {
+func (this *CaptchaValidator) show(actionConfig *CaptchaAction, req requests.Request, writer http.ResponseWriter) {
 	// show captcha
-	var captchaId = captcha.NewLen(6)
+	var countLetters = 6
+	if actionConfig.CountLetters > 0 && actionConfig.CountLetters <= 10 {
+		countLetters = int(actionConfig.CountLetters)
+	}
+	var captchaId = captcha.NewLen(countLetters)
 	var buf = bytes.NewBuffer([]byte{})
 	err := captcha.WriteImage(buf, captchaId, 200, 100)
 	if err != nil {
@@ -77,9 +89,9 @@ func (this *CaptchaValidator) show(actionConfig *CaptchaAction, request requests
 		return
 	}
 
-	var lang = actionConfig.Language
+	var lang = actionConfig.Lang
 	if len(lang) == 0 {
-		acceptLanguage := request.WAFRaw().Header.Get("Accept-Language")
+		var acceptLanguage = req.WAFRaw().Header.Get("Accept-Language")
 		if len(acceptLanguage) > 0 {
 			langIndex := strings.Index(acceptLanguage, ",")
 			if langIndex > 0 {
@@ -114,12 +126,62 @@ func (this *CaptchaValidator) show(actionConfig *CaptchaAction, request requests
 		msgRequestId = "Request ID"
 	}
 
+	var msgCss = ""
+	var requestIdBox = `<address>` + msgRequestId + `: ` + req.Format("${requestId}") + `</address>`
+	var msgFooter = ""
+	var body = `<form method="POST">
+	<input type="hidden" name="GOEDGE_WAF_CAPTCHA_ID" value="` + captchaId + `"/>
+	<div class="ui-image">
+		<img src="data:image/png;base64, ` + base64.StdEncoding.EncodeToString(buf.Bytes()) + `"/>` + `
+	</div>
+	<div class="ui-input">
+		<p>` + msgPrompt + `</p>
+		<input type="text" name="GOEDGE_WAF_CAPTCHA_CODE" id="GOEDGE_WAF_CAPTCHA_CODE" maxlength="6" autocomplete="off" z-index="1" class="input"/>
+	</div>
+	<div class="ui-button">
+		<button type="submit" style="line-height:24px;margin-top:10px">` + msgButtonTitle + `</button>
+	</div>
+</form>
+` + requestIdBox + `
+` + msgFooter + ``
+
+	// 默认设置
+	if actionConfig.UIIsOn {
+		if len(actionConfig.UITitle) > 0 {
+			msgTitle = actionConfig.UITitle
+		}
+		if len(actionConfig.UIPrompt) > 0 {
+			msgPrompt = actionConfig.UIPrompt
+		}
+		if len(actionConfig.UIButtonTitle) > 0 {
+			msgButtonTitle = actionConfig.UIButtonTitle
+		}
+		if len(actionConfig.UICss) > 0 {
+			msgCss = actionConfig.UICss
+		}
+		if !actionConfig.UIShowRequestId {
+			requestIdBox = ""
+		}
+		if len(actionConfig.UIFooter) > 0 {
+			msgFooter = actionConfig.UIFooter
+		}
+		if len(actionConfig.UIBody) > 0 {
+			var index = strings.Index(actionConfig.UIBody, "${body}")
+			if index < 0 {
+				body = actionConfig.UIBody + body
+			} else {
+				body = actionConfig.UIBody[:index] + body + actionConfig.UIBody[index+7:] // 7是"${body}"的长度
+			}
+		}
+	}
+
 	writer.Header().Set("Content-Type", "text/html; charset=utf-8")
 	_, _ = writer.Write([]byte(`<!DOCTYPE html>
 <html>
 <head>
 	<title>` + msgTitle + `</title>
 	<meta name="viewport" content="width=device-width, initial-scale=1, user-scalable=0">
+	<meta charset="UTF-8"/>
 	<script type="text/javascript">
 	if (window.addEventListener != null) {
 		window.addEventListener("load", function () {
@@ -131,32 +193,22 @@ func (this *CaptchaValidator) show(actionConfig *CaptchaAction, request requests
 	form { width: 20em; margin: 0 auto; text-align: center; }
 	.input { font-size:16px;line-height:24px; letter-spacing: 15px; padding-left: 10px; width: 140px; }
 	address { margin-top: 1em; padding-top: 0.5em; border-top: 1px #ccc solid; text-align: center; }
+` + msgCss + `
 	</style>
 </head>
-<body>
+<body>` + body + `
-<form method="POST">
-	<input type="hidden" name="GOEDGE_WAF_CAPTCHA_ID" value="` + captchaId + `"/>
-	<img src="data:image/png;base64, ` + base64.StdEncoding.EncodeToString(buf.Bytes()) + `"/>` + `
-	<div>
-		<p>` + msgPrompt + `</p>
-		<input type="text" name="GOEDGE_WAF_CAPTCHA_CODE" id="GOEDGE_WAF_CAPTCHA_CODE" maxlength="6" autocomplete="off" z-index="1" class="input"/>
-	</div>
-	<div>
-		<button type="submit" style="line-height:24px;margin-top:10px">` + msgButtonTitle + `</button>
-	</div>
-</form>
-<address>` + msgRequestId + `: ` + request.Format("${requestId}") + `</address>
 </body>
 </html>`))
 }
 
-func (this *CaptchaValidator) validate(actionConfig *CaptchaAction, maxFails int, failBlockTimeout int, policyId int64, groupId int64, setId int64, originURL string, request requests.Request, writer http.ResponseWriter) (allow bool) {
+func (this *CaptchaValidator) validate(actionConfig *CaptchaAction, policyId int64, groupId int64, setId int64, originURL string, req requests.Request, writer http.ResponseWriter) (allow bool) {
-	var captchaId = request.WAFRaw().FormValue("GOEDGE_WAF_CAPTCHA_ID")
+
+	var captchaId = req.WAFRaw().FormValue("GOEDGE_WAF_CAPTCHA_ID")
 	if len(captchaId) > 0 {
-		var captchaCode = request.WAFRaw().FormValue("GOEDGE_WAF_CAPTCHA_CODE")
+		var captchaCode = req.WAFRaw().FormValue("GOEDGE_WAF_CAPTCHA_CODE")
 		if captcha.VerifyString(captchaId, captchaCode) {
 			// 清除计数
-			ttlcache.SharedCache.Delete("CAPTCHA:FAILS:" + request.WAFRemoteIP())
+			ttlcache.SharedCache.Delete(this.cacheKey(req))
 
 			var life = CaptchaSeconds
 			if actionConfig.Life > 0 {
@@ -164,18 +216,18 @@ func (this *CaptchaValidator) validate(actionConfig *CaptchaAction, maxFails int
 			}
 
 			// 加入到白名单
-			SharedIPWhiteList.RecordIP("set:"+strconv.FormatInt(setId, 10), actionConfig.Scope, request.WAFServerId(), request.WAFRemoteIP(), time.Now().Unix()+int64(life), policyId, false, groupId, setId, "")
+			SharedIPWhiteList.RecordIP("set:"+strconv.FormatInt(setId, 10), actionConfig.Scope, req.WAFServerId(), req.WAFRemoteIP(), time.Now().Unix()+int64(life), policyId, false, groupId, setId, "")
 
-			http.Redirect(writer, request.WAFRaw(), originURL, http.StatusSeeOther)
+			http.Redirect(writer, req.WAFRaw(), originURL, http.StatusSeeOther)
 
 			return false
 		} else {
 			// 增加计数
-			if !this.IncreaseFails(request, maxFails, failBlockTimeout, policyId, groupId, setId) {
+			if !this.IncreaseFails(req, actionConfig, policyId, groupId, setId) {
 				return false
 			}
 
-			http.Redirect(writer, request.WAFRaw(), request.WAFRaw().URL.String(), http.StatusSeeOther)
+			http.Redirect(writer, req.WAFRaw(), req.WAFRaw().URL.String(), http.StatusSeeOther)
 		}
 	}
 
@@ -183,17 +235,28 @@ func (this *CaptchaValidator) validate(actionConfig *CaptchaAction, maxFails int
 }
 
 // IncreaseFails 增加失败次数，以便后续操作
-func (this *CaptchaValidator) IncreaseFails(request requests.Request, maxFails int, failBlockTimeout int, policyId int64, groupId int64, setId int64) (goNext bool) {
+func (this *CaptchaValidator) IncreaseFails(req requests.Request, actionConfig *CaptchaAction, policyId int64, groupId int64, setId int64) (goNext bool) {
+	var maxFails = actionConfig.MaxFails
+	var failBlockTimeout = actionConfig.FailBlockTimeout
 	if maxFails > 0 && failBlockTimeout > 0 {
 		// 加上展示的计数
 		maxFails *= 2
 
-		var countFails = ttlcache.SharedCache.IncreaseInt64("CAPTCHA:FAILS:"+request.WAFRemoteIP(), 1, time.Now().Unix()+300, true)
+		var countFails = ttlcache.SharedCache.IncreaseInt64(this.cacheKey(req), 1, time.Now().Unix()+300, true)
 		if int(countFails) >= maxFails {
 			var useLocalFirewall = false
-			SharedIPBlackList.RecordIP(IPTypeAll, firewallconfigs.FirewallScopeService, request.WAFServerId(), request.WAFRemoteIP(), time.Now().Unix()+int64(failBlockTimeout), policyId, useLocalFirewall, groupId, setId, "CAPTCHA验证连续失败")
+
+			if actionConfig.FailBlockScopeAll {
+				useLocalFirewall = true
+			}
+
+			SharedIPBlackList.RecordIP(IPTypeAll, firewallconfigs.FirewallScopeService, req.WAFServerId(), req.WAFRemoteIP(), time.Now().Unix()+int64(failBlockTimeout), policyId, useLocalFirewall, groupId, setId, "CAPTCHA验证连续失败")
 			return false
 		}
 	}
 	return true
 }
+
+func (this *CaptchaValidator) cacheKey(req requests.Request) string {
+	return "CAPTCHA:FAILS:" + req.WAFRemoteIP() + ":" + types.String(req.WAFServerId())
+}

internal/waf/rule_set.go

@@ -66,7 +66,7 @@ func (this *RuleSet) Init(waf *WAF) error {
 	// action instances
 	this.actionInstances = []ActionInterface{}
 	for _, action := range this.Actions {
-		instance := FindActionInstance(action.Code, action.Options)
+		var instance = FindActionInstance(action.Code, action.Options)
 		if instance == nil {
 			remotelogs.Error("WAF_RULE_SET", "can not find instance for action '"+action.Code+"'")
 			continue
@@ -79,6 +79,7 @@ func (this *RuleSet) Init(waf *WAF) error {
 		}
 
 		this.actionInstances = append(this.actionInstances, instance)
+		waf.AddAction(instance)
 	}
 
 	// sort actions

internal/waf/waf.go

@@ -27,11 +27,13 @@ type WAF struct {
 	SYNFlood         *firewallconfigs.SYNFloodConfig `yaml:"synFlood" json:"synFlood"`
 
 	DefaultBlockAction   *BlockAction
+	DefaultCaptchaAction *CaptchaAction
 
 	hasInboundRules  bool
 	hasOutboundRules bool
 
 	checkpointsMap map[string]checkpoints.CheckpointInterface // prefix => checkpoint
+	actionMap      map[int64]ActionInterface                  // actionId => ActionInterface
 }
 
 func NewWAF() *WAF {
@@ -74,6 +76,9 @@ func (this *WAF) Init() (resultErrors []error) {
 		this.checkpointsMap[def.Prefix] = instance
 	}
 
+	// action map
+	this.actionMap = map[int64]ActionInterface{}
+
 	// rules
 	this.hasInboundRules = len(this.Inbound) > 0
 	this.hasOutboundRules = len(this.Outbound) > 0
@@ -324,8 +329,16 @@ func (this *WAF) ContainsGroupCode(code string) bool {
 	return false
 }
 
+func (this *WAF) AddAction(action ActionInterface) {
+	this.actionMap[action.ActionId()] = action
+}
+
+func (this *WAF) FindAction(actionId int64) ActionInterface {
+	return this.actionMap[actionId]
+}
+
 func (this *WAF) Copy() *WAF {
-	waf := &WAF{
+	var waf = &WAF{
 		Id:       this.Id,
 		IsOn:     this.IsOn,
 		Name:     this.Name,

internal/waf/waf_manager.go

@@ -1,26 +1,25 @@
-package nodes
+package waf
 
 import (
 	"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs/firewallconfigs"
 	"github.com/TeaOSLab/EdgeNode/internal/errors"
 	"github.com/TeaOSLab/EdgeNode/internal/remotelogs"
-	"github.com/TeaOSLab/EdgeNode/internal/waf"
 	"strconv"
 	"sync"
 )
 
-var sharedWAFManager = NewWAFManager()
+var SharedWAFManager = NewWAFManager()
 
 // WAFManager WAF管理器
 type WAFManager struct {
-	mapping map[int64]*waf.WAF // policyId => WAF
+	mapping map[int64]*WAF // policyId => WAF
 	locker  sync.RWMutex
 }
 
 // NewWAFManager 获取新对象
 func NewWAFManager() *WAFManager {
 	return &WAFManager{
-		mapping: map[int64]*waf.WAF{},
+		mapping: map[int64]*WAF{},
 	}
 }
 
@@ -29,9 +28,9 @@ func (this *WAFManager) UpdatePolicies(policies []*firewallconfigs.HTTPFirewallP
 	this.locker.Lock()
 	defer this.locker.Unlock()
 
-	m := map[int64]*waf.WAF{}
+	m := map[int64]*WAF{}
 	for _, p := range policies {
-		w, err := this.convertWAF(p)
+		w, err := this.ConvertWAF(p)
 		if w != nil {
 			m[p.Id] = w
 		}
@@ -44,22 +43,22 @@ func (this *WAFManager) UpdatePolicies(policies []*firewallconfigs.HTTPFirewallP
 }
 
 // FindWAF 查找WAF
-func (this *WAFManager) FindWAF(policyId int64) *waf.WAF {
+func (this *WAFManager) FindWAF(policyId int64) *WAF {
 	this.locker.RLock()
 	w, _ := this.mapping[policyId]
 	this.locker.RUnlock()
 	return w
 }
 
-// 将Policy转换为WAF
+// ConvertWAF 将Policy转换为WAF
-func (this *WAFManager) convertWAF(policy *firewallconfigs.HTTPFirewallPolicy) (*waf.WAF, error) {
+func (this *WAFManager) ConvertWAF(policy *firewallconfigs.HTTPFirewallPolicy) (*WAF, error) {
 	if policy == nil {
 		return nil, errors.New("policy should not be nil")
 	}
 	if len(policy.Mode) == 0 {
 		policy.Mode = firewallconfigs.FirewallModeDefend
 	}
-	w := &waf.WAF{
+	var w = &WAF{
 		Id:               policy.Id,
 		IsOn:             policy.IsOn,
 		Name:             policy.Name,
@@ -71,7 +70,7 @@ func (this *WAFManager) convertWAF(policy *firewallconfigs.HTTPFirewallPolicy) (
 	// inbound
 	if policy.Inbound != nil && policy.Inbound.IsOn {
 		for _, group := range policy.Inbound.Groups {
-			g := &waf.RuleGroup{
+			g := &RuleGroup{
 				Id:          group.Id,
 				IsOn:        group.IsOn,
 				Name:        group.Name,
@@ -82,7 +81,7 @@ func (this *WAFManager) convertWAF(policy *firewallconfigs.HTTPFirewallPolicy) (
 
 			// rule sets
 			for _, set := range group.Sets {
-				s := &waf.RuleSet{
+				s := &RuleSet{
 					Id:          set.Id,
 					Code:        set.Code,
 					IsOn:        set.IsOn,
@@ -97,10 +96,10 @@ func (this *WAFManager) convertWAF(policy *firewallconfigs.HTTPFirewallPolicy) (
 
 				// rules
 				for _, rule := range set.Rules {
-					r := &waf.Rule{
+					r := &Rule{
 						Description:       rule.Description,
 						Param:             rule.Param,
-						ParamFilters:      []*waf.ParamFilter{},
+						ParamFilters:      []*ParamFilter{},
 						Operator:          rule.Operator,
 						Value:             rule.Value,
 						IsCaseInsensitive: rule.IsCaseInsensitive,
@@ -108,7 +107,7 @@ func (this *WAFManager) convertWAF(policy *firewallconfigs.HTTPFirewallPolicy) (
 					}
 
 					for _, paramFilter := range rule.ParamFilters {
-						r.ParamFilters = append(r.ParamFilters, &waf.ParamFilter{
+						r.ParamFilters = append(r.ParamFilters, &ParamFilter{
 							Code:    paramFilter.Code,
 							Options: paramFilter.Options,
 						})
@@ -127,7 +126,7 @@ func (this *WAFManager) convertWAF(policy *firewallconfigs.HTTPFirewallPolicy) (
 	// outbound
 	if policy.Outbound != nil && policy.Outbound.IsOn {
 		for _, group := range policy.Outbound.Groups {
-			g := &waf.RuleGroup{
+			g := &RuleGroup{
 				Id:          group.Id,
 				IsOn:        group.IsOn,
 				Name:        group.Name,
@@ -138,7 +137,7 @@ func (this *WAFManager) convertWAF(policy *firewallconfigs.HTTPFirewallPolicy) (
 
 			// rule sets
 			for _, set := range group.Sets {
-				s := &waf.RuleSet{
+				s := &RuleSet{
 					Id:          set.Id,
 					Code:        set.Code,
 					IsOn:        set.IsOn,
@@ -154,7 +153,7 @@ func (this *WAFManager) convertWAF(policy *firewallconfigs.HTTPFirewallPolicy) (
 
 				// rules
 				for _, rule := range set.Rules {
-					r := &waf.Rule{
+					r := &Rule{
 						Description:       rule.Description,
 						Param:             rule.Param,
 						Operator:          rule.Operator,
@@ -172,9 +171,9 @@ func (this *WAFManager) convertWAF(policy *firewallconfigs.HTTPFirewallPolicy) (
 		}
 	}
 
-	// action
+	// block action
 	if policy.BlockOptions != nil {
-		w.DefaultBlockAction = &waf.BlockAction{
+		w.DefaultBlockAction = &BlockAction{
 			StatusCode: policy.BlockOptions.StatusCode,
 			Body:       policy.BlockOptions.Body,
 			URL:        policy.BlockOptions.URL,
@@ -182,6 +181,26 @@ func (this *WAFManager) convertWAF(policy *firewallconfigs.HTTPFirewallPolicy) (
 		}
 	}
 
+	// captcha action
+	if policy.CaptchaOptions != nil {
+		w.DefaultCaptchaAction = &CaptchaAction{
+			Life:              policy.CaptchaOptions.Life,
+			MaxFails:          policy.CaptchaOptions.MaxFails,
+			FailBlockTimeout:  policy.CaptchaOptions.FailBlockTimeout,
+			FailBlockScopeAll: policy.CaptchaOptions.FailBlockScopeAll,
+			CountLetters:      policy.CaptchaOptions.CountLetters,
+			UIIsOn:            policy.CaptchaOptions.UIIsOn,
+			UITitle:           policy.CaptchaOptions.UITitle,
+			UIPrompt:          policy.CaptchaOptions.UIPrompt,
+			UIButtonTitle:     policy.CaptchaOptions.UIButtonTitle,
+			UIShowRequestId:   policy.CaptchaOptions.UIShowRequestId,
+			UICss:             policy.CaptchaOptions.UICss,
+			UIFooter:          policy.CaptchaOptions.UIFooter,
+			UIBody:            policy.CaptchaOptions.UIBody,
+			Lang:              policy.CaptchaOptions.Lang,
+		}
+	}
+
 	errorList := w.Init()
 	if len(errorList) > 0 {
 		return w, errorList[0]

internal/waf/waf_manager_test.go

@@ -1,7 +1,8 @@
-package nodes
+package waf_test
 
 import (
 	"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs/firewallconfigs"
+	"github.com/TeaOSLab/EdgeNode/internal/waf"
 	"github.com/iwind/TeaGo/logs"
 	"testing"
 )
@@ -35,7 +36,7 @@ func TestWAFManager_convert(t *testing.T) {
 			},
 		},
 	}
-	w, err := sharedWAFManager.convertWAF(p)
+	w, err := waf.SharedWAFManager.ConvertWAF(p)
 	if err != nil {
 		t.Fatal(err)
 	}