update session package

Makefile

@@ -3,7 +3,7 @@ BUILD_DIR=$(shell pwd)/build
 CMD_DIR=$(shell pwd)/cmd
 VERSION=$(shell git describe --tags --long || echo "unknown version")
 BUILD_TIME=$(shell date -u)
-BUILD_TAGS='fakeDNS stats'
+BUILD_TAGS='fakedns session'
 GOBUILD=go build -ldflags '-s -w -X "github.com/xjasonlyu/tun2socks/constant.Version=$(VERSION)"'
 
 all: build

cmd/main.go

@@ -12,25 +12,26 @@ import (
 	"syscall"
 	"time"
 
+	D "github.com/xjasonlyu/tun2socks/component/fakedns"
+	S "github.com/xjasonlyu/tun2socks/component/session"
 	C "github.com/xjasonlyu/tun2socks/constant"
 	"github.com/xjasonlyu/tun2socks/core"
 	"github.com/xjasonlyu/tun2socks/filter"
 	"github.com/xjasonlyu/tun2socks/log"
 	"github.com/xjasonlyu/tun2socks/proxy"
 	"github.com/xjasonlyu/tun2socks/tun"
-
-	D "github.com/xjasonlyu/tun2socks/component/fakedns"
-	S "github.com/xjasonlyu/tun2socks/component/stats"
 )
 
 const MTU = 1500
 
 var (
 	args = new(CmdArgs)
-	postFlagsInitFn []func()
+
+	// Modules init func
+	registeredInitFn []func()
 
 	fakeDNS D.FakeDNS
-	sessionStater S.SessionStater
+	monitor S.Monitor
 )
 
 type CmdArgs struct {
@@ -53,13 +54,13 @@ type CmdArgs struct {
 	FakeDNSAddr   *string
 	FakeDNSHosts  *string
 
-	// Stats
+	// Session Stats
-	Stats     *bool
+	EnableStats *bool
 	StatsAddr   *string
 }
 
-func addPostFlagsInitFn(fn func()) {
+func registerInitFn(fn func()) {
-	postFlagsInitFn = append(postFlagsInitFn, fn)
+	registeredInitFn = append(registeredInitFn, fn)
 }
 
 func init() {
@@ -109,7 +110,7 @@ func main() {
 	}
 
 	// Initialization modules
-	for _, fn := range postFlagsInitFn {
+	for _, fn := range registeredInitFn {
 		if fn != nil {
 			fn()
 		}
@@ -136,8 +137,8 @@ func main() {
 	lwipWriter = filter.NewICMPFilter(lwipWriter).(io.Writer)
 
 	// Register TCP and UDP handlers to handle accepted connections.
-	core.RegisterTCPConnHandler(proxy.NewTCPHandler(proxyHost, proxyPort, fakeDNS, sessionStater))
+	core.RegisterTCPConnHandler(proxy.NewTCPHandler(proxyHost, proxyPort, fakeDNS, monitor))
-	core.RegisterUDPConnHandler(proxy.NewUDPHandler(proxyHost, proxyPort, *args.UdpTimeout, fakeDNS, sessionStater))
+	core.RegisterUDPConnHandler(proxy.NewUDPHandler(proxyHost, proxyPort, *args.UdpTimeout, fakeDNS, monitor))
 
 	// Register an output callback to write packets output from lwip stack to tun
 	// device, output function should be set before input any packets.
@@ -164,7 +165,7 @@ func main() {
 	}
 
 	// Stop session stater
-	if sessionStater != nil {
+	if monitor != nil {
-		sessionStater.Stop()
+		monitor.Stop()
 	}
 }

cmd/main_fakedns.go

@@ -1,4 +1,4 @@
-// +build fakeDNS
+// +build fakedns
 
 package main
 
@@ -15,7 +15,7 @@ func init() {
 	args.FakeIPRange = flag.String("fakeIPRange", "198.18.0.0/15", "Fake IP CIDR range for DNS")
 	args.FakeDNSHosts = flag.String("fakeDNSHosts", "", "DNS hosts mapping, e.g. 'example.com=1.1.1.1,example.net=2.2.2.2'")
 
-	addPostFlagsInitFn(func() {
+	registerInitFn(func() {
 		if *args.EnableFakeDNS {
 			var err error
 			fakeDNS, err = fakedns.NewServer(*args.FakeIPRange, *args.FakeDNSHosts)
@@ -28,7 +28,7 @@ func init() {
 
 			// Start fakeDNS server
 			if err := fakeDNS.Start(); err != nil {
-				log.Fatalf("Start fake DNS server failed: %v", err)
+				log.Fatalf("Start fake DNS failed: %v", err)
 			}
 			log.Infof("Fake DNS serving at %v", fakedns.ServeAddr)
 		} else {

cmd/main_session.go

@@ -0,0 +1,32 @@
+// +build session
+
+package main
+
+import (
+	"flag"
+
+	"github.com/xjasonlyu/tun2socks/component/session"
+	"github.com/xjasonlyu/tun2socks/log"
+)
+
+func init() {
+	args.EnableStats = flag.Bool("stats", false, "Enable session statistics monitor")
+	args.StatsAddr = flag.String("statsAddr", "localhost:6001", "Listen address of session monitor, open in your browser to view statistics")
+
+	registerInitFn(func() {
+		if *args.EnableStats {
+			monitor = session.NewServer()
+
+			// Set stats variables
+			session.ServeAddr = *args.StatsAddr
+
+			// Start session stater
+			if err := monitor.Start(); err != nil {
+				log.Fatalf("Start session monitor failed: %v", err)
+			}
+			log.Infof("Session monitor serving at %v", session.ServeAddr)
+		} else {
+			monitor = nil
+		}
+	})
+}

cmd/main_stats.go

@@ -1,32 +0,0 @@
-// +build stats
-
-package main
-
-import (
-	"flag"
-
-	"github.com/xjasonlyu/tun2socks/component/stats/session"
-	"github.com/xjasonlyu/tun2socks/log"
-)
-
-func init() {
-	args.Stats = flag.Bool("stats", false, "Enable session statistics")
-	args.StatsAddr = flag.String("statsAddr", "localhost:6001", "Listen address of stats, open in your browser to view statistics")
-
-	addPostFlagsInitFn(func() {
-		if *args.Stats {
-			sessionStater = session.NewSimpleSessionStater()
-
-			// Set stats variables
-			session.ServeAddr = *args.StatsAddr
-
-			// Start session stater
-			if err := sessionStater.Start(); err != nil {
-				log.Fatalf("Start session stater failed: %v", err)
-			}
-			log.Infof("Session stater serving at %v", session.ServeAddr)
-		} else {
-			sessionStater = nil
-		}
-	})
-}

component/session/server.go

@@ -19,38 +19,38 @@ const maxCompletedSessions = 100
 
 var (
 	ServeAddr = "localhost:6001"
-	ServePath = "/stats/session/plain"
+	ServePath = "/session/plain"
 )
 
-type simpleSessionStater struct {
+type Server struct {
 	server                *http.Server
 	activeSessionMap      sync.Map
 	completedSessionQueue *queue.Queue
 }
 
-func NewSimpleSessionStater() *simpleSessionStater {
+func NewServer() *Server {
-	return &simpleSessionStater{
+	return &Server{
 		completedSessionQueue: queue.New(maxCompletedSessions),
 	}
 }
 
-func (s *simpleSessionStater) sessionStatsHandler(resp http.ResponseWriter, req *http.Request) {
+func (s *Server) handler(resp http.ResponseWriter, req *http.Request) {
 	// Slice of active sessions
-	var activeSessions []*C.Session
+	var activeSessions []*Session
 	s.activeSessionMap.Range(func(key, value interface{}) bool {
-		activeSessions = append(activeSessions, value.(*C.Session))
+		activeSessions = append(activeSessions, value.(*Session))
 		return true
 	})
 
 	// Slice of completed sessions
-	var completedSessions []*C.Session
+	var completedSessions []*Session
 	for _, item := range s.completedSessionQueue.Copy() {
-		if sess, ok := item.(*C.Session); ok {
+		if sess, ok := item.(*Session); ok {
 			completedSessions = append(completedSessions, sess)
 		}
 	}
 
-	tablePrint := func(w io.Writer, sessions []*C.Session) {
+	tablePrint := func(w io.Writer, sessions []*Session) {
 		// Sort by session start time.
 		sort.Slice(sessions, func(i, j int) bool {
 			return sessions[i].SessionStart.Sub(sessions[j].SessionStart) < 0
@@ -97,7 +97,7 @@ table, th, td {
 	_ = w.Flush()
 }
 
-func (s *simpleSessionStater) Start() error {
+func (s *Server) Start() error {
 	_, port, err := net.SplitHostPort(ServeAddr)
 	if port == "0" || port == "" || err != nil {
 		return errors.New("address format error")
@@ -117,7 +117,7 @@ func (s *simpleSessionStater) Start() error {
 	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
 		http.Redirect(w, r, ServePath, 301)
 	})
-	mux.HandleFunc(ServePath, s.sessionStatsHandler)
+	mux.HandleFunc(ServePath, s.handler)
 	s.server = &http.Server{Addr: ServeAddr, Handler: mux}
 	go func() {
 		s.server.Serve(c)
@@ -126,22 +126,22 @@ func (s *simpleSessionStater) Start() error {
 	return nil
 }
 
-func (s *simpleSessionStater) Stop() error {
+func (s *Server) Stop() error {
 	return s.server.Close()
 }
 
-func (s *simpleSessionStater) AddSession(key interface{}, session *C.Session) {
+func (s *Server) AddSession(key interface{}, session *Session) {
 	s.activeSessionMap.Store(key, session)
 }
 
-func (s *simpleSessionStater) GetSession(key interface{}) *C.Session {
+func (s *Server) GetSession(key interface{}) *Session {
 	if sess, ok := s.activeSessionMap.Load(key); ok {
-		return sess.(*C.Session)
+		return sess.(*Session)
 	}
 	return nil
 }
 
-func (s *simpleSessionStater) RemoveSession(key interface{}) {
+func (s *Server) RemoveSession(key interface{}) {
 	if sess, ok := s.activeSessionMap.Load(key); ok {
 		// move to completed sessions
 		s.completedSessionQueue.Put(sess)

component/session/session.go

@@ -0,0 +1,90 @@
+package session
+
+import (
+	"net"
+	"sync"
+	"sync/atomic"
+	"time"
+)
+
+type Monitor interface {
+	Start() error
+	Stop() error
+
+	// METHODS
+	AddSession(key interface{}, session *Session)
+	GetSession(key interface{}) *Session
+	RemoveSession(key interface{})
+}
+
+type Session struct {
+	Process       string
+	Network       string
+	DialerAddr    string
+	ClientAddr    string
+	TargetAddr    string
+	UploadBytes   int64
+	DownloadBytes int64
+	SessionStart  time.Time
+	SessionClose  time.Time
+}
+
+// Track SessionConn
+type Conn struct {
+	*Session
+	net.Conn
+	once sync.Once
+}
+
+func (c *Conn) Read(b []byte) (n int, err error) {
+	n, err = c.Conn.Read(b)
+	if n > 0 {
+		atomic.AddInt64(&c.DownloadBytes, int64(n))
+	}
+	return
+}
+
+func (c *Conn) Write(b []byte) (n int, err error) {
+	n, err = c.Conn.Write(b)
+	if n > 0 {
+		atomic.AddInt64(&c.UploadBytes, int64(n))
+	}
+	return
+}
+
+func (c *Conn) Close() error {
+	c.once.Do(func() {
+		c.SessionClose = time.Now()
+	})
+	return c.Conn.Close()
+}
+
+// Track SessionPacketConn
+type PacketConn struct {
+	*Session
+	net.PacketConn
+	once sync.Once
+}
+
+func (c *PacketConn) ReadFrom(b []byte) (n int, addr net.Addr, err error) {
+	n, addr, err = c.PacketConn.ReadFrom(b)
+	if n > 0 {
+		atomic.AddInt64(&c.DownloadBytes, int64(n))
+	}
+	return
+}
+
+func (c *PacketConn) WriteTo(b []byte, addr net.Addr) (n int, err error) {
+	n, err = c.PacketConn.WriteTo(b, addr)
+	if n > 0 {
+		atomic.AddInt64(&c.UploadBytes, int64(n))
+	}
+	return
+}
+
+func (c *PacketConn) Close() error {
+	c.once.Do(func() {
+		c.SessionClose = time.Now()
+	})
+	return c.PacketConn.Close()
+}

component/stats/stats.go

@@ -1,13 +0,0 @@
-package stats
-
-import (
-	C "github.com/xjasonlyu/tun2socks/constant"
-)
-
-type SessionStater interface {
-	Start() error
-	Stop() error
-	AddSession(key interface{}, session *C.Session)
-	GetSession(key interface{}) *C.Session
-	RemoveSession(key interface{})
-}

constant/session.go

@@ -1,94 +0,0 @@
-package constant
-
-import (
-	"net"
-	"sync"
-	"sync/atomic"
-	"time"
-)
-
-type Session struct {
-	Process       string
-	Network       string
-	DialerAddr    string
-	ClientAddr    string
-	TargetAddr    string
-	UploadBytes   int64
-	DownloadBytes int64
-	SessionStart  time.Time
-	SessionClose  time.Time
-}
-
-// Track SessionConn
-type SessionConn struct {
-	net.Conn
-	once    sync.Once
-	session *Session
-}
-
-func NewSessionConn(conn net.Conn, session *Session) net.Conn {
-	return &SessionConn{
-		Conn:    conn,
-		session: session,
-	}
-}
-
-func (c *SessionConn) Read(b []byte) (n int, err error) {
-	n, err = c.Conn.Read(b)
-	if n > 0 {
-		atomic.AddInt64(&c.session.DownloadBytes, int64(n))
-	}
-	return
-}
-
-func (c *SessionConn) Write(b []byte) (n int, err error) {
-	n, err = c.Conn.Write(b)
-	if n > 0 {
-		atomic.AddInt64(&c.session.UploadBytes, int64(n))
-	}
-	return
-}
-
-func (c *SessionConn) Close() error {
-	c.once.Do(func() {
-		c.session.SessionClose = time.Now()
-	})
-	return c.Conn.Close()
-}
-
-// Track SessionPacketConn
-type SessionPacketConn struct {
-	net.PacketConn
-	once    sync.Once
-	session *Session
-}
-
-func NewSessionPacketConn(conn net.PacketConn, session *Session) net.PacketConn {
-	return &SessionPacketConn{
-		PacketConn: conn,
-		session:    session,
-	}
-}
-
-func (c *SessionPacketConn) ReadFrom(b []byte) (n int, addr net.Addr, err error) {
-	n, addr, err = c.PacketConn.ReadFrom(b)
-	if n > 0 {
-		atomic.AddInt64(&c.session.DownloadBytes, int64(n))
-	}
-	return
-}
-
-func (c *SessionPacketConn) WriteTo(b []byte, addr net.Addr) (n int, err error) {
-	n, err = c.PacketConn.WriteTo(b, addr)
-	if n > 0 {
-		atomic.AddInt64(&c.session.UploadBytes, int64(n))
-	}
-	return
-}
-
-func (c *SessionPacketConn) Close() error {
-	c.once.Do(func() {
-		c.session.SessionClose = time.Now()
-	})
-	return c.PacketConn.Close()
-}

proxy/tcp.go

@@ -10,7 +10,7 @@ import (
 	"github.com/xjasonlyu/tun2socks/common/lsof"
 	"github.com/xjasonlyu/tun2socks/common/pool"
 	D "github.com/xjasonlyu/tun2socks/component/fakedns"
-	"github.com/xjasonlyu/tun2socks/component/stats"
+	S "github.com/xjasonlyu/tun2socks/component/session"
 	C "github.com/xjasonlyu/tun2socks/constant"
 	"github.com/xjasonlyu/tun2socks/core"
 	"github.com/xjasonlyu/tun2socks/log"
@@ -21,15 +21,15 @@ type tcpHandler struct {
 	proxyPort int
 
 	fakeDNS D.FakeDNS
-	sessionStater stats.SessionStater
+	monitor S.Monitor
 }
 
-func NewTCPHandler(proxyHost string, proxyPort int, fakeDNS D.FakeDNS, sessionStater stats.SessionStater) core.TCPConnHandler {
+func NewTCPHandler(proxyHost string, proxyPort int, fakeDNS D.FakeDNS, monitor S.Monitor) core.TCPConnHandler {
 	return &tcpHandler{
 		proxyHost: proxyHost,
 		proxyPort: proxyPort,
 		fakeDNS:   fakeDNS,
-		sessionStater: sessionStater,
+		monitor:   monitor,
 	}
 }
 
@@ -77,8 +77,8 @@ func (h *tcpHandler) relay(localConn, remoteConn net.Conn) {
 	wg.Wait() // Wait for Up Link done
 
 	// Remove session
-	if h.sessionStater != nil {
+	if h.monitor != nil {
-		h.sessionStater.RemoveSession(localConn)
+		h.monitor.RemoveSession(localConn)
 	}
 }
 
@@ -104,8 +104,8 @@ func (h *tcpHandler) Handle(conn net.Conn, target *net.TCPAddr) error {
 
 	// Get name of the process
 	var process = lsof.GetProcessName(localConn.LocalAddr())
-	if h.sessionStater != nil {
+	if h.monitor != nil {
-		sess := &C.Session{
+		session := &S.Session{
 			Process:       process,
 			Network:       localConn.LocalAddr().Network(),
 			DialerAddr:    remoteConn.LocalAddr().String(),
@@ -115,9 +115,9 @@ func (h *tcpHandler) Handle(conn net.Conn, target *net.TCPAddr) error {
 			DownloadBytes: 0,
 			SessionStart:  time.Now(),
 		}
-		h.sessionStater.AddSession(localConn, sess)
+		h.monitor.AddSession(localConn, session)
 
-		remoteConn = C.NewSessionConn(remoteConn, sess)
+		remoteConn = &S.Conn{Session: session, Conn: remoteConn}
 	}
 
 	// Set keepalive

proxy/udp.go

@@ -10,7 +10,7 @@ import (
 	"github.com/xjasonlyu/tun2socks/common/lsof"
 	"github.com/xjasonlyu/tun2socks/common/pool"
 	D "github.com/xjasonlyu/tun2socks/component/fakedns"
-	"github.com/xjasonlyu/tun2socks/component/stats"
+	S "github.com/xjasonlyu/tun2socks/component/session"
 	C "github.com/xjasonlyu/tun2socks/constant"
 	"github.com/xjasonlyu/tun2socks/core"
 	"github.com/xjasonlyu/tun2socks/log"
@@ -25,15 +25,15 @@ type udpHandler struct {
 	remoteConnMap sync.Map
 
 	fakeDNS D.FakeDNS
-	sessionStater stats.SessionStater
+	monitor S.Monitor
 }
 
-func NewUDPHandler(proxyHost string, proxyPort int, timeout time.Duration, fakeDNS D.FakeDNS, sessionStater stats.SessionStater) core.UDPConnHandler {
+func NewUDPHandler(proxyHost string, proxyPort int, timeout time.Duration, fakeDNS D.FakeDNS, monitor S.Monitor) core.UDPConnHandler {
 	return &udpHandler{
 		proxyHost: proxyHost,
 		proxyPort: proxyPort,
 		fakeDNS:   fakeDNS,
-		sessionStater: sessionStater,
+		monitor:   monitor,
 		timeout:   timeout,
 	}
 }
@@ -82,8 +82,8 @@ func (h *udpHandler) Connect(conn core.UDPConn, target *net.UDPAddr) error {
 
 	// Get name of the process
 	var process = lsof.GetProcessName(conn.LocalAddr())
-	if h.sessionStater != nil {
+	if h.monitor != nil {
-		sess := &C.Session{
+		session := &S.Session{
 			Process:       process,
 			Network:       conn.LocalAddr().Network(),
 			DialerAddr:    remoteConn.LocalAddr().String(),
@@ -93,9 +93,9 @@ func (h *udpHandler) Connect(conn core.UDPConn, target *net.UDPAddr) error {
 			DownloadBytes: 0,
 			SessionStart:  time.Now(),
 		}
-		h.sessionStater.AddSession(conn, sess)
+		h.monitor.AddSession(conn, session)
 
-		remoteConn = C.NewSessionPacketConn(remoteConn, sess)
+		remoteConn = &S.PacketConn{Session: session, PacketConn: remoteConn}
 	}
 
 	h.remoteAddrMap.Store(conn, remoteAddr)
@@ -152,7 +152,7 @@ func (h *udpHandler) Close(conn core.UDPConn) {
 	conn.Close()
 
 	// Remove session
-	if h.sessionStater != nil {
+	if h.monitor != nil {
-		h.sessionStater.RemoveSession(conn)
+		h.monitor.RemoveSession(conn)
 	}
 }

scripts/run_socks_linux.sh

@@ -11,4 +11,4 @@ config_route() {
 }
 
 config_route &
-sudo ./build/tun2socks -tunAddr 240.0.0.2 -tunGw 240.0.0.1 -proxyServer 192.168.1.1:1080 -fakeDNS -loglevel info -stats
+sudo ./build/tun2socks -tunAddr 240.0.0.2 -tunGw 240.0.0.1 -proxyServer 192.168.1.1:1080 -fakeDNS

scripts/run_socks_macos.sh

@@ -7,4 +7,4 @@ config_route() {
 }
 
 config_route &
-sudo ./build/tun2socks -tunAddr 240.0.0.2 -tunGw 240.0.0.1 -proxyServer 192.168.1.1:1080 -fakeDNS -loglevel info -stats
+sudo ./build/tun2socks -tunAddr 240.0.0.2 -tunGw 240.0.0.1 -proxyServer 192.168.1.1:1080 -fakeDNS