diff --git a/cmd/main.go b/cmd/main.go index f903c4e..72cacef 100644 --- a/cmd/main.go +++ b/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.RegisterUDPConnHandler(proxy.NewUDPHandler(proxyHost, proxyPort, *args.UdpTimeout, fakeDNS, monitor)) + core.RegisterTCPConnHandler(proxy.NewTCPHandler(proxyHost, proxyPort)) + 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. diff --git a/component/session/server.go b/component/session/server.go index a2aed6e..65911af 100644 --- a/component/session/server.go +++ b/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 diff --git a/component/session/session.go b/component/session/session.go index b46f0d4..778c5ed 100644 --- a/component/session/session.go +++ b/component/session/session.go @@ -13,7 +13,6 @@ type Monitor interface { // METHODS AddSession(key interface{}, session *Session) - GetSession(key interface{}) *Session RemoveSession(key interface{}) } diff --git a/proxy/proxy.go b/proxy/proxy.go new file mode 100644 index 0000000..cdc997d --- /dev/null +++ b/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) + } +} diff --git a/proxy/tcp.go b/proxy/tcp.go index fabf8b6..a2ddc09 100644 --- a/proxy/tcp.go +++ b/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 { - h.monitor.RemoveSession(localConn) - } + 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} } diff --git a/proxy/udp.go b/proxy/udp.go index b0f39db..715af60 100644 --- a/proxy/udp.go +++ b/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 { - h.monitor.RemoveSession(conn) - } + removeSession(conn) } diff --git a/proxy/utils.go b/proxy/utils.go index 53253c7..a97d33b 100644 --- a/proxy/utils.go +++ b/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