use buildin dns to connect websocket

README.md

@@ -7,8 +7,8 @@
 ```
 # git clone https://github.com/rkonfj/toh.git
 # go build -ldflags "-s -w"
-
 ```
+
 - Run
 ```
 # ./toh server --help

client/client.go

@@ -2,6 +2,7 @@ package client
 
 import (
 	"context"
+	"fmt"
 	"math/rand"
 	"net"
 	"net/http"
@@ -14,7 +15,8 @@ import (
 )
 
 type TohClient struct {
-	options Options
+	options    Options
+	httpClient *http.Client
 }
 
 type Options struct {
@@ -23,8 +25,32 @@ type Options struct {
 }
 
 func NewTohClient(options Options) (*TohClient, error) {
+	httpClient := &http.Client{
+		Transport: &http.Transport{
+			DialContext: func(ctx context.Context, network, addr string) (conn net.Conn, err error) {
+				dialer := &net.Dialer{}
+				host, port, err := net.SplitHostPort(addr)
+				if err != nil {
+					return
+				}
+				dnsLookupCtx, cancel := context.WithTimeout(ctx, 3*time.Second)
+				defer cancel()
+				ips, err := (&net.Resolver{
+					Dial: func(ctx context.Context, network, address string) (net.Conn, error) {
+						return dialer.DialContext(ctx, "tcp", "8.8.8.8:53")
+					},
+				}).LookupIP(dnsLookupCtx, "ip", host)
+				if err != nil {
+					return
+				}
+				conn, err = dialer.DialContext(ctx, network, fmt.Sprintf("%s:%s", ips[rand.Intn(len(ips))], port))
+				return
+			},
+		},
+	}
 	return &TohClient{
-		options: options,
+		options:    options,
+		httpClient: httpClient,
 	}, nil
 }
 
@@ -72,7 +98,9 @@ func (c *TohClient) dial(ctx context.Context, network, addr string) (conn *webso
 	handshake.Add("x-toh-addr", addr)
 
 	t1 := time.Now()
-	conn, _, err = websocket.Dial(ctx, c.options.ServerAddr, &websocket.DialOptions{HTTPHeader: handshake})
+	conn, _, err = websocket.Dial(ctx, c.options.ServerAddr, &websocket.DialOptions{
+		HTTPHeader: handshake, HTTPClient: c.httpClient,
+	})
 	if err != nil {
 		return
 	}

cmd/pf/tunnel.go

@@ -84,7 +84,7 @@ func (t *TunnelManager) Run() {
 	}
 
 	for _, f := range t.forwards {
-		logrus.Infof("listen %s://%s for %s now", f.network, f.local, f.remote)
+		logrus.Infof("listen on %s for %s://%s now", f.local, f.network, f.remote)
 		go t.forward(f)
 	}
 	t.wg.Wait()

server/server.go

@@ -67,6 +67,9 @@ func (s *TohServer) Run() {
 }
 
 func (s *TohServer) pipe(wsConn *websocket.Conn, netConn net.Conn) {
+	if wsConn == nil || netConn == nil {
+		return
+	}
 	go func() {
 		io.Copy(netConn, RWWS(wsConn))
 		logrus.Debugf("ws conn closed, close remote conn(%s) now", netConn.RemoteAddr().String())

socks5/socks5.go

@@ -32,7 +32,7 @@ func (s *Socks5Server) Run() error {
 	if err != nil {
 		return err
 	}
-	logrus.Infof("socks5 listen on %s now", s.opts.Listen)
+	logrus.Infof("listen on %s for socks5 now", s.opts.Listen)
 	defer l.Close()
 	for {
 		conn, err := l.Accept()
@@ -167,7 +167,9 @@ func (s *Socks5Server) handshake(conn net.Conn, buf []byte) (tcpConn, udpConn ne
 }
 
 func (s *Socks5Server) pipe(conn, rConn net.Conn) {
-	logrus.Debugf("strat pipeline %s<->%s", conn.RemoteAddr().String(), rConn.RemoteAddr().String())
+	if conn == nil || rConn == nil {
+		return
+	}
 	go func() {
 		io.Copy(conn, rConn)
 		conn.Close()