优化代码

internal/const/const.go

@@ -1,7 +1,7 @@
 package teaconst
 
 const (
-	Version = "0.6.0"
+	Version = "0.6.1"
 
 	ProductName = "Edge Node"
 	ProcessName = "edge-node"

internal/nodes/client_conn.go

@@ -42,7 +42,11 @@ type ClientConn struct {
 	lastErr     error
 
 	readDeadlineTime int64
-	isShortReading   bool // header or handshake
+	isShortReading   bool // reading header or tls handshake
+
+	isDebugging      bool
+	autoReadTimeout  bool
+	autoWriteTimeout bool
 }
 
 func NewClientConn(rawConn net.Conn, isHTTP bool, isTLS bool, isInAllowList bool) net.Conn {
@@ -59,6 +63,14 @@ func NewClientConn(rawConn net.Conn, isHTTP bool, isTLS bool, isInAllowList bool
 		createdAt:      time.Now().Unix(),
 	}
 
+	var globalServerConfig = sharedNodeConfig.GlobalServerConfig
+	if globalServerConfig != nil {
+		var performanceConfig = globalServerConfig.Performance
+		conn.isDebugging = performanceConfig.Debug
+		conn.autoReadTimeout = performanceConfig.AutoReadTimeout
+		conn.autoWriteTimeout = performanceConfig.AutoWriteTimeout
+	}
+
 	if isHTTP {
 		// TODO 可以在配置中设置此值
 		_ = conn.SetLinger(nodeconfigs.DefaultTCPLinger)
@@ -71,9 +83,7 @@ func NewClientConn(rawConn net.Conn, isHTTP bool, isTLS bool, isInAllowList bool
 }
 
 func (this *ClientConn) Read(b []byte) (n int, err error) {
-	var globalServerConfig = sharedNodeConfig.GlobalServerConfig
-
-	if globalServerConfig != nil && globalServerConfig.Performance.Debug {
+	if this.isDebugging {
 		this.lastReadAt = time.Now().Unix()
 
 		defer func() {
@@ -93,8 +103,7 @@ func (this *ClientConn) Read(b []byte) (n int, err error) {
 	}
 
 	// 设置读超时时间
-	var autoReadTimeout = globalServerConfig != nil && globalServerConfig.Performance.AutoReadTimeout
-	if this.isHTTP && !this.isWebsocket && !this.isShortReading && autoReadTimeout {
+	if this.isHTTP && !this.isWebsocket && !this.isShortReading && this.autoReadTimeout {
 		this.setHTTPReadTimeout()
 	}
 
@@ -134,9 +143,7 @@ func (this *ClientConn) Read(b []byte) (n int, err error) {
 }
 
 func (this *ClientConn) Write(b []byte) (n int, err error) {
-	var globalServerConfig = sharedNodeConfig.GlobalServerConfig
-
-	if globalServerConfig != nil && globalServerConfig.Performance.Debug {
+	if this.isDebugging {
 		this.lastWriteAt = time.Now().Unix()
 
 		defer func() {
@@ -147,7 +154,7 @@ func (this *ClientConn) Write(b []byte) (n int, err error) {
 	}
 
 	// 设置写超时时间
-	if globalServerConfig != nil && globalServerConfig.Performance.AutoWriteTimeout {
+	if this.autoWriteTimeout {
 		// TODO L2 -> L1 写入时不限制时间
 		var timeoutSeconds = len(b) / 1024
 		if timeoutSeconds < 3 {
@@ -157,7 +164,7 @@ func (this *ClientConn) Write(b []byte) (n int, err error) {
 	}
 
 	// 延长读超时时间
-	if this.isHTTP && !this.isWebsocket && globalServerConfig != nil && globalServerConfig.Performance.AutoReadTimeout {
+	if this.isHTTP && !this.isWebsocket && this.autoReadTimeout {
 		this.setHTTPReadTimeout()
 	}
 
@@ -216,8 +223,7 @@ func (this *ClientConn) SetDeadline(t time.Time) error {
 
 func (this *ClientConn) SetReadDeadline(t time.Time) error {
 	// 如果开启了HTTP自动读超时选项，则自动控制超时时间
-	var globalServerConfig = sharedNodeConfig.GlobalServerConfig
-	if this.isHTTP && !this.isWebsocket && globalServerConfig != nil && globalServerConfig.Performance.AutoReadTimeout {
+	if this.isHTTP && !this.isWebsocket && this.autoReadTimeout {
 		this.isShortReading = false
 
 		var unixTime = t.Unix()

internal/nodes/http_request_host_redirect.go

@@ -146,6 +146,13 @@ func (this *HTTPRequest) doHostRedirect() (blocked bool) {
 					u.Status = http.StatusTemporaryRedirect
 				}
 				this.processResponseHeaders(this.writer.Header(), u.Status)
+
+				// 参数
+				var qIndex = strings.Index(this.uri, "?")
+				if qIndex >= 0 {
+					afterURL += this.uri[qIndex:]
+				}
+
 				http.Redirect(this.RawWriter, this.RawReq, afterURL, u.Status)
 				return true
 			}

internal/waf/action_base.go

@@ -2,10 +2,6 @@
 
 package waf
 
-import (
-	"net/http"
-)
-
 type BaseAction struct {
 	currentActionId int64
 }
@@ -19,16 +15,3 @@ func (this *BaseAction) ActionId() int64 {
 func (this *BaseAction) SetActionId(actionId int64) {
 	this.currentActionId = actionId
 }
-
-// CloseConn 关闭连接
-func (this *BaseAction) CloseConn(writer http.ResponseWriter) error {
-	// 断开连接
-	hijack, ok := writer.(http.Hijacker)
-	if ok {
-		conn, _, err := hijack.Hijack()
-		if err == nil && conn != nil {
-			return conn.Close()
-		}
-	}
-	return nil
-}

internal/waf/action_get_302.go

@@ -70,7 +70,12 @@ func (this *Get302Action) Perform(waf *WAF, group *RuleGroup, set *RuleSet, requ
 	http.Redirect(writer, request.WAFRaw(), Get302Path+"?info="+url.QueryEscape(info), http.StatusFound)
 
 	if request.WAFRaw().ProtoMajor == 1 {
-		_ = this.CloseConn(writer)
+		flusher, ok := writer.(http.Flusher)
+		if ok {
+			flusher.Flush()
+		}
+
+		request.WAFClose()
 	}
 
 	return false, false

internal/waf/action_post_307.go

@@ -87,7 +87,12 @@ 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)
+		flusher, ok := writer.(http.Flusher)
+		if ok {
+			flusher.Flush()
+		}
+
+		request.WAFClose()
 	}
 
 	return false, false