cleanup dns & stats codes

cmd/main_fakedns.go

@@ -4,9 +4,9 @@ package main
 
 import (
 	"flag"
-	"fmt"
 
 	"github.com/xjasonlyu/tun2socks/component/dns/fakedns"
+	"github.com/xjasonlyu/tun2socks/log"
 )
 
 func init() {
@@ -21,14 +21,17 @@ func init() {
 			var err error
 			fakeDNS, err = fakedns.NewServer(*args.FakeIPRange, *args.FakeDNSHosts, *args.DNSCacheSize)
 			if err != nil {
-				panic(fmt.Errorf("create fake dns server failed: %v", err))
+				log.Fatalf("Create fake DNS server failed: %v", err)
 			}
 
 			// Set fakeDNS variables
 			fakedns.ServeAddr = *args.FakeDNSAddr
 
 			// Start fakeDNS server
-			fakeDNS.Start()
+			if err := fakeDNS.Start(); err != nil {
+				log.Fatalf("Start fake DNS server failed: %v", err)
+			}
+			log.Infof("Fake DNS serving at %v", fakedns.ServeAddr)
 		} else {
 			fakeDNS = nil
 		}

cmd/main_stats.go

@@ -6,6 +6,7 @@ import (
 	"flag"
 
 	"github.com/xjasonlyu/tun2socks/component/stats/session"
+	"github.com/xjasonlyu/tun2socks/log"
 )
 
 func init() {
@@ -20,7 +21,10 @@ func init() {
 			session.ServeAddr = *args.StatsAddr
 
 			// Start session stater
-			sessionStater.Start()
+			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/dns/fakedns/middleware.go

@@ -1,7 +1,6 @@
 package fakedns
 
 import (
-	"fmt"
 	"net"
 	"strings"
 
@@ -80,28 +79,6 @@ func withHost(hosts *trie.Trie, next handler) handler {
 	}
 }
 
-func lineToHosts(str string) *trie.Trie {
-	// trim `'` `"` ` ` char
-	str = strings.Trim(str, "' \"")
-	if str == "" {
-		return nil
-	}
-	tree := trie.New()
-	s := strings.Split(str, ",")
-	for _, host := range s {
-		m := strings.Split(host, "=")
-		if len(m) != 2 {
-			continue
-		}
-		domain := strings.TrimSpace(m[0])
-		target := strings.TrimSpace(m[1])
-		if err := tree.Insert(domain, net.ParseIP(target)); err != nil {
-			panic(fmt.Sprintf("add hosts error: %v", err))
-		}
-	}
-	return tree
-}
-
 func newHandler(hosts *trie.Trie, pool *fakeip.Pool) handler {
 	if hosts != nil {
 		return withHost(hosts, withFakeIP(pool))

component/dns/fakedns/server.go

@@ -2,11 +2,14 @@ package fakedns
 
 import (
 	"errors"
+	"fmt"
 	"net"
+	"strings"
 
-	D "github.com/miekg/dns"
+	trie "github.com/xjasonlyu/tun2socks/common/domain-trie"
 	"github.com/xjasonlyu/tun2socks/common/fakeip"
-	"github.com/xjasonlyu/tun2socks/log"
+
+	D "github.com/miekg/dns"
 )
 
 const (
@@ -33,7 +36,6 @@ func (s *Server) ServeDNS(w D.ResponseWriter, r *D.Msg) {
 }
 
 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")
@@ -44,18 +46,20 @@ func (s *Server) Start() error {
 		return err
 	}
 
-	p, err := net.ListenUDP("udp", udpAddr)
+	pc, err := net.ListenUDP("udp", udpAddr)
 	if err != nil {
 		return err
 	}
 
-	s.Server = &D.Server{Addr: ServeAddr, PacketConn: p, Handler: s}
-	go s.ActivateAndServe()
+	s.Server = &D.Server{Addr: ServeAddr, PacketConn: pc, Handler: s}
+	go func() {
+		s.ActivateAndServe()
+	}()
+
 	return nil
 }
 
 func (s *Server) Stop() error {
-	log.Debugf("Stop fake DNS server")
 	return s.Shutdown()
 }
 
@@ -63,7 +67,7 @@ func (s *Server) IPToHost(ip net.IP) (string, bool) {
 	return s.p.LookBack(ip)
 }
 
-func NewServer(fakeIPRange, hostsLine string, size int) (*Server, error) {
+func NewServer(fakeIPRange, hosts string, size int) (*Server, error) {
 	_, ipnet, err := net.ParseCIDR(fakeIPRange)
 	if err != nil {
 		return nil, err
@@ -73,8 +77,29 @@ func NewServer(fakeIPRange, hostsLine string, size int) (*Server, error) {
 		return nil, err
 	}
 
-	hosts := lineToHosts(hostsLine)
-	handler := newHandler(hosts, pool)
+	hostsTree := func(str string) *trie.Trie {
+		// trim `'` `"` ` ` char
+		str = strings.Trim(str, "' \"")
+		if str == "" {
+			return nil
+		}
+		tree := trie.New()
+		s := strings.Split(str, ",")
+		for _, host := range s {
+			m := strings.Split(host, "=")
+			if len(m) != 2 {
+				continue
+			}
+			domain := strings.TrimSpace(m[0])
+			target := strings.TrimSpace(m[1])
+			if err := tree.Insert(domain, net.ParseIP(target)); err != nil {
+				panic(fmt.Sprintf("add hosts error: %v", err))
+			}
+		}
+		return tree
+	}(hosts)
+
+	handler := newHandler(hostsTree, pool)
 
 	return &Server{
 		p: pool,

component/stats/session/session.go

@@ -2,8 +2,10 @@ package session
 
 import (
 	"bufio"
+	"errors"
 	"fmt"
 	"io"
+	"net"
 	"net/http"
 	"sort"
 	"sync"
@@ -11,7 +13,6 @@ import (
 
 	"github.com/xjasonlyu/tun2socks/common/queue"
 	C "github.com/xjasonlyu/tun2socks/constant"
-	"github.com/xjasonlyu/tun2socks/log"
 )
 
 const maxCompletedSessions = 100
@@ -97,19 +98,35 @@ table, th, td {
 }
 
 func (s *simpleSessionStater) Start() error {
-	log.Debugf("Start session stater")
+	_, port, err := net.SplitHostPort(ServeAddr)
+	if port == "0" || port == "" || err != nil {
+		return errors.New("address format error")
+	}
+
+	tcpAddr, err := net.ResolveTCPAddr("tcp", ServeAddr)
+	if err != nil {
+		return err
+	}
+
+	c, err := net.ListenTCP("tcp", tcpAddr)
+	if err != nil {
+		return err
+	}
+
 	mux := http.NewServeMux()
 	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
 		http.Redirect(w, r, ServePath, 301)
 	})
 	mux.HandleFunc(ServePath, s.sessionStatsHandler)
 	s.server = &http.Server{Addr: ServeAddr, Handler: mux}
-	go s.server.ListenAndServe()
+	go func() {
+		s.server.Serve(c)
+	}()
+
 	return nil
 }
 
 func (s *simpleSessionStater) Stop() error {
-	log.Debugf("Stop session stater")
 	return s.server.Close()
 }