From 5bb81b70d6afa925467b72dc21c20924eee2f2a5 Mon Sep 17 00:00:00 2001
From: gVisor bot <gvisor-bot@google.com>
Date: Sat, 28 Mar 2020 20:05:38 +0800
Subject: [PATCH] Fix: udp traffic track (#608)

---
 tunnel/connection.go | 2 --
 tunnel/tracker.go    | 8 ++++++++
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/tunnel/connection.go b/tunnel/connection.go
index 9553ed73..ce70a18b 100644
--- a/tunnel/connection.go
+++ b/tunnel/connection.go
@@ -85,7 +85,6 @@ func handleUDPToRemote(packet C.UDPPacket, pc C.PacketConn, metadata *C.Metadata
 	if _, err := pc.WriteWithMetadata(packet.Data(), metadata); err != nil {
 		return
 	}
-	DefaultManager.Upload() <- int64(len(packet.Data()))
 }
 
 func handleUDPToLocal(packet C.UDPPacket, pc net.PacketConn, key string, fAddr net.Addr) {
@@ -109,7 +108,6 @@ func handleUDPToLocal(packet C.UDPPacket, pc net.PacketConn, key string, fAddr n
 		if err != nil {
 			return
 		}
-		DefaultManager.Download() <- int64(n)
 	}
 }
 
diff --git a/tunnel/tracker.go b/tunnel/tracker.go
index e96c2835..2be52203 100644
--- a/tunnel/tracker.go
+++ b/tunnel/tracker.go
@@ -103,6 +103,14 @@ func (ut *udpTracker) WriteTo(b []byte, addr net.Addr) (int, error) {
 	return n, err
 }
 
+func (ut *udpTracker) WriteWithMetadata(p []byte, metadata *C.Metadata) (int, error) {
+	n, err := ut.PacketConn.WriteWithMetadata(p, metadata)
+	upload := int64(n)
+	ut.manager.Upload() <- upload
+	ut.UploadTotal += upload
+	return n, err
+}
+
 func (ut *udpTracker) Close() error {
 	ut.manager.Leave(ut)
 	return ut.PacketConn.Close()