From cd19a4a7bc26efaefed6d49eab95916594bba3aa 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: Mon, 25 Oct 2021 19:00:42 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96=E8=BF=9E=E6=8E=A5=E5=85=B3?=
 =?UTF-8?q?=E9=97=AD=E9=80=9F=E5=BA=A6?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/nodes/client_conn.go       | 12 +++++++-----
 internal/nodes/client_conn_utils.go | 22 ++++++++++++++++++++++
 internal/nodes/client_listener.go   | 25 +++++++++++++++++--------
 internal/nodes/http_request_waf.go  | 12 ++++--------
 internal/nodes/listener.go          |  2 +-
 5 files changed, 51 insertions(+), 22 deletions(-)
 create mode 100644 internal/nodes/client_conn_utils.go

diff --git a/internal/nodes/client_conn.go b/internal/nodes/client_conn.go
index 81219a4..e0d7b56 100644
--- a/internal/nodes/client_conn.go
+++ b/internal/nodes/client_conn.go
@@ -48,11 +48,13 @@ type ClientConn struct {
 	isClosed bool
 }
 
-func NewClientConn(conn net.Conn) net.Conn {
-	tcpConn, ok := conn.(*net.TCPConn)
-	if ok {
-		// TODO 可以设置此值
-		_ = tcpConn.SetLinger(0)
+func NewClientConn(conn net.Conn, quickClose bool) net.Conn {
+	if quickClose {
+		tcpConn, ok := conn.(*net.TCPConn)
+		if ok {
+			// TODO 可以设置此值
+			_ = tcpConn.SetLinger(0)
+		}
 	}
 
 	return &ClientConn{rawConn: conn}
diff --git a/internal/nodes/client_conn_utils.go b/internal/nodes/client_conn_utils.go
new file mode 100644
index 0000000..b00b0e7
--- /dev/null
+++ b/internal/nodes/client_conn_utils.go
@@ -0,0 +1,22 @@
+// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package nodes
+
+import (
+	"net"
+)
+
+// 判断客户端连接是否已关闭
+func isClientConnClosed(conn net.Conn) bool {
+	if conn == nil {
+		return true
+	}
+	clientConn, ok := conn.(*ClientConn)
+	if ok {
+		return clientConn.IsClosed()
+	}
+
+	// TODO 解决tls.Conn无法获取底层连接对象的问题
+
+	return false
+}
diff --git a/internal/nodes/client_listener.go b/internal/nodes/client_listener.go
index f2588fd..5033da4 100644
--- a/internal/nodes/client_listener.go
+++ b/internal/nodes/client_listener.go
@@ -11,10 +11,14 @@ import (
 // ClientListener 客户端网络监听
 type ClientListener struct {
 	rawListener net.Listener
+	quickClose  bool
 }
 
-func NewClientListener(listener net.Listener) net.Listener {
-	return &ClientListener{rawListener: listener}
+func NewClientListener(listener net.Listener, quickClose bool) net.Listener {
+	return &ClientListener{
+		rawListener: listener,
+		quickClose:  quickClose,
+	}
 }
 
 func (this *ClientListener) Accept() (net.Conn, error) {
@@ -25,15 +29,20 @@ func (this *ClientListener) Accept() (net.Conn, error) {
 	// 是否在WAF名单中
 	ip, _, err := net.SplitHostPort(conn.RemoteAddr().String())
 	if err == nil {
-		if !waf.SharedIPWhiteList.Contains(waf.IPTypeAll, firewallconfigs.FirewallScopeGlobal, 0, ip) && waf.SharedIPBlackList.Contains(waf.IPTypeAll, firewallconfigs.FirewallScopeGlobal, 0, ip) {
-			defer func() {
-				_ = conn.Close()
-			}()
-			return conn, nil
+		if !waf.SharedIPWhiteList.Contains(waf.IPTypeAll, firewallconfigs.FirewallScopeGlobal, 0, ip) &&
+			waf.SharedIPBlackList.Contains(waf.IPTypeAll, firewallconfigs.FirewallScopeGlobal, 0, ip) {
+
+			tcpConn, ok := conn.(*net.TCPConn)
+			if ok {
+				_ = tcpConn.SetLinger(0)
+			}
+
+			_ = conn.Close()
+			return this.Accept()
 		}
 	}
 
-	return NewClientConn(conn), nil
+	return NewClientConn(conn, this.quickClose), nil
 }
 
 func (this *ClientListener) Close() error {
diff --git a/internal/nodes/http_request_waf.go b/internal/nodes/http_request_waf.go
index 7446de3..d3d3a5d 100644
--- a/internal/nodes/http_request_waf.go
+++ b/internal/nodes/http_request_waf.go
@@ -20,22 +20,18 @@ func (this *HTTPRequest) doWAFRequest() (blocked bool) {
 	// 当前连接是否已关闭
 	var conn = this.RawReq.Context().Value(HTTPConnContextKey)
 	if conn != nil {
-		trafficConn, ok := conn.(*ClientConn)
-		if ok && trafficConn.IsClosed() {
+		if isClientConnClosed(conn.(net.Conn)) {
 			this.disableLog = true
 			return true
 		}
 	}
 
 	// 检查是否在临时黑名单中
-	if waf.SharedIPBlackList.Contains(waf.IPTypeAll, firewallconfigs.FirewallScopeService, this.Server.Id, this.WAFRemoteIP()) {
+	var remoteAddr = this.WAFRemoteIP()
+	if waf.SharedIPBlackList.Contains(waf.IPTypeAll, firewallconfigs.FirewallScopeService, this.Server.Id, remoteAddr) || waf.SharedIPBlackList.Contains(waf.IPTypeAll, firewallconfigs.FirewallScopeGlobal, 0, remoteAddr) {
 		this.disableLog = true
-
 		if conn != nil {
-			trafficConn, ok := conn.(*ClientConn)
-			if ok && !trafficConn.IsClosed() {
-				_ = trafficConn.Close()
-			}
+			_ = conn.(net.Conn).Close()
 		}
 
 		return true
diff --git a/internal/nodes/listener.go b/internal/nodes/listener.go
index d02fa6b..5967e84 100644
--- a/internal/nodes/listener.go
+++ b/internal/nodes/listener.go
@@ -59,7 +59,7 @@ func (this *Listener) listenTCP() error {
 	if err != nil {
 		return err
 	}
-	netListener = NewClientListener(netListener)
+	netListener = NewClientListener(netListener, protocol.IsHTTPFamily() || protocol.IsHTTPSFamily())
 	events.On(events.EventQuit, func() {
 		remotelogs.Println("LISTENER", "quit "+this.group.FullAddr())
 		_ = netListener.Close()