From 86db3dfc49b0df448e98d68e2bb0e451cbdb70af Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=88=98=E7=A5=A5=E8=B6=85?= <iwind.liu@gmail.com>
Date: Tue, 12 Oct 2021 09:06:28 +0800
Subject: [PATCH] =?UTF-8?q?WAF=E5=8A=A8=E4=BD=9Crecord=5Fip=E8=BF=94?=
 =?UTF-8?q?=E5=9B=9E403/=E4=BC=98=E5=8C=96=E5=85=B3=E9=97=AD=E8=BF=9E?=
 =?UTF-8?q?=E6=8E=A5=E6=96=B9=E6=B3=95?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/waf/action_base.go      | 15 ---------------
 internal/waf/action_get_302.go   |  2 +-
 internal/waf/action_post_307.go  |  2 +-
 internal/waf/action_record_ip.go |  4 +++-
 4 files changed, 5 insertions(+), 18 deletions(-)

diff --git a/internal/waf/action_base.go b/internal/waf/action_base.go
index e0e6bec..d3a0026 100644
--- a/internal/waf/action_base.go
+++ b/internal/waf/action_base.go
@@ -2,20 +2,5 @@
 
 package waf
 
-import "net/http"
-
 type BaseAction struct {
 }
-
-// CloseConn 关闭连接
-func (this *BaseAction) CloseConn(writer http.ResponseWriter) error {
-	// 断开连接
-	hijack, ok := writer.(http.Hijacker)
-	if ok {
-		conn, _, err := hijack.Hijack()
-		if err == nil {
-			return conn.Close()
-		}
-	}
-	return nil
-}
diff --git a/internal/waf/action_get_302.go b/internal/waf/action_get_302.go
index 8cf37d6..9dffd05 100644
--- a/internal/waf/action_get_302.go
+++ b/internal/waf/action_get_302.go
@@ -66,7 +66,7 @@ func (this *Get302Action) Perform(waf *WAF, group *RuleGroup, set *RuleSet, requ
 
 	// 关闭连接
 	if request.WAFRaw().ProtoMajor == 1 {
-		_ = this.CloseConn(writer)
+		request.WAFClose()
 	}
 
 	return true
diff --git a/internal/waf/action_post_307.go b/internal/waf/action_post_307.go
index cfc0bb5..4eb4393 100644
--- a/internal/waf/action_post_307.go
+++ b/internal/waf/action_post_307.go
@@ -82,7 +82,7 @@ func (this *Post307Action) Perform(waf *WAF, group *RuleGroup, set *RuleSet, req
 	http.Redirect(writer, request.WAFRaw(), request.WAFRaw().URL.String(), http.StatusTemporaryRedirect)
 
 	if request.WAFRaw().ProtoMajor == 1 {
-		_ = this.CloseConn(writer)
+		request.WAFClose()
 	}
 
 	return true
diff --git a/internal/waf/action_record_ip.go b/internal/waf/action_record_ip.go
index 353b6f9..58546b5 100644
--- a/internal/waf/action_record_ip.go
+++ b/internal/waf/action_record_ip.go
@@ -90,7 +90,9 @@ func (this *RecordIPAction) Perform(waf *WAF, group *RuleGroup, set *RuleSet, re
 	expiredAt := time.Now().Unix() + int64(timeout)
 
 	if this.Type == "black" {
-		_ = this.CloseConn(writer)
+		writer.WriteHeader(http.StatusForbidden)
+
+		request.WAFClose()
 
 		SharedIPBlackList.Add(IPTypeAll, request.WAFRemoteIP(), expiredAt)
 	} else {