修复在HTTPS下无法连接Websocket的问题

internal/const/const.go

@@ -1,7 +1,7 @@
 package teaconst
 
 const (
-	Version = "0.6.0"
+	Version = "0.6.2"
 
 	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 this.isDebugging {
-
-	if globalServerConfig != nil && globalServerConfig.Performance.Debug {
 		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 && this.autoReadTimeout {
-	if this.isHTTP && !this.isShortReading && 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 this.isDebugging {
-
-	if globalServerConfig != nil && globalServerConfig.Performance.Debug {
 		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 {
+	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 && this.autoReadTimeout {
-	if this.isHTTP && globalServerConfig != nil && globalServerConfig.Performance.AutoReadTimeout {
 		this.isShortReading = false
 
 		var unixTime = t.Unix()

internal/nodes/client_conn_base.go

@@ -16,6 +16,8 @@ type BaseClientConn struct {
 	remoteAddr string
 	hasLimit   bool
 
+	isWebsocket bool
+
 	isClosed bool
 
 	rawIP string
@@ -122,3 +124,7 @@ func (this *BaseClientConn) SetLinger(seconds int) error {
 	}
 	return nil
 }
+
+func (this *BaseClientConn) SetIsWebsocket(isWebsocket bool) {
+	this.isWebsocket = isWebsocket
+}

internal/nodes/client_conn_interface.go

@@ -23,4 +23,7 @@ type ClientConnInterface interface {
 
 	// UserId 获取当前连接所属服务的用户ID
 	UserId() int64
+
+	// SetIsWebsocket 设置是否为Websocket
+	SetIsWebsocket(isWebsocket bool)
 }

internal/nodes/client_tls_conn.go

@@ -55,3 +55,16 @@ func (this *ClientTLSConn) SetReadDeadline(t time.Time) error {
 func (this *ClientTLSConn) SetWriteDeadline(t time.Time) error {
 	return this.rawConn.SetWriteDeadline(t)
 }
+
+func (this *ClientTLSConn) SetIsWebsocket(isWebsocket bool) {
+	tlsConn, ok := this.rawConn.(*tls.Conn)
+	if ok {
+		var rawConn = tlsConn.NetConn()
+		if rawConn != nil {
+			clientConn, ok := rawConn.(*ClientConn)
+			if ok {
+				clientConn.SetIsWebsocket(isWebsocket)
+			}
+		}
+	}
+}

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/nodes/http_request_websocket.go

@@ -70,6 +70,13 @@ func (this *HTTPRequest) doWebsocket(requestHost string, isLastRetry bool) (shou
 		this.RawReq.Header.Set("Origin", newRequestOrigin)
 	}
 
+	// 获取当前连接
+	var requestConn = this.RawReq.Context().Value(HTTPConnContextKey)
+	if requestConn == nil {
+		return
+	}
+
+	// 连接源站
 	// TODO 增加N次错误重试，重试的时候需要尝试不同的源站
 	originConn, _, err := OriginConnect(this.origin, this.requestServerPort(), this.RawReq.RemoteAddr, requestHost)
 	if err != nil {
@@ -102,6 +109,11 @@ func (this *HTTPRequest) doWebsocket(requestHost string, isLastRetry bool) (shou
 		return
 	}
 
+	requestClientConn, ok := requestConn.(ClientConnInterface)
+	if ok {
+		requestClientConn.SetIsWebsocket(true)
+	}
+
 	clientConn, _, err := this.writer.Hijack()
 	if err != nil || clientConn == nil {
 		this.write50x(err, http.StatusInternalServerError, "Failed to get origin site connection", "获取源站连接失败", false)

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