chore: proxyDialer don't push flow to manager in statistic

xishang0128 · Apr 11, 2023 · bad7340 · bad7340
1 parent 7beb091
commit bad7340
Show file tree

Hide file tree

Showing 3 changed files with 48 additions and 12 deletions.
diff --git a/component/proxydialer/proxydialer.go b/component/proxydialer/proxydialer.go
@@ -64,7 +64,7 @@ func (p proxyDialer) DialContext(ctx context.Context, network, address string) (
 		}
 	}
 	if p.statistic {
-		conn = statistic.NewTCPTracker(conn, statistic.DefaultManager, currentMeta, nil, 0, 0)
+		conn = statistic.NewTCPTracker(conn, statistic.DefaultManager, currentMeta, nil, 0, 0, false)
 	}
 	return conn, err
 }
@@ -100,7 +100,7 @@ func (p proxyDialer) listenPacket(ctx context.Context, currentMeta *C.Metadata)
 		}
 	}
 	if p.statistic {
-		pc = statistic.NewUDPTracker(pc, statistic.DefaultManager, currentMeta, nil, 0, 0)
+		pc = statistic.NewUDPTracker(pc, statistic.DefaultManager, currentMeta, nil, 0, 0, false)
 	}
 	return pc, nil
 }

diff --git a/tunnel/statistic/tracker.go b/tunnel/statistic/tracker.go
@@ -32,6 +32,8 @@ type tcpTracker struct {
 	C.Conn `json:"-"`
 	*trackerInfo
 	manager *Manager
+
+	pushToManager bool `json:"-"`
 }
 
 func (tt *tcpTracker) ID() string {
@@ -41,31 +43,39 @@ func (tt *tcpTracker) ID() string {
 func (tt *tcpTracker) Read(b []byte) (int, error) {
 	n, err := tt.Conn.Read(b)
 	download := int64(n)
-	tt.manager.PushDownloaded(download)
+	if tt.pushToManager {
+		tt.manager.PushDownloaded(download)
+	}
 	tt.DownloadTotal.Add(download)
 	return n, err
 }
 
 func (tt *tcpTracker) ReadBuffer(buffer *buf.Buffer) (err error) {
 	err = tt.Conn.ReadBuffer(buffer)
 	download := int64(buffer.Len())
-	tt.manager.PushDownloaded(download)
+	if tt.pushToManager {
+		tt.manager.PushDownloaded(download)
+	}
 	tt.DownloadTotal.Add(download)
 	return
 }
 
 func (tt *tcpTracker) Write(b []byte) (int, error) {
 	n, err := tt.Conn.Write(b)
 	upload := int64(n)
-	tt.manager.PushUploaded(upload)
+	if tt.pushToManager {
+		tt.manager.PushUploaded(upload)
+	}
 	tt.UploadTotal.Add(upload)
 	return n, err
 }
 
 func (tt *tcpTracker) WriteBuffer(buffer *buf.Buffer) (err error) {
 	upload := int64(buffer.Len())
 	err = tt.Conn.WriteBuffer(buffer)
-	tt.manager.PushUploaded(upload)
+	if tt.pushToManager {
+		tt.manager.PushUploaded(upload)
+	}
 	tt.UploadTotal.Add(upload)
 	return
 }
@@ -79,7 +89,7 @@ func (tt *tcpTracker) Upstream() any {
 	return tt.Conn
 }
 
-func NewTCPTracker(conn C.Conn, manager *Manager, metadata *C.Metadata, rule C.Rule, uploadTotal int64, downloadTotal int64) *tcpTracker {
+func NewTCPTracker(conn C.Conn, manager *Manager, metadata *C.Metadata, rule C.Rule, uploadTotal int64, downloadTotal int64, pushToManager bool) *tcpTracker {
 	if conn != nil {
 		if tcpAddr, ok := conn.RemoteAddr().(*net.TCPAddr); ok {
 			metadata.RemoteDst = tcpAddr.IP.String()
@@ -100,6 +110,16 @@ func NewTCPTracker(conn C.Conn, manager *Manager, metadata *C.Metadata, rule C.R
 			UploadTotal:   atomic.NewInt64(uploadTotal),
 			DownloadTotal: atomic.NewInt64(downloadTotal),
 		},
+		pushToManager: pushToManager,
+	}
+
+	if pushToManager {
+		if uploadTotal > 0 {
+			manager.PushUploaded(uploadTotal)
+		}
+		if downloadTotal > 0 {
+			manager.PushDownloaded(downloadTotal)
+		}
 	}
 
 	if rule != nil {
@@ -115,6 +135,8 @@ type udpTracker struct {
 	C.PacketConn `json:"-"`
 	*trackerInfo
 	manager *Manager
+
+	pushToManager bool `json:"-"`
 }
 
 func (ut *udpTracker) ID() string {
@@ -124,15 +146,19 @@ func (ut *udpTracker) ID() string {
 func (ut *udpTracker) ReadFrom(b []byte) (int, net.Addr, error) {
 	n, addr, err := ut.PacketConn.ReadFrom(b)
 	download := int64(n)
-	ut.manager.PushDownloaded(download)
+	if ut.pushToManager {
+		ut.manager.PushDownloaded(download)
+	}
 	ut.DownloadTotal.Add(download)
 	return n, addr, err
 }
 
 func (ut *udpTracker) WriteTo(b []byte, addr net.Addr) (int, error) {
 	n, err := ut.PacketConn.WriteTo(b, addr)
 	upload := int64(n)
-	ut.manager.PushUploaded(upload)
+	if ut.pushToManager {
+		ut.manager.PushUploaded(upload)
+	}
 	ut.UploadTotal.Add(upload)
 	return n, err
 }
@@ -142,7 +168,7 @@ func (ut *udpTracker) Close() error {
 	return ut.PacketConn.Close()
 }
 
-func NewUDPTracker(conn C.PacketConn, manager *Manager, metadata *C.Metadata, rule C.Rule, uploadTotal int64, downloadTotal int64) *udpTracker {
+func NewUDPTracker(conn C.PacketConn, manager *Manager, metadata *C.Metadata, rule C.Rule, uploadTotal int64, downloadTotal int64, pushToManager bool) *udpTracker {
 	metadata.RemoteDst = conn.RemoteDestination()
 
 	ut := &udpTracker{
@@ -157,6 +183,16 @@ func NewUDPTracker(conn C.PacketConn, manager *Manager, metadata *C.Metadata, ru
 			UploadTotal:   atomic.NewInt64(uploadTotal),
 			DownloadTotal: atomic.NewInt64(downloadTotal),
 		},
+		pushToManager: pushToManager,
+	}
+
+	if pushToManager {
+		if uploadTotal > 0 {
+			manager.PushUploaded(uploadTotal)
+		}
+		if downloadTotal > 0 {
+			manager.PushDownloaded(downloadTotal)
+		}
 	}
 
 	if rule != nil {

diff --git a/tunnel/tunnel.go b/tunnel/tunnel.go
@@ -364,7 +364,7 @@ func handleUDPConn(packet C.PacketAdapter) {
 		}
 		pCtx.InjectPacketConn(rawPc)
 
-		pc := statistic.NewUDPTracker(rawPc, statistic.DefaultManager, metadata, rule, 0, 0)
+		pc := statistic.NewUDPTracker(rawPc, statistic.DefaultManager, metadata, rule, 0, 0, true)
 
 		switch true {
 		case metadata.SpecialProxy != "":
@@ -494,7 +494,7 @@ func handleTCPConn(connCtx C.ConnContext) {
 		return
 	}
 
-	remoteConn = statistic.NewTCPTracker(remoteConn, statistic.DefaultManager, metadata, rule, 0, int64(peekLen))
+	remoteConn = statistic.NewTCPTracker(remoteConn, statistic.DefaultManager, metadata, rule, 0, int64(peekLen), true)
 	defer func(remoteConn C.Conn) {
 		_ = remoteConn.Close()
 	}(remoteConn)