修复Websocket无法正常交互的问题

internal/const/const.go

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

internal/nodes/http_request_cache.go

@@ -19,6 +19,11 @@ import (
 
 // 读取缓存
 func (this *HTTPRequest) doCacheRead(useStale bool) (shouldStop bool) {
+	// 需要动态Upgrade的不缓存
+	if len(this.RawReq.Header.Get("Upgrade")) > 0 {
+		return
+	}
+
 	this.cacheCanTryStale = false
 
 	var cachePolicy = this.ReqServer.HTTPCachePolicy

internal/nodes/http_request_waf.go

@@ -218,7 +218,7 @@ func (this *HTTPRequest) checkWAFRequest(firewallPolicy *firewallconfigs.HTTPFir
 	}
 
 	goNext, hasRequestBody, ruleGroup, ruleSet, err := w.MatchRequest(this, this.writer)
-	if forceLog && logRequestBody && hasRequestBody {
+	if forceLog && logRequestBody && hasRequestBody && ruleSet != nil && ruleSet.HasAttackActions() {
 		this.wafHasRequestBody = true
 	}
 	if err != nil {
@@ -294,7 +294,7 @@ func (this *HTTPRequest) checkWAFResponse(firewallPolicy *firewallconfigs.HTTPFi
 	}
 
 	goNext, hasRequestBody, ruleGroup, ruleSet, err := w.MatchResponse(this, resp, this.writer)
-	if forceLog && logRequestBody && hasRequestBody {
+	if forceLog && logRequestBody && hasRequestBody && ruleSet != nil && ruleSet.HasAttackActions() {
 		this.wafHasRequestBody = true
 	}
 	if err != nil {

internal/nodes/http_request_websocket.go

@@ -2,6 +2,7 @@ package nodes
 
 import (
 	"bufio"
+	"bytes"
 	"errors"
 	"github.com/TeaOSLab/EdgeNode/internal/utils"
 	"io"
@@ -9,8 +10,36 @@ import (
 	"net/url"
 )
 
+// WebsocketResponseReader Websocket响应Reader
+type WebsocketResponseReader struct {
+	rawReader io.Reader
+	buf       []byte
+}
+
+func NewWebsocketResponseReader(rawReader io.Reader) *WebsocketResponseReader {
+	return &WebsocketResponseReader{
+		rawReader: rawReader,
+	}
+}
+
+func (this *WebsocketResponseReader) Read(p []byte) (n int, err error) {
+	n, err = this.rawReader.Read(p)
+	if n > 0 {
+		if len(this.buf) == 0 {
+			this.buf = make([]byte, n)
+			copy(this.buf, p[:n])
+		} else {
+			this.buf = append(this.buf, p[:n]...)
+		}
+	}
+	return
+}
+
 // 处理Websocket请求
 func (this *HTTPRequest) doWebsocket(requestHost string, isLastRetry bool) (shouldRetry bool) {
+	// 设置不缓存
+	this.web.Cache = nil
+
 	if this.web.WebsocketRef == nil || !this.web.WebsocketRef.IsOn || this.web.Websocket == nil || !this.web.Websocket.IsOn {
 		this.writer.WriteHeader(http.StatusForbidden)
 		this.addError(errors.New("websocket have not been enabled yet"))
@@ -84,14 +113,20 @@ func (this *HTTPRequest) doWebsocket(requestHost string, isLastRetry bool) (shou
 
 	go func() {
 		// 读取第一个响应
-		resp, err := http.ReadResponse(bufio.NewReader(originConn), this.RawReq)
+		var respReader = NewWebsocketResponseReader(originConn)
+		resp, err := http.ReadResponse(bufio.NewReader(respReader), this.RawReq)
 		if err != nil {
+			if resp.Body != nil {
+				_ = resp.Body.Close()
+			}
+
 			_ = clientConn.Close()
 			_ = originConn.Close()
 			return
 		}
 
 		this.processResponseHeaders(resp.Header, resp.StatusCode)
+		this.writer.statusCode = resp.StatusCode
 
 		// 将响应写回客户端
 		err = resp.Write(clientConn)
@@ -105,6 +140,25 @@ func (this *HTTPRequest) doWebsocket(requestHost string, isLastRetry bool) (shou
 			return
 		}
 
+		// 剩余已经从源站读取的内容
+		var headerBytes = respReader.buf
+		var headerIndex = bytes.Index(headerBytes, []byte{'\r', '\n', '\r', '\n'}) // CRLF
+		if headerIndex > 0 {
+			var leftBytes = headerBytes[headerIndex+4:]
+			if len(leftBytes) > 0 {
+				_, err = clientConn.Write(leftBytes)
+				if err != nil {
+					if resp.Body != nil {
+						_ = resp.Body.Close()
+					}
+
+					_ = clientConn.Close()
+					_ = originConn.Close()
+					return
+				}
+			}
+		}
+
 		if resp.Body != nil {
 			_ = resp.Body.Close()
 		}

internal/nodes/node.go

@@ -759,6 +759,7 @@ func (this *Node) listenSock() error {
 
 				// 退出主进程
 				events.Notify(events.EventQuit)
+				time.Sleep(100 * time.Millisecond)
 				utils.Exit()
 			case "quit":
 				_ = cmd.ReplyOk()

internal/utils/dbs/db.go

@@ -20,6 +20,9 @@ func NewDB(rawDB *sql.DB) *DB {
 		rawDB: rawDB,
 	}
 
+	events.OnKey(events.EventQuit, fmt.Sprintf("db_%p", db), func() {
+		_ = rawDB.Close()
+	})
 	events.OnKey(events.EventTerminated, fmt.Sprintf("db_%p", db), func() {
 		_ = rawDB.Close()
 	})

internal/utils/readers/reader_print.go

@@ -0,0 +1,28 @@
+// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn .
+
+package readers
+
+import (
+	"io"
+	"log"
+)
+
+type PrintReader struct {
+	rawReader io.Reader
+	tag       string
+}
+
+func NewPrintReader(rawReader io.Reader, tag string) io.Reader {
+	return &PrintReader{
+		rawReader: rawReader,
+		tag:       tag,
+	}
+}
+
+func (this *PrintReader) Read(p []byte) (n int, err error) {
+	n, err = this.rawReader.Read(p)
+	if n > 0 {
+		log.Println("[" + this.tag + "]" + string(p[:n]))
+	}
+	return
+}

internal/utils/service_linux.go

@@ -114,7 +114,8 @@ func (this *ServiceManager) installSystemdService(systemd, exePath string, args
 	var shortName = teaconst.SystemdServiceName
 	var longName = "GoEdge Node" // TODO 将来可以修改
 
-	var desc = `# Provides:          ` + shortName + `
+	var desc = `### BEGIN INIT INFO
+# Provides:          ` + shortName + `
 # Required-Start:    $all
 # Required-Stop:
 # Default-Start:     2 3 4 5

internal/utils/writers/writer_print.go

@@ -0,0 +1,28 @@
+// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn .
+
+package writers
+
+import (
+	"io"
+	"log"
+)
+
+type PrintWriter struct {
+	rawWriter io.Writer
+	tag       string
+}
+
+func NewPrintWriter(rawWriter io.Writer, tag string) io.Writer {
+	return &PrintWriter{
+		rawWriter: rawWriter,
+		tag:       tag,
+	}
+}
+
+func (this *PrintWriter) Write(p []byte) (n int, err error) {
+	n, err = this.rawWriter.Write(p)
+	if n > 0 {
+		log.Println("[" + this.tag + "]" + string(p[:n]))
+	}
+	return
+}