diff --git a/stats/connections.go b/stats/connections.go index d25f08c..818b98e 100755 --- a/stats/connections.go +++ b/stats/connections.go @@ -7,7 +7,7 @@ import ( "strconv" "time" - "github.com/xjasonlyu/tun2socks/tunnel" + "github.com/xjasonlyu/tun2socks/tunnel/statistic" "github.com/go-chi/chi" "github.com/go-chi/render" @@ -26,7 +26,7 @@ func connectionRouter() http.Handler { func getConnections(w http.ResponseWriter, r *http.Request) { if !websocket.IsWebSocketUpgrade(r) { - snapshot := tunnel.DefaultManager.Snapshot() + snapshot := statistic.DefaultManager.Snapshot() render.JSON(w, r, snapshot) return } @@ -52,7 +52,7 @@ func getConnections(w http.ResponseWriter, r *http.Request) { buf := &bytes.Buffer{} sendSnapshot := func() error { buf.Reset() - snapshot := tunnel.DefaultManager.Snapshot() + snapshot := statistic.DefaultManager.Snapshot() if err := json.NewEncoder(buf).Encode(snapshot); err != nil { return err } @@ -75,7 +75,7 @@ func getConnections(w http.ResponseWriter, r *http.Request) { func closeConnection(w http.ResponseWriter, r *http.Request) { id := chi.URLParam(r, "id") - snapshot := tunnel.DefaultManager.Snapshot() + snapshot := statistic.DefaultManager.Snapshot() for _, c := range snapshot.Connections { if id == c.ID() { _ = c.Close() @@ -86,7 +86,7 @@ func closeConnection(w http.ResponseWriter, r *http.Request) { } func closeAllConnections(w http.ResponseWriter, r *http.Request) { - snapshot := tunnel.DefaultManager.Snapshot() + snapshot := statistic.DefaultManager.Snapshot() for _, c := range snapshot.Connections { _ = c.Close() } diff --git a/stats/server.go b/stats/server.go index 5396364..e0648ad 100755 --- a/stats/server.go +++ b/stats/server.go @@ -10,7 +10,7 @@ import ( "github.com/xjasonlyu/tun2socks/constant" "github.com/xjasonlyu/tun2socks/log" - "github.com/xjasonlyu/tun2socks/tunnel" + "github.com/xjasonlyu/tun2socks/tunnel/statistic" "github.com/go-chi/chi" "github.com/go-chi/cors" @@ -176,7 +176,7 @@ func traffic(w http.ResponseWriter, r *http.Request) { tick := time.NewTicker(time.Second) defer tick.Stop() - t := tunnel.DefaultManager + t := statistic.DefaultManager buf := &bytes.Buffer{} var err error for range tick.C { diff --git a/tunnel/manager.go b/tunnel/manager.go deleted file mode 100755 index 261e685..0000000 --- a/tunnel/manager.go +++ /dev/null @@ -1,22 +0,0 @@ -package tunnel - -import ( - "net" - - "github.com/xjasonlyu/tun2socks/common/adapter" - "github.com/xjasonlyu/tun2socks/component/manager" -) - -var ( - // DefaultManager is the default traffic and connections - // manager used by tunnel. - DefaultManager = manager.New() -) - -func newTCPTracker(conn net.Conn, metadata *adapter.Metadata) net.Conn { - return manager.NewTCPTracker(conn, metadata, DefaultManager) -} - -func newUDPTracker(conn net.PacketConn, metadata *adapter.Metadata) net.PacketConn { - return manager.NewUDPTracker(conn, metadata, DefaultManager) -} diff --git a/component/manager/manager.go b/tunnel/statistic/manager.go similarity index 92% rename from component/manager/manager.go rename to tunnel/statistic/manager.go index 9a0d54f..df85b50 100755 --- a/component/manager/manager.go +++ b/tunnel/statistic/manager.go @@ -1,5 +1,4 @@ -// Package manager provides statistic management of connections. -package manager +package statistic import ( "sync" @@ -8,6 +7,21 @@ import ( "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 @@ -18,20 +32,6 @@ type Manager struct { downloadTotal *atomic.Int64 } -func New() *Manager { - manager := &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 manager.handle() - - return manager -} - func (m *Manager) Join(c tracker) { m.connections.Store(c.ID(), c) } diff --git a/component/manager/tracker.go b/tunnel/statistic/tracker.go similarity index 99% rename from component/manager/tracker.go rename to tunnel/statistic/tracker.go index c35e58a..5315993 100755 --- a/component/manager/tracker.go +++ b/tunnel/statistic/tracker.go @@ -1,4 +1,4 @@ -package manager +package statistic import ( "net" diff --git a/tunnel/tcp.go b/tunnel/tcp.go index 51ebb28..5d8f04b 100755 --- a/tunnel/tcp.go +++ b/tunnel/tcp.go @@ -11,12 +11,17 @@ import ( "github.com/xjasonlyu/tun2socks/common/pool" "github.com/xjasonlyu/tun2socks/log" "github.com/xjasonlyu/tun2socks/proxy" + "github.com/xjasonlyu/tun2socks/tunnel/statistic" ) const ( tcpWaitTimeout = 5 * time.Second ) +func newTCPTracker(conn net.Conn, metadata *adapter.Metadata) net.Conn { + return statistic.NewTCPTracker(conn, metadata, statistic.DefaultManager) +} + func handleTCP(localConn adapter.TCPConn) { defer localConn.Close() diff --git a/tunnel/udp.go b/tunnel/udp.go index 875ebf7..01d99a0 100755 --- a/tunnel/udp.go +++ b/tunnel/udp.go @@ -12,6 +12,7 @@ import ( "github.com/xjasonlyu/tun2socks/component/nat" "github.com/xjasonlyu/tun2socks/log" "github.com/xjasonlyu/tun2socks/proxy" + "github.com/xjasonlyu/tun2socks/tunnel/statistic" ) const ( @@ -24,6 +25,10 @@ var ( natTable = nat.NewTable() ) +func newUDPTracker(conn net.PacketConn, metadata *adapter.Metadata) net.PacketConn { + return statistic.NewUDPTracker(conn, metadata, statistic.DefaultManager) +} + func handleUDP(packet adapter.UDPPacket) { metadata := packet.Metadata() if !metadata.Valid() {