修复HTTPS连接无法记录带宽的问题，优化带宽计算方法

internal/nodes/client_conn.go

@@ -11,6 +11,7 @@ import (
 	"github.com/TeaOSLab/EdgeNode/internal/ttlcache"
 	"github.com/TeaOSLab/EdgeNode/internal/utils"
 	"github.com/TeaOSLab/EdgeNode/internal/waf"
+	"github.com/iwind/TeaGo/Tea"
 	"github.com/iwind/TeaGo/types"
 	"net"
 	"os"
@@ -22,6 +23,8 @@ import (
 
 // ClientConn 客户端连接
 type ClientConn struct {
+	BaseClientConn
+
 	once sync.Once
 
 	isTLS       bool
@@ -31,8 +34,6 @@ type ClientConn struct {
 	isLO bool // 是否为环路
 
 	hasResetSYNFlood bool
-
-	BaseClientConn
 }
 
 func NewClientConn(conn net.Conn, isTLS bool, quickClose bool) net.Conn {
@@ -110,11 +111,10 @@ func (this *ClientConn) Read(b []byte) (n int, err error) {
 func (this *ClientConn) Write(b []byte) (n int, err error) {
 	n, err = this.rawConn.Write(b)
 	if n > 0 {
-		atomic.AddUint64(&teaconst.OutTrafficBytes, uint64(n))
-
 		// 统计当前服务带宽
 		if this.serverId > 0 {
-			if !this.isLO { // 环路不统计带宽，避免缓存预热等行为产生带宽
+			if !this.isLO || Tea.IsTesting() { // 环路不统计带宽，避免缓存预热等行为产生带宽
+				atomic.AddUint64(&teaconst.OutTrafficBytes, uint64(n))
 				stats.SharedBandwidthStatManager.Add(this.userId, this.serverId, int64(n))
 			}
 		}

internal/nodes/client_conn_base.go

@@ -2,7 +2,10 @@
 
 package nodes
 
-import "net"
+import (
+	"crypto/tls"
+	"net"
+)
 
 type BaseClientConn struct {
 	rawConn net.Conn
@@ -42,6 +45,17 @@ func (this *BaseClientConn) Bind(serverId int64, remoteAddr string, maxConnsPerS
 // SetServerId 设置服务ID
 func (this *BaseClientConn) SetServerId(serverId int64) {
 	this.serverId = serverId
+
+	// 设置包装前连接
+	switch conn := this.rawConn.(type) {
+	case *tls.Conn:
+		nativeConn, ok := conn.NetConn().(ClientConnInterface)
+		if ok {
+			nativeConn.SetServerId(serverId)
+		}
+	case *ClientConn:
+		conn.SetServerId(serverId)
+	}
 }
 
 // ServerId 读取当前连接绑定的服务ID
@@ -52,6 +66,17 @@ func (this *BaseClientConn) ServerId() int64 {
 // SetUserId 设置所属服务的用户ID
 func (this *BaseClientConn) SetUserId(userId int64) {
 	this.userId = userId
+
+	// 设置包装前连接
+	switch conn := this.rawConn.(type) {
+	case *tls.Conn:
+		nativeConn, ok := conn.NetConn().(ClientConnInterface)
+		if ok {
+			nativeConn.SetUserId(userId)
+		}
+	case *ClientConn:
+		conn.SetUserId(userId)
+	}
 }
 
 // UserId 获取当前连接所属服务的用户ID
@@ -66,9 +91,15 @@ func (this *BaseClientConn) RawIP() string {
 }
 
 // TCPConn 转换为TCPConn
-func (this *BaseClientConn) TCPConn() (*net.TCPConn, bool) {
-	conn, ok := this.rawConn.(*net.TCPConn)
-	return conn, ok
+func (this *BaseClientConn) TCPConn() (tcpConn *net.TCPConn, ok bool) {
+	// 设置包装前连接
+	switch conn := this.rawConn.(type) {
+	case *tls.Conn:
+		tcpConn, ok = conn.NetConn().(*net.TCPConn)
+	default:
+		tcpConn, ok = this.rawConn.(*net.TCPConn)
+	}
+	return
 }
 
 // SetLinger 设置Linger

internal/nodes/node.go

@@ -869,6 +869,11 @@ func (this *Node) listenSock() error {
 				} else {
 					_ = cmd.ReplyOk()
 				}
+			case "bandwidth":
+				var m = stats.SharedBandwidthStatManager.Map()
+				_ = cmd.Reply(&gosock.Command{Params: maps.Map{
+					"stats": m,
+				}})
 			}
 		})