Refactor: make inbound request contextual

tunnel/statistic/manager.go

@@ -0,0 +1,95 @@
+package statistic
+
+import (
+	"sync"
+	"time"
+
+	"go.uber.org/atomic"
+)
+
+var DefaultManager *Manager
+
+func init() {
+	DefaultManager = &Manager{
+		uploadTemp:    atomic.NewInt64(0),
+		downloadTemp:  atomic.NewInt64(0),
+		uploadBlip:    atomic.NewInt64(0),
+		downloadBlip:  atomic.NewInt64(0),
+		uploadTotal:   atomic.NewInt64(0),
+		downloadTotal: atomic.NewInt64(0),
+	}
+
+	go DefaultManager.handle()
+}
+
+type Manager struct {
+	connections   sync.Map
+	uploadTemp    *atomic.Int64
+	downloadTemp  *atomic.Int64
+	uploadBlip    *atomic.Int64
+	downloadBlip  *atomic.Int64
+	uploadTotal   *atomic.Int64
+	downloadTotal *atomic.Int64
+}
+
+func (m *Manager) Join(c tracker) {
+	m.connections.Store(c.ID(), c)
+}
+
+func (m *Manager) Leave(c tracker) {
+	m.connections.Delete(c.ID())
+}
+
+func (m *Manager) PushUploaded(size int64) {
+	m.uploadTemp.Add(size)
+	m.uploadTotal.Add(size)
+}
+
+func (m *Manager) PushDownloaded(size int64) {
+	m.downloadTemp.Add(size)
+	m.downloadTotal.Add(size)
+}
+
+func (m *Manager) Now() (up int64, down int64) {
+	return m.uploadBlip.Load(), m.downloadBlip.Load()
+}
+
+func (m *Manager) Snapshot() *Snapshot {
+	connections := []tracker{}
+	m.connections.Range(func(key, value interface{}) bool {
+		connections = append(connections, value.(tracker))
+		return true
+	})
+
+	return &Snapshot{
+		UploadTotal:   m.uploadTotal.Load(),
+		DownloadTotal: m.downloadTotal.Load(),
+		Connections:   connections,
+	}
+}
+
+func (m *Manager) ResetStatistic() {
+	m.uploadTemp.Store(0)
+	m.uploadBlip.Store(0)
+	m.uploadTotal.Store(0)
+	m.downloadTemp.Store(0)
+	m.downloadBlip.Store(0)
+	m.downloadTotal.Store(0)
+}
+
+func (m *Manager) handle() {
+	ticker := time.NewTicker(time.Second)
+
+	for range ticker.C {
+		m.uploadBlip.Store(m.uploadTemp.Load())
+		m.uploadTemp.Store(0)
+		m.downloadBlip.Store(m.downloadTemp.Load())
+		m.downloadTemp.Store(0)
+	}
+}
+
+type Snapshot struct {
+	DownloadTotal int64     `json:"downloadTotal"`
+	UploadTotal   int64     `json:"uploadTotal"`
+	Connections   []tracker `json:"connections"`
+}

tunnel/statistic/tracker.go

@@ -0,0 +1,141 @@
+package statistic
+
+import (
+	"net"
+	"time"
+
+	C "github.com/Dreamacro/clash/constant"
+
+	"github.com/gofrs/uuid"
+	"go.uber.org/atomic"
+)
+
+type tracker interface {
+	ID() string
+	Close() error
+}
+
+type trackerInfo struct {
+	UUID          uuid.UUID     `json:"id"`
+	Metadata      *C.Metadata   `json:"metadata"`
+	UploadTotal   *atomic.Int64 `json:"upload"`
+	DownloadTotal *atomic.Int64 `json:"download"`
+	Start         time.Time     `json:"start"`
+	Chain         C.Chain       `json:"chains"`
+	Rule          string        `json:"rule"`
+	RulePayload   string        `json:"rulePayload"`
+}
+
+type tcpTracker struct {
+	C.Conn `json:"-"`
+	*trackerInfo
+	manager *Manager
+}
+
+func (tt *tcpTracker) ID() string {
+	return tt.UUID.String()
+}
+
+func (tt *tcpTracker) Read(b []byte) (int, error) {
+	n, err := tt.Conn.Read(b)
+	download := int64(n)
+	tt.manager.PushDownloaded(download)
+	tt.DownloadTotal.Add(download)
+	return n, err
+}
+
+func (tt *tcpTracker) Write(b []byte) (int, error) {
+	n, err := tt.Conn.Write(b)
+	upload := int64(n)
+	tt.manager.PushUploaded(upload)
+	tt.UploadTotal.Add(upload)
+	return n, err
+}
+
+func (tt *tcpTracker) Close() error {
+	tt.manager.Leave(tt)
+	return tt.Conn.Close()
+}
+
+func NewTCPTracker(conn C.Conn, manager *Manager, metadata *C.Metadata, rule C.Rule) *tcpTracker {
+	uuid, _ := uuid.NewV4()
+
+	t := &tcpTracker{
+		Conn:    conn,
+		manager: manager,
+		trackerInfo: &trackerInfo{
+			UUID:          uuid,
+			Start:         time.Now(),
+			Metadata:      metadata,
+			Chain:         conn.Chains(),
+			Rule:          "",
+			UploadTotal:   atomic.NewInt64(0),
+			DownloadTotal: atomic.NewInt64(0),
+		},
+	}
+
+	if rule != nil {
+		t.trackerInfo.Rule = rule.RuleType().String()
+		t.trackerInfo.RulePayload = rule.Payload()
+	}
+
+	manager.Join(t)
+	return t
+}
+
+type udpTracker struct {
+	C.PacketConn `json:"-"`
+	*trackerInfo
+	manager *Manager
+}
+
+func (ut *udpTracker) ID() string {
+	return ut.UUID.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)
+	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)
+	ut.UploadTotal.Add(upload)
+	return n, err
+}
+
+func (ut *udpTracker) Close() error {
+	ut.manager.Leave(ut)
+	return ut.PacketConn.Close()
+}
+
+func NewUDPTracker(conn C.PacketConn, manager *Manager, metadata *C.Metadata, rule C.Rule) *udpTracker {
+	uuid, _ := uuid.NewV4()
+
+	ut := &udpTracker{
+		PacketConn: conn,
+		manager:    manager,
+		trackerInfo: &trackerInfo{
+			UUID:          uuid,
+			Start:         time.Now(),
+			Metadata:      metadata,
+			Chain:         conn.Chains(),
+			Rule:          "",
+			UploadTotal:   atomic.NewInt64(0),
+			DownloadTotal: atomic.NewInt64(0),
+		},
+	}
+
+	if rule != nil {
+		ut.trackerInfo.Rule = rule.RuleType().String()
+		ut.trackerInfo.RulePayload = rule.Payload()
+	}
+
+	manager.Join(ut)
+	return ut
+}