Imrpove flow tracking & sniff action

common/trafficcontrol/tracker.go

@@ -7,6 +7,7 @@ import (
 	"time"
 
 	"github.com/sagernet/sing-box/adapter"
+	"github.com/sagernet/sing-tun"
 	"github.com/sagernet/sing/common"
 	"github.com/sagernet/sing/common/bufio"
 	N "github.com/sagernet/sing/common/network"
@@ -68,6 +69,13 @@ func (m *Manager) RoutedPacketConnection(ctx context.Context, conn N.PacketConn,
 	return tracker
 }
 
+func (m *Manager) RoutedFlow(ctx context.Context, metadata adapter.InboundContext, matchedRule adapter.Rule, matchOutbound adapter.Outbound) tun.FlowTracker {
+	return &flowTracker{
+		metadata: m.newTrackerMetadata(metadata, matchedRule, matchOutbound, new(atomic.Int64), new(atomic.Int64)),
+		manager:  m,
+	}
+}
+
 func (m *Manager) newTrackerMetadata(metadata adapter.InboundContext, matchedRule adapter.Rule, matchOutbound adapter.Outbound, upload *atomic.Int64, download *atomic.Int64) TrackerMetadata {
 	id, _ := uuid.NewV4()
 	var (
@@ -135,6 +143,53 @@ func (t *connTracker) WriterReplaceable() bool {
 	return true
 }
 
+var (
+	_ Tracker         = (*flowTracker)(nil)
+	_ tun.FlowTracker = (*flowTracker)(nil)
+)
+
+type flowTracker struct {
+	metadata TrackerMetadata
+	manager  *Manager
+	handle   tun.FlowHandle
+}
+
+func (t *flowTracker) Metadata() *TrackerMetadata {
+	return &t.metadata
+}
+
+func (t *flowTracker) AttachFlow(handle tun.FlowHandle) {
+	t.handle = handle
+	t.manager.join(t)
+}
+
+func (t *flowTracker) CountForward(n int) {
+	t.metadata.Upload.Add(int64(n))
+	t.manager.uploadTotal.Add(int64(n))
+}
+
+func (t *flowTracker) CountReverse(n int) {
+	t.metadata.Download.Add(int64(n))
+	t.manager.downloadTotal.Add(int64(n))
+}
+
+func (t *flowTracker) FlowEstablished() {
+}
+
+func (t *flowTracker) CloseFlow(reason tun.FlowCloseReason) {
+	t.manager.leave(t)
+}
+
+func (t *flowTracker) Close() error {
+	handle := t.handle
+	if handle != nil {
+		handle.CloseFlow()
+	} else {
+		t.manager.leave(t)
+	}
+	return nil
+}
+
 type packetConnTracker struct {
 	N.PacketConn
 	metadata TrackerMetadata