Add SSM API service

service/ssmapi/traffic.go

@@ -0,0 +1,215 @@
+package ssmapi
+
+import (
+	"net"
+	"sync"
+
+	"github.com/sagernet/sing-box/adapter"
+	"github.com/sagernet/sing/common/atomic"
+	"github.com/sagernet/sing/common/bufio"
+	N "github.com/sagernet/sing/common/network"
+)
+
+var _ adapter.SSMTracker = (*TrafficManager)(nil)
+
+type TrafficManager struct {
+	globalUplink          atomic.Int64
+	globalDownlink        atomic.Int64
+	globalUplinkPackets   atomic.Int64
+	globalDownlinkPackets atomic.Int64
+	globalTCPSessions     atomic.Int64
+	globalUDPSessions     atomic.Int64
+	userAccess            sync.Mutex
+	userUplink            map[string]*atomic.Int64
+	userDownlink          map[string]*atomic.Int64
+	userUplinkPackets     map[string]*atomic.Int64
+	userDownlinkPackets   map[string]*atomic.Int64
+	userTCPSessions       map[string]*atomic.Int64
+	userUDPSessions       map[string]*atomic.Int64
+}
+
+func NewTrafficManager() *TrafficManager {
+	manager := &TrafficManager{
+		userUplink:          make(map[string]*atomic.Int64),
+		userDownlink:        make(map[string]*atomic.Int64),
+		userUplinkPackets:   make(map[string]*atomic.Int64),
+		userDownlinkPackets: make(map[string]*atomic.Int64),
+		userTCPSessions:     make(map[string]*atomic.Int64),
+		userUDPSessions:     make(map[string]*atomic.Int64),
+	}
+	return manager
+}
+
+func (s *TrafficManager) UpdateUsers(users []string) {
+	s.userAccess.Lock()
+	defer s.userAccess.Unlock()
+	newUserUplink := make(map[string]*atomic.Int64)
+	newUserDownlink := make(map[string]*atomic.Int64)
+	newUserUplinkPackets := make(map[string]*atomic.Int64)
+	newUserDownlinkPackets := make(map[string]*atomic.Int64)
+	newUserTCPSessions := make(map[string]*atomic.Int64)
+	newUserUDPSessions := make(map[string]*atomic.Int64)
+	for _, user := range users {
+		newUserUplink[user] = s.userUplinkPackets[user]
+		newUserDownlink[user] = s.userDownlinkPackets[user]
+		newUserUplinkPackets[user] = s.userUplinkPackets[user]
+		newUserDownlinkPackets[user] = s.userDownlinkPackets[user]
+		newUserTCPSessions[user] = s.userTCPSessions[user]
+		newUserUDPSessions[user] = s.userUDPSessions[user]
+	}
+	s.userUplink = newUserUplink
+	s.userDownlink = newUserDownlink
+	s.userUplinkPackets = newUserUplinkPackets
+	s.userDownlinkPackets = newUserDownlinkPackets
+	s.userTCPSessions = newUserTCPSessions
+	s.userUDPSessions = newUserUDPSessions
+}
+
+func (s *TrafficManager) userCounter(user string) (*atomic.Int64, *atomic.Int64, *atomic.Int64, *atomic.Int64, *atomic.Int64, *atomic.Int64) {
+	s.userAccess.Lock()
+	defer s.userAccess.Unlock()
+	upCounter, loaded := s.userUplink[user]
+	if !loaded {
+		upCounter = new(atomic.Int64)
+		s.userUplink[user] = upCounter
+	}
+	downCounter, loaded := s.userDownlink[user]
+	if !loaded {
+		downCounter = new(atomic.Int64)
+		s.userDownlink[user] = downCounter
+	}
+	upPacketsCounter, loaded := s.userUplinkPackets[user]
+	if !loaded {
+		upPacketsCounter = new(atomic.Int64)
+		s.userUplinkPackets[user] = upPacketsCounter
+	}
+	downPacketsCounter, loaded := s.userDownlinkPackets[user]
+	if !loaded {
+		downPacketsCounter = new(atomic.Int64)
+		s.userDownlinkPackets[user] = downPacketsCounter
+	}
+	tcpSessionsCounter, loaded := s.userTCPSessions[user]
+	if !loaded {
+		tcpSessionsCounter = new(atomic.Int64)
+		s.userTCPSessions[user] = tcpSessionsCounter
+	}
+	udpSessionsCounter, loaded := s.userUDPSessions[user]
+	if !loaded {
+		udpSessionsCounter = new(atomic.Int64)
+		s.userUDPSessions[user] = udpSessionsCounter
+	}
+	return upCounter, downCounter, upPacketsCounter, downPacketsCounter, tcpSessionsCounter, udpSessionsCounter
+}
+
+func (s *TrafficManager) TrackConnection(conn net.Conn, metadata adapter.InboundContext) net.Conn {
+	s.globalTCPSessions.Add(1)
+	var readCounter []*atomic.Int64
+	var writeCounter []*atomic.Int64
+	readCounter = append(readCounter, &s.globalUplink)
+	writeCounter = append(writeCounter, &s.globalDownlink)
+	upCounter, downCounter, _, _, tcpSessionCounter, _ := s.userCounter(metadata.User)
+	readCounter = append(readCounter, upCounter)
+	writeCounter = append(writeCounter, downCounter)
+	tcpSessionCounter.Add(1)
+	return bufio.NewInt64CounterConn(conn, readCounter, writeCounter)
+}
+
+func (s *TrafficManager) TrackPacketConnection(conn N.PacketConn, metadata adapter.InboundContext) N.PacketConn {
+	s.globalUDPSessions.Add(1)
+	var readCounter []*atomic.Int64
+	var readPacketCounter []*atomic.Int64
+	var writeCounter []*atomic.Int64
+	var writePacketCounter []*atomic.Int64
+	readCounter = append(readCounter, &s.globalUplink)
+	writeCounter = append(writeCounter, &s.globalDownlink)
+	readPacketCounter = append(readPacketCounter, &s.globalUplinkPackets)
+	writePacketCounter = append(writePacketCounter, &s.globalDownlinkPackets)
+	upCounter, downCounter, upPacketsCounter, downPacketsCounter, _, udpSessionCounter := s.userCounter(metadata.User)
+	readCounter = append(readCounter, upCounter)
+	writeCounter = append(writeCounter, downCounter)
+	readPacketCounter = append(readPacketCounter, upPacketsCounter)
+	writePacketCounter = append(writePacketCounter, downPacketsCounter)
+	udpSessionCounter.Add(1)
+	return bufio.NewInt64CounterPacketConn(conn, append(readCounter, readPacketCounter...), append(writeCounter, writePacketCounter...))
+}
+
+func (s *TrafficManager) ReadUser(user *UserObject) {
+	s.userAccess.Lock()
+	defer s.userAccess.Unlock()
+	s.readUser(user)
+}
+
+func (s *TrafficManager) readUser(user *UserObject) {
+	if counter, loaded := s.userUplink[user.UserName]; loaded {
+		user.UplinkBytes = counter.Load()
+	}
+	if counter, loaded := s.userDownlink[user.UserName]; loaded {
+		user.DownlinkBytes = counter.Load()
+	}
+	if counter, loaded := s.userUplinkPackets[user.UserName]; loaded {
+		user.UplinkPackets = counter.Load()
+	}
+	if counter, loaded := s.userDownlinkPackets[user.UserName]; loaded {
+		user.DownlinkPackets = counter.Load()
+	}
+	if counter, loaded := s.userTCPSessions[user.UserName]; loaded {
+		user.TCPSessions = counter.Load()
+	}
+	if counter, loaded := s.userUDPSessions[user.UserName]; loaded {
+		user.UDPSessions = counter.Load()
+	}
+}
+
+func (s *TrafficManager) ReadUsers(users []*UserObject) {
+	s.userAccess.Lock()
+	defer s.userAccess.Unlock()
+	for _, user := range users {
+		s.readUser(user)
+	}
+	return
+}
+
+func (s *TrafficManager) ReadGlobal() (
+	uplinkBytes int64,
+	downlinkBytes int64,
+	uplinkPackets int64,
+	downlinkPackets int64,
+	tcpSessions int64,
+	udpSessions int64,
+) {
+	return s.globalUplink.Load(),
+		s.globalDownlink.Load(),
+		s.globalUplinkPackets.Load(),
+		s.globalDownlinkPackets.Load(),
+		s.globalTCPSessions.Load(),
+		s.globalUDPSessions.Load()
+}
+
+func (s *TrafficManager) Clear() {
+	s.globalUplink.Store(0)
+	s.globalDownlink.Store(0)
+	s.globalUplinkPackets.Store(0)
+	s.globalDownlinkPackets.Store(0)
+	s.globalTCPSessions.Store(0)
+	s.globalUDPSessions.Store(0)
+	s.userAccess.Lock()
+	defer s.userAccess.Unlock()
+	for _, counter := range s.userUplink {
+		counter.Store(0)
+	}
+	for _, counter := range s.userDownlink {
+		counter.Store(0)
+	}
+	for _, counter := range s.userUplinkPackets {
+		counter.Store(0)
+	}
+	for _, counter := range s.userDownlinkPackets {
+		counter.Store(0)
+	}
+	for _, counter := range s.userTCPSessions {
+		counter.Store(0)
+	}
+	for _, counter := range s.userUDPSessions {
+		counter.Store(0)
+	}
+}