update session & fakeDNS

cmd/main.go

@@ -30,11 +30,9 @@ var (
 	description = "A tun2socks implementation written in Go."
 
 	args            = new(CmdArgs)
-
 	postFlagsInitFn []func()
 
-	lwipWriter    io.Writer
-	fakeDns       dns.FakeDns
+	fakeDNS       dns.FakeDNS
 	sessionStater stats.SessionStater
 )
 
@@ -45,7 +43,7 @@ type CmdArgs struct {
 	TunAddr  *string
 	TunGw    *string
 	TunMask  *string
-	TunDns   *string
+	TunDNS   *string
 	LogLevel *string
 
 	// Proxy
@@ -53,11 +51,11 @@ type CmdArgs struct {
 	UdpTimeout  *time.Duration
 
 	// FakeDNS
-	EnableFakeDns *bool
-	DnsCacheSize  *int
+	EnableFakeDNS *bool
+	DNSCacheSize  *int
 	FakeIPRange   *string
-	FakeDnsAddr   *string
-	FakeDnsHosts  *string
+	FakeDNSAddr   *string
+	FakeDNSHosts  *string
 
 	// Stats
 	Stats     *bool
@@ -76,7 +74,7 @@ func init() {
 	args.TunAddr = flag.String("tunAddr", "240.0.0.2", "TUN interface address")
 	args.TunGw = flag.String("tunGw", "240.0.0.1", "TUN interface gateway")
 	args.TunMask = flag.String("tunMask", "255.255.255.0", "TUN interface netmask, it should be a prefix length (a number) for IPv6 address")
-	args.TunDns = flag.String("tunDns", "1.1.1.1", "DNS resolvers for TUN interface (Windows Only)")
+	args.TunDNS = flag.String("tunDNS", "1.1.1.1", "DNS resolvers for TUN interface (Windows Only)")
 
 	// Proxy
 	args.ProxyServer = flag.String("proxyServer", "", "Proxy server address")
@@ -117,16 +115,16 @@ func main() {
 	}
 
 	// Open the tun device.
-	dnsServers := strings.Split(*args.TunDns, ",")
+	dnsServers := strings.Split(*args.TunDNS, ",")
 	tunDev, err := tun.OpenTunDevice(*args.TunName, *args.TunAddr, *args.TunGw, *args.TunMask, dnsServers)
 	if err != nil {
 		log.Fatalf("failed to open tun device: %v", err)
 	}
 
 	// Setup TCP/IP stack.
+	var lwipWriter = core.NewLWIPStack().(io.Writer)
 	// Wrap a writer to delay ICMP packets.
-	w := core.NewLWIPStack().(io.Writer)
-	lwipWriter = filter.NewICMPFilter(w).(io.Writer)
+	lwipWriter = filter.NewICMPFilter(lwipWriter).(io.Writer)
 
 	// Register TCP and UDP handlers to handle accepted connections.
 	proxyAddr, err := net.ResolveTCPAddr("tcp", *args.ProxyServer)
@@ -135,8 +133,8 @@ func main() {
 	}
 	proxyHost := proxyAddr.IP.String()
 	proxyPort := proxyAddr.Port
-	core.RegisterTCPConnHandler(proxy.NewTCPHandler(proxyHost, proxyPort, fakeDns, sessionStater))
-	core.RegisterUDPConnHandler(proxy.NewUDPHandler(proxyHost, proxyPort, *args.UdpTimeout, fakeDns, sessionStater))
+	core.RegisterTCPConnHandler(proxy.NewTCPHandler(proxyHost, proxyPort, fakeDNS, sessionStater))
+	core.RegisterUDPConnHandler(proxy.NewUDPHandler(proxyHost, proxyPort, *args.UdpTimeout, fakeDNS, sessionStater))
 
 	// Register an output callback to write packets output from lwip stack to tun
 	// device, output function should be set before input any packets.
@@ -157,6 +155,12 @@ func main() {
 	signal.Notify(osSignals, os.Interrupt, os.Kill, syscall.SIGTERM, syscall.SIGHUP)
 	<-osSignals
 
+	// Stop fakeDNS server
+	if fakeDNS != nil {
+		fakeDNS.Stop()
+	}
+
+	// Stop session stater
 	if sessionStater != nil {
 		sessionStater.Stop()
 	}

cmd/main_stats.go

@@ -17,7 +17,7 @@ func init() {
 			sessionStater = session.NewSimpleSessionStater()
 
 			// stats variables
-			session.StatsAddr = *args.StatsAddr
+			session.ServeAddr = *args.StatsAddr
 			session.StatsVersion = version
 
 			// start session stater

common/dns/fakedns/server.go

@@ -6,6 +6,7 @@ import (
 
 	D "github.com/miekg/dns"
 	"github.com/xjasonlyu/tun2socks/common/fakeip"
+	"github.com/xjasonlyu/tun2socks/common/log"
 )
 
 const (
@@ -13,6 +14,10 @@ const (
 	dnsDefaultTTL uint32 = 600
 )
 
+var (
+	ServeAddr = "127.0.0.1:5353"
+)
+
 type Server struct {
 	*D.Server
 	p *fakeip.Pool
@@ -28,13 +33,14 @@ func (s *Server) ServeDNS(w D.ResponseWriter, r *D.Msg) {
 	s.h(w, r)
 }
 
-func (s *Server) StartServer(addr string) error {
-	_, port, err := net.SplitHostPort(addr)
+func (s *Server) Start() error {
+	log.Debugf("Start fake DNS server")
+	_, port, err := net.SplitHostPort(ServeAddr)
 	if port == "0" || port == "" || err != nil {
 		return errors.New("address format error")
 	}
 
-	udpAddr, err := net.ResolveUDPAddr("udp", addr)
+	udpAddr, err := net.ResolveUDPAddr("udp", ServeAddr)
 	if err != nil {
 		return err
 	}
@@ -44,12 +50,16 @@ func (s *Server) StartServer(addr string) error {
 		return err
 	}
 
-	s.Server = &D.Server{Addr: addr, PacketConn: p, Handler: s}
+	s.Server = &D.Server{Addr: ServeAddr, PacketConn: p, Handler: s}
 	go s.ActivateAndServe()
-
 	return nil
 }
 
+func (s *Server) Stop() error {
+	log.Debugf("Stop fake DNS server")
+	return s.Shutdown()
+}
+
 func (s *Server) IPToHost(ip net.IP) (string, bool) {
 	return s.p.LookBack(ip)
 }

common/stats/session/session.go

@@ -17,8 +17,8 @@ import (
 const maxCompletedSessions = 100
 
 var (
-	StatsAddr = "localhost:6001"
-	StatsPath = "/stats/session/plain"
+	ServeAddr = "localhost:6001"
+	ServePath = "/stats/session/plain"
 
 	StatsVersion = ""
 )
@@ -98,20 +98,21 @@ table, th, td {
 	_ = w.Flush()
 }
 
-func (s *simpleSessionStater) Start() {
+func (s *simpleSessionStater) Start() error {
 	log.Debugf("Start session stater")
 	mux := http.NewServeMux()
 	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
-		http.Redirect(w, r, StatsPath, 301)
+		http.Redirect(w, r, ServePath, 301)
 	})
-	mux.HandleFunc(StatsPath, s.sessionStatsHandler)
-	s.server = &http.Server{Addr: StatsAddr, Handler: mux}
+	mux.HandleFunc(ServePath, s.sessionStatsHandler)
+	s.server = &http.Server{Addr: ServeAddr, Handler: mux}
 	go s.server.ListenAndServe()
+	return nil
 }
 
-func (s *simpleSessionStater) Stop() {
+func (s *simpleSessionStater) Stop() error {
 	log.Debugf("Stop session stater")
-	s.server.Close()
+	return s.server.Close()
 }
 
 func (s *simpleSessionStater) AddSession(key interface{}, session *stats.Session) {

common/stats/stats.go

@@ -8,8 +8,8 @@ import (
 )
 
 type SessionStater interface {
-	Start()
-	Stop()
+	Start() error
+	Stop() error
 	AddSession(key interface{}, session *Session)
 	GetSession(key interface{}) *Session
 	RemoveSession(key interface{})