update register

cmd/main.go

@@ -137,9 +137,12 @@ func main() {
 	// Wrap a writer to delay ICMP packets
 	lwipWriter = filter.NewICMPFilter(lwipWriter).(io.Writer)
 
+	// Register modules to proxy
+	proxy.RegisterFakeDNS(fakeDNS)
+	proxy.RegisterMonitor(monitor)
 	// Register TCP and UDP handlers to handle accepted connections.
-	core.RegisterTCPConnHandler(proxy.NewTCPHandler(proxyHost, proxyPort, fakeDNS, monitor))
+	core.RegisterTCPConnHandler(proxy.NewTCPHandler(proxyHost, proxyPort))
-	core.RegisterUDPConnHandler(proxy.NewUDPHandler(proxyHost, proxyPort, *args.UdpTimeout, fakeDNS, monitor))
+	core.RegisterUDPConnHandler(proxy.NewUDPHandler(proxyHost, proxyPort, *args.UdpTimeout))
 
 	// Register an output callback to write packets output from lwip stack to tun
 	// device, output function should be set before input any packets.

component/session/server.go

@@ -134,13 +134,6 @@ func (s *Server) AddSession(key interface{}, session *Session) {
 	s.activeSessionMap.Store(key, session)
 }
 
-func (s *Server) GetSession(key interface{}) *Session {
-	if sess, ok := s.activeSessionMap.Load(key); ok {
-		return sess.(*Session)
-	}
-	return nil
-}
-
 func (s *Server) RemoveSession(key interface{}) {
 	if sess, ok := s.activeSessionMap.Load(key); ok {
 		// move to completed sessions

component/session/session.go

@@ -13,7 +13,6 @@ type Monitor interface {
 
 	// METHODS
 	AddSession(key interface{}, session *Session)
-	GetSession(key interface{}) *Session
 	RemoveSession(key interface{})
 }
 

proxy/proxy.go

@@ -0,0 +1,58 @@
+package proxy
+
+import (
+	"errors"
+	"net"
+
+	D "github.com/xjasonlyu/tun2socks/component/fakedns"
+	S "github.com/xjasonlyu/tun2socks/component/session"
+)
+
+var (
+	fakeDNS D.FakeDNS
+	monitor S.Monitor
+)
+
+func RegisterFakeDNS(d D.FakeDNS) {
+	fakeDNS = d
+}
+
+func RegisterMonitor(m S.Monitor) {
+	monitor = m
+}
+
+// DNS lookup
+func lookupHost(target net.Addr) (targetHost string, err error) {
+	var targetIP net.IP
+	switch addr := target.(type) {
+	case *net.TCPAddr:
+		targetIP = addr.IP
+	case *net.UDPAddr:
+		targetIP = addr.IP
+	default:
+		err = errors.New("invalid target type")
+		return
+	}
+
+	targetHost = targetIP.String()
+	// Replace with a domain name if target address IP is a fake IP.
+	if fakeDNS != nil {
+		if host, exist := fakeDNS.IPToHost(targetIP); exist {
+			targetHost = host
+		}
+	}
+	return
+}
+
+// Session Operation
+func addSession(key interface{}, session *S.Session) {
+	if monitor != nil {
+		monitor.AddSession(key, session)
+	}
+}
+
+func removeSession(key interface{}) {
+	if monitor != nil {
+		monitor.RemoveSession(key)
+	}
+}

proxy/tcp.go

@@ -12,24 +12,18 @@ import (
 	"github.com/xjasonlyu/tun2socks/core"
 	"github.com/xjasonlyu/tun2socks/log"
 
-	D "github.com/xjasonlyu/tun2socks/component/fakedns"
 	S "github.com/xjasonlyu/tun2socks/component/session"
 )
 
 type tcpHandler struct {
 	proxyHost string
 	proxyPort int
-
-	fakeDNS D.FakeDNS
-	monitor S.Monitor
 }
 
-func NewTCPHandler(proxyHost string, proxyPort int, fakeDNS D.FakeDNS, monitor S.Monitor) core.TCPConnHandler {
+func NewTCPHandler(proxyHost string, proxyPort int) core.TCPConnHandler {
 	return &tcpHandler{
 		proxyHost: proxyHost,
 		proxyPort: proxyPort,
-		fakeDNS:   fakeDNS,
-		monitor:   monitor,
 	}
 }
 
@@ -77,9 +71,7 @@ func (h *tcpHandler) relay(localConn, remoteConn net.Conn) {
 	wg.Wait() // Wait for Up Link done
 
 	// Remove session
-	if h.monitor != nil {
+	removeSession(localConn)
-		h.monitor.RemoveSession(localConn)
-	}
 }
 
 func (h *tcpHandler) Handle(conn net.Conn, target *net.TCPAddr) error {
@@ -87,7 +79,7 @@ func (h *tcpHandler) Handle(conn net.Conn, target *net.TCPAddr) error {
 	var localConn = conn
 
 	// Lookup fakeDNS host record
-	targetHost, err := lookupHost(h.fakeDNS, target)
+	targetHost, err := lookupHost(target)
 	if err != nil {
 		log.Warnf("lookup target host error: %v", err)
 		return err
@@ -104,7 +96,7 @@ func (h *tcpHandler) Handle(conn net.Conn, target *net.TCPAddr) error {
 
 	// Get name of the process
 	var process = lsof.GetProcessName(localConn.LocalAddr())
-	if h.monitor != nil {
+	if monitor != nil {
 		session := &S.Session{
 			Process:       process,
 			Network:       localConn.LocalAddr().Network(),
@@ -115,8 +107,7 @@ func (h *tcpHandler) Handle(conn net.Conn, target *net.TCPAddr) error {
 			DownloadBytes: 0,
 			SessionStart:  time.Now(),
 		}
-		h.monitor.AddSession(localConn, session)
+		addSession(localConn, session)
-
 		remoteConn = &S.Conn{Session: session, Conn: remoteConn}
 	}
 

proxy/udp.go

@@ -12,7 +12,6 @@ import (
 	"github.com/xjasonlyu/tun2socks/core"
 	"github.com/xjasonlyu/tun2socks/log"
 
-	D "github.com/xjasonlyu/tun2socks/component/fakedns"
 	S "github.com/xjasonlyu/tun2socks/component/session"
 )
 
@@ -23,17 +22,12 @@ type udpHandler struct {
 
 	remoteAddrMap sync.Map
 	remoteConnMap sync.Map
-
-	fakeDNS D.FakeDNS
-	monitor S.Monitor
 }
 
-func NewUDPHandler(proxyHost string, proxyPort int, timeout time.Duration, fakeDNS D.FakeDNS, monitor S.Monitor) core.UDPConnHandler {
+func NewUDPHandler(proxyHost string, proxyPort int, timeout time.Duration) core.UDPConnHandler {
 	return &udpHandler{
 		proxyHost: proxyHost,
 		proxyPort: proxyPort,
-		fakeDNS:   fakeDNS,
-		monitor:   monitor,
 		timeout:   timeout,
 	}
 }
@@ -65,7 +59,7 @@ func (h *udpHandler) fetchUDPInput(conn core.UDPConn, input net.PacketConn, addr
 
 func (h *udpHandler) Connect(conn core.UDPConn, target *net.UDPAddr) error {
 	// Lookup fakeDNS host record
-	targetHost, err := lookupHost(h.fakeDNS, target)
+	targetHost, err := lookupHost(target)
 	if err != nil {
 		log.Warnf("lookup target host error: %v", err)
 		return err
@@ -82,7 +76,7 @@ func (h *udpHandler) Connect(conn core.UDPConn, target *net.UDPAddr) error {
 
 	// Get name of the process
 	var process = lsof.GetProcessName(conn.LocalAddr())
-	if h.monitor != nil {
+	if monitor != nil {
 		session := &S.Session{
 			Process:       process,
 			Network:       conn.LocalAddr().Network(),
@@ -93,8 +87,7 @@ func (h *udpHandler) Connect(conn core.UDPConn, target *net.UDPAddr) error {
 			DownloadBytes: 0,
 			SessionStart:  time.Now(),
 		}
-		h.monitor.AddSession(conn, session)
+		addSession(conn, session)
-
 		remoteConn = &S.PacketConn{Session: session, PacketConn: remoteConn}
 	}
 
@@ -152,7 +145,5 @@ func (h *udpHandler) Close(conn core.UDPConn) {
 	conn.Close()
 
 	// Remove session
-	if h.monitor != nil {
+	removeSession(conn)
-		h.monitor.RemoveSession(conn)
-	}
 }

proxy/utils.go

@@ -1,40 +1,15 @@
 package proxy
 
 import (
-	"errors"
 	"fmt"
 	"io"
 	"io/ioutil"
 	"net"
 	"time"
 
-	D "github.com/xjasonlyu/tun2socks/component/fakedns"
 	"github.com/xjasonlyu/tun2socks/proxy/socks"
 )
 
-// DNS lookup
-func lookupHost(fakeDNS D.FakeDNS, target net.Addr) (targetHost string, err error) {
-	var targetIP net.IP
-	switch addr := target.(type) {
-	case *net.TCPAddr:
-		targetIP = addr.IP
-	case *net.UDPAddr:
-		targetIP = addr.IP
-	default:
-		err = errors.New("invalid target type")
-		return
-	}
-
-	targetHost = targetIP.String()
-	// Replace with a domain name if target address IP is a fake IP.
-	if fakeDNS != nil {
-		if host, exist := fakeDNS.IPToHost(targetIP); exist {
-			targetHost = host
-		}
-	}
-	return
-}
-
 // TCP functions
 type duplexConn interface {
 	net.Conn