增加服务带宽统计

internal/nodes/client_conn.go

@@ -7,12 +7,14 @@ import (
 	"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs/firewallconfigs"
 	teaconst "github.com/TeaOSLab/EdgeNode/internal/const"
 	"github.com/TeaOSLab/EdgeNode/internal/iplibrary"
+	"github.com/TeaOSLab/EdgeNode/internal/stats"
 	"github.com/TeaOSLab/EdgeNode/internal/ttlcache"
 	"github.com/TeaOSLab/EdgeNode/internal/utils"
 	"github.com/TeaOSLab/EdgeNode/internal/waf"
 	"github.com/iwind/TeaGo/types"
 	"net"
 	"os"
+	"strings"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -26,6 +28,8 @@ type ClientConn struct {
 	hasDeadline bool
 	hasRead     bool
 
+	isLO bool // 是否为环路
+
 	hasResetSYNFlood bool
 
 	BaseClientConn
@@ -41,10 +45,28 @@ func NewClientConn(conn net.Conn, isTLS bool, quickClose bool) net.Conn {
 		}
 	}
 
-	return &ClientConn{BaseClientConn: BaseClientConn{rawConn: conn}, isTLS: isTLS}
+	// 是否为环路
+	var remoteAddr = conn.RemoteAddr().String()
+	var isLO = strings.HasPrefix(remoteAddr, "127.0.0.1:") || strings.HasPrefix(remoteAddr, "[::1]:")
+
+	return &ClientConn{
+		BaseClientConn: BaseClientConn{rawConn: conn},
+		isTLS:          isTLS,
+		isLO:           isLO,
+	}
 }
 
 func (this *ClientConn) Read(b []byte) (n int, err error) {
+	// 环路直接读取
+	if this.isLO {
+		n, err = this.rawConn.Read(b)
+		if n > 0 {
+			atomic.AddUint64(&teaconst.InTrafficBytes, uint64(n))
+		}
+		return
+	}
+
+	// TLS
 	if this.isTLS {
 		if !this.hasDeadline {
 			_ = this.rawConn.SetReadDeadline(time.Now().Add(time.Duration(nodeconfigs.DefaultTLSHandshakeTimeout) * time.Second)) // TODO 握手超时时间可以设置
@@ -55,6 +77,7 @@ func (this *ClientConn) Read(b []byte) (n int, err error) {
 		}
 	}
 
+	// 开始读取
 	n, err = this.rawConn.Read(b)
 	if n > 0 {
 		atomic.AddUint64(&teaconst.InTrafficBytes, uint64(n))
@@ -85,7 +108,15 @@ 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 { // 环路不统计带宽，避免缓存预热等行为产生带宽
+				stats.SharedBandwidthStatManager.Add(this.serverId, int64(n))
+			}
+		}
 	}
+
 	return
 }
 

internal/nodes/client_conn_base.go

@@ -36,6 +36,16 @@ func (this *BaseClientConn) Bind(serverId int64, remoteAddr string, maxConnsPerS
 	return sharedClientConnLimiter.Add(this.rawConn.RemoteAddr().String(), serverId, remoteAddr, maxConnsPerServer, maxConnsPerIP)
 }
 
+// SetServerId 设置服务ID
+func (this *BaseClientConn) SetServerId(serverId int64) {
+	this.serverId = serverId
+}
+
+// ServerId 读取当前连接绑定的服务ID
+func (this *BaseClientConn) ServerId() int64 {
+	return this.serverId
+}
+
 // RawIP 原本IP
 func (this *BaseClientConn) RawIP() string {
 	ip, _, _ := net.SplitHostPort(this.rawConn.RemoteAddr().String())

internal/nodes/client_conn_interface.go

@@ -11,4 +11,10 @@ type ClientConnInterface interface {
 
 	// Bind 绑定服务
 	Bind(serverId int64, remoteAddr string, maxConnsPerServer int, maxConnsPerIP int) bool
+
+	// ServerId 获取服务ID
+	ServerId() int64
+
+	// SetServerId 设置服务ID
+	SetServerId(serverId int64)
 }

internal/nodes/http_request.go

@@ -356,7 +356,7 @@ func (this *HTTPRequest) doEnd() {
 
 	// 流量统计
 	// TODO 增加是否开启开关
-	if this.ReqServer != nil {
+	if this.ReqServer != nil && this.ReqServer.Id > 0 {
 		var countCached int64 = 0
 		var cachedBytes int64 = 0
 
@@ -373,17 +373,17 @@ func (this *HTTPRequest) doEnd() {
 		}
 
 		stats.SharedTrafficStatManager.Add(this.ReqServer.Id, this.ReqHost, this.writer.SentBodyBytes()+this.writer.SentHeaderBytes(), cachedBytes, 1, countCached, countAttacks, attackBytes, this.ReqServer.ShouldCheckTrafficLimit(), this.ReqServer.PlanId())
-	}
 
-	// 指标
-	if metrics.SharedManager.HasHTTPMetrics() {
-		this.doMetricsResponse()
-	}
+		// 指标
+		if metrics.SharedManager.HasHTTPMetrics() {
+			this.doMetricsResponse()
+		}
 
-	// 统计
-	if this.web.StatRef != nil && this.web.StatRef.IsOn {
-		// 放到最后执行
-		this.doStat()
+		// 统计
+		if this.web.StatRef != nil && this.web.StatRef.IsOn {
+			// 放到最后执行
+			this.doStat()
+		}
 	}
 }
 

internal/nodes/http_request_limit.go

@@ -24,7 +24,7 @@ func (this *HTTPRequest) doRequestLimit() (shouldStop bool) {
 
 	// 设置连接相关参数
 	if this.web.RequestLimit.MaxConns > 0 || this.web.RequestLimit.MaxConnsPerIP > 0 {
-		requestConn := this.RawReq.Context().Value(HTTPConnContextKey)
+		var requestConn = this.RawReq.Context().Value(HTTPConnContextKey)
 		if requestConn != nil {
 			clientConn, ok := requestConn.(ClientConnInterface)
 			if ok && !clientConn.IsBound() {

internal/nodes/listener.go

@@ -43,7 +43,7 @@ func (this *Listener) Listen() error {
 	if this.group == nil {
 		return nil
 	}
-	protocol := this.group.Protocol()
+	var protocol = this.group.Protocol()
 	if protocol.IsUDPFamily() {
 		return this.listenUDP()
 	}
@@ -54,7 +54,7 @@ func (this *Listener) listenTCP() error {
 	if this.group == nil {
 		return nil
 	}
-	protocol := this.group.Protocol()
+	var protocol = this.group.Protocol()
 
 	tcpListener, err := this.createTCPListener()
 	if err != nil {

internal/nodes/listener_http.go

@@ -178,6 +178,17 @@ func (this *HTTPListener) ServeHTTP(rawWriter http.ResponseWriter, rawReq *http.
 		}
 	}
 
+	// 绑定连接
+	if server != nil && server.Id > 0 {
+		var requestConn = rawReq.Context().Value(HTTPConnContextKey)
+		if requestConn != nil {
+			clientConn, ok := requestConn.(ClientConnInterface)
+			if ok {
+				clientConn.SetServerId(server.Id)
+			}
+		}
+	}
+
 	// 包装新请求对象
 	var req = &HTTPRequest{
 		RawReq:     rawReq,

internal/nodes/listener_tcp.go

@@ -63,7 +63,6 @@ func (this *TCPListener) Reload(group *serverconfigs.ServerAddressGroup) {
 }
 
 func (this *TCPListener) handleConn(conn net.Conn) error {
-
 	var server = this.Group.FirstServer()
 	if server == nil {
 		return errors.New("no server available")
@@ -72,6 +71,23 @@ func (this *TCPListener) handleConn(conn net.Conn) error {
 		return errors.New("no ReverseProxy configured for the server")
 	}
 
+	// 绑定连接和服务
+	clientConn, ok := conn.(ClientConnInterface)
+	if ok {
+		clientConn.SetServerId(server.Id)
+	} else {
+		tlsConn, ok := conn.(*tls.Conn)
+		if ok {
+			var internalConn = tlsConn.NetConn()
+			if internalConn != nil {
+				clientConn, ok = internalConn.(ClientConnInterface)
+				if ok {
+					clientConn.SetServerId(server.Id)
+				}
+			}
+		}
+	}
+
 	// 是否已达到流量限制
 	if this.reachedTrafficLimit() {
 		// 关闭连接