diff --git a/internal/nodes/http_request_websocket.go b/internal/nodes/http_request_websocket.go index 5bb4b25..91de38f 100644 --- a/internal/nodes/http_request_websocket.go +++ b/internal/nodes/http_request_websocket.go @@ -56,7 +56,7 @@ func (this *HTTPRequest) doWebsocket() { } clientConn, _, err := this.writer.Hijack() - if err != nil { + if err != nil || clientConn == nil { this.write50x(err, http.StatusInternalServerError) return } diff --git a/internal/waf/action_base.go b/internal/waf/action_base.go index b3cc227..93bf304 100644 --- a/internal/waf/action_base.go +++ b/internal/waf/action_base.go @@ -15,7 +15,7 @@ func (this *BaseAction) CloseConn(writer http.ResponseWriter) error { hijack, ok := writer.(http.Hijacker) if ok { conn, _, err := hijack.Hijack() - if err == nil { + if err == nil && conn != nil { return conn.Close() } } diff --git a/internal/waf/action_page.go b/internal/waf/action_page.go new file mode 100644 index 0000000..ab5b905 --- /dev/null +++ b/internal/waf/action_page.go @@ -0,0 +1,41 @@ +// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved. + +package waf + +import ( + "github.com/TeaOSLab/EdgeNode/internal/waf/requests" + "net/http" +) + +type PageAction struct { + BaseAction + + Status int `yaml:"status" json:"status"` + Body string `yaml:"body" json:"body"` +} + +func (this *PageAction) Init(waf *WAF) error { + return nil +} + +func (this *PageAction) Code() string { + return ActionPage +} + +func (this *PageAction) IsAttack() bool { + return false +} + +// WillChange determine if the action will change the request +func (this *PageAction) WillChange() bool { + return true +} + +// Perform the action +func (this *PageAction) Perform(waf *WAF, group *RuleGroup, set *RuleSet, request requests.Request, writer http.ResponseWriter) (allow bool) { + writer.Header().Set("Content-Type", "text/html; charset=utf-8") + writer.WriteHeader(this.Status) + _, _ = writer.Write([]byte(request.Format(this.Body))) + + return false +} diff --git a/internal/waf/action_types.go b/internal/waf/action_types.go index 3a9d512..75aeace 100644 --- a/internal/waf/action_types.go +++ b/internal/waf/action_types.go @@ -13,6 +13,7 @@ const ( ActionPost307 ActionString = "post_307" // 针对POST的307重定向认证 ActionRecordIP ActionString = "record_ip" // 记录IP ActionTag ActionString = "tag" // 标签 + ActionPage ActionString = "page" // 显示网页 ActionAllow ActionString = "allow" // allow ActionGoGroup ActionString = "go_group" // go to next rule group ActionGoSet ActionString = "go_set" // go to next rule set @@ -73,6 +74,12 @@ var AllActions = []*ActionDefinition{ Instance: new(TagAction), Type: reflect.TypeOf(new(TagAction)).Elem(), }, + { + Name: "显示页面", + Code: ActionPage, + Instance: new(PageAction), + Type: reflect.TypeOf(new(PageAction)).Elem(), + }, { Name: "跳到下一个规则分组", Code: ActionGoGroup, diff --git a/internal/waf/rule_set.go b/internal/waf/rule_set.go index 2e2274a..35169d3 100644 --- a/internal/waf/rule_set.go +++ b/internal/waf/rule_set.go @@ -67,13 +67,16 @@ func (this *RuleSet) Init(waf *WAF) error { 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) + continue } + err := instance.Init(waf) if err != nil { remotelogs.Error("WAF_RULE_SET", "init action '"+action.Code+"' failed: "+err.Error()) + continue } + + this.actionInstances = append(this.actionInstances, instance) } // sort actions