feat: use internal resolver for wg network

internal/proxy/proxy.go

@@ -0,0 +1,94 @@
+package proxy
+
+import (
+	"context"
+	"encoding/json"
+	"net"
+	"net/http"
+
+	"github.com/zhsj/wghttp/internal/resolver"
+	"github.com/zhsj/wghttp/internal/third_party/tailscale/httpproxy"
+	"github.com/zhsj/wghttp/internal/third_party/tailscale/proxymux"
+	"github.com/zhsj/wghttp/internal/third_party/tailscale/socks5"
+)
+
+type dialer func(ctx context.Context, network, address string) (net.Conn, error)
+
+type Proxy struct {
+	Dial  dialer
+	DNS   string
+	Stats func() (any, error)
+}
+
+func statsHandler(next http.Handler, stats func() (any, error)) http.Handler {
+	return http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
+		if r.URL.Host != "" || r.URL.Path != "/stats" {
+			next.ServeHTTP(rw, r)
+			return
+		}
+		s, err := stats()
+		if err != nil {
+			rw.WriteHeader(http.StatusInternalServerError)
+		} else {
+			resp, _ := json.MarshalIndent(s, "", "  ")
+			rw.Header().Set("Content-Type", "application/json")
+			_, _ = rw.Write(append(resp, '\n'))
+		}
+	})
+}
+
+func dialWithDNS(dial dialer, dns string) dialer {
+	resolv := resolver.New(dns, dial)
+
+	return func(ctx context.Context, network, address string) (net.Conn, error) {
+		host, port, err := net.SplitHostPort(address)
+		if err != nil {
+			return nil, err
+		}
+		if err == nil {
+			if ip := net.ParseIP(host); ip != nil {
+				return dial(ctx, network, address)
+			}
+		}
+
+		ips, err := resolv.LookupHost(ctx, host)
+		if err != nil {
+			return nil, err
+		}
+
+		var (
+			lastErr error
+			conn    net.Conn
+		)
+		for _, ip := range ips {
+			addr := net.JoinHostPort(ip, port)
+			conn, lastErr = dial(ctx, network, addr)
+			if lastErr == nil {
+				return conn, nil
+			}
+		}
+		return nil, lastErr
+	}
+}
+
+func (p Proxy) Serve(ln net.Listener) {
+	d := dialWithDNS(p.Dial, p.DNS)
+
+	socksListener, httpListener := proxymux.SplitSOCKSAndHTTP(ln)
+
+	httpProxy := &http.Server{Handler: statsHandler(httpproxy.Handler(d), p.Stats)}
+	socksProxy := &socks5.Server{Dialer: d}
+
+	errc := make(chan error, 2)
+	go func() {
+		if err := httpProxy.Serve(httpListener); err != nil {
+			errc <- err
+		}
+	}()
+	go func() {
+		if err := socksProxy.Serve(socksListener); err != nil {
+			errc <- err
+		}
+	}()
+	<-errc
+}

internal/proxy/proxy_test.go

@@ -0,0 +1,45 @@
+package proxy
+
+import (
+	"context"
+	"fmt"
+	"net"
+	"testing"
+)
+
+func TestDialWithDNS(t *testing.T) {
+	if testing.Short() {
+		t.Skip()
+	}
+
+	stdDiar := net.Dialer{
+		Resolver: &net.Resolver{
+			Dial: func(ctx context.Context, network, address string) (net.Conn, error) {
+				return nil, fmt.Errorf("dial to %s", address)
+			},
+		},
+	}
+
+	// d := dialWithDNS(stdDiar.DialContext, "https://223.5.5.5")
+	d := dialWithDNS(func(ctx context.Context, network, address string) (net.Conn, error) {
+		t.Logf("dial to %s:%s", network, address)
+		return stdDiar.DialContext(ctx, network, address)
+	}, "tls://223.5.5.5")
+
+	for _, addr := range []string{
+		"example.com:80",
+		"223.6.6.6:80",
+		"localhost:22",
+		"127.0.0.1:22",
+	} {
+		t.Run(addr, func(t *testing.T) {
+			conn, err := d(context.Background(), "tcp", addr)
+			if err != nil {
+				t.Error(err)
+			} else {
+				t.Logf("remote %s", conn.RemoteAddr())
+				conn.Close()
+			}
+		})
+	}
+}

main.go

@@ -6,7 +6,6 @@ import (
 	_ "embed"
 	"fmt"
 	"net"
-	"net/http"
 	"net/netip"
 	"os"
 	"strings"
@@ -16,9 +15,7 @@ import (
 	"golang.zx2c4.com/wireguard/device"
 	"golang.zx2c4.com/wireguard/tun/netstack"
 
-	"github.com/zhsj/wghttp/internal/third_party/tailscale/httpproxy"
+	"github.com/zhsj/wghttp/internal/proxy"
-	"github.com/zhsj/wghttp/internal/third_party/tailscale/proxymux"
-	"github.com/zhsj/wghttp/internal/third_party/tailscale/socks5"
 )
 
 //go:embed README.md
@@ -33,7 +30,7 @@ type options struct {
 	ClientIPs  []string `long:"client-ip" env:"CLIENT_IP" env-delim:"," description:"[Interface].Address\tfor WireGuard client (can be set multiple times)"`
 	ClientPort int      `long:"client-port" env:"CLIENT_PORT" description:"[Interface].ListenPort\tfor WireGuard client (optional)"`
 	PrivateKey string   `long:"private-key" env:"PRIVATE_KEY" description:"[Interface].PrivateKey\tfor WireGuard client (format: base64)"`
-	DNS        string   `long:"dns" env:"DNS" description:"[Interface].DNS\tfor WireGuard network (format: IP)"`
+	DNS        string   `long:"dns" env:"DNS" description:"[Interface].DNS\tfor WireGuard network (format: protocol://ip:port)\nProtocol includes udp(default), tcp, tls(DNS over TLS) and https(DNS over HTTPS)"`
 	MTU        int      `long:"mtu" env:"MTU" default:"1280" description:"[Interface].MTU\tfor WireGuard network"`
 
 	PeerEndpoint      string        `long:"peer-endpoint" env:"PEER_ENDPOINT" description:"[Peer].Endpoint\tfor WireGuard server (format: host:port)"`
@@ -41,7 +38,7 @@ type options struct {
 	PresharedKey      string        `long:"preshared-key" env:"PRESHARED_KEY" description:"[Peer].PresharedKey\tfor WireGuard network (optional, format: base64)"`
 	KeepaliveInterval time.Duration `long:"keepalive-interval" env:"KEEPALIVE_INTERVAL" description:"[Peer].PersistentKeepalive\tfor WireGuard network (optional)"`
 
-	ResolveDNS      string        `long:"resolve-dns" env:"RESOLVE_DNS" description:"DNS for resolving WireGuard server address (optional, format: protocol://ip:port)\nProtocol includes tcp, udp, tls(DNS over TLS) and https(DNS over HTTPS)"`
+	ResolveDNS      string        `long:"resolve-dns" env:"RESOLVE_DNS" description:"DNS for resolving WireGuard server address (optional, format: protocol://ip:port)\nProtocol includes udp(default), tcp, tls(DNS over TLS) and https(DNS over HTTPS)"`
 	ResolveInterval time.Duration `long:"resolve-interval" env:"RESOLVE_INTERVAL" default:"1m" description:"Interval for resolving WireGuard server address (set 0 to disable)"`
 
 	Listen   string `long:"listen" env:"LISTEN" default:"localhost:8080" description:"HTTP & SOCKS5 server address"`
@@ -89,26 +86,11 @@ Description:`
 		os.Exit(1)
 	}
 
-	socksListener, httpListener := proxymux.SplitSOCKSAndHTTP(listener)
+	proxier := proxy.Proxy{
-	dialer := proxyDialer(tnet)
+		Dial: proxyDialer(tnet), DNS: opts.DNS, Stats: stats(dev),
-
-	httpProxy := &http.Server{Handler: statsHandler(httpproxy.Handler(dialer), dev)}
-	socksProxy := &socks5.Server{Dialer: dialer}
-
-	errc := make(chan error, 2)
-	go func() {
-		if err := httpProxy.Serve(httpListener); err != nil {
-			logger.Errorf("Serving http proxy: %v", err)
-			errc <- err
 	}
-	}()
+	proxier.Serve(listener)
-	go func() {
+
-		if err := socksProxy.Serve(socksListener); err != nil {
-			logger.Errorf("Serving socks5 proxy: %v", err)
-			errc <- err
-		}
-	}()
-	<-errc
 	os.Exit(1)
 }
 
@@ -151,7 +133,7 @@ func setupNet() (*device.Device, *netstack.Net, error) {
 	if len(opts.ClientIPs) == 0 {
 		return nil, nil, fmt.Errorf("client IP is required")
 	}
-	var clientIPs, dnsServers []netip.Addr
+	var clientIPs []netip.Addr
 	for _, s := range opts.ClientIPs {
 		ip, err := netip.ParseAddr(s)
 		if err != nil {
@@ -159,14 +141,8 @@ func setupNet() (*device.Device, *netstack.Net, error) {
 		}
 		clientIPs = append(clientIPs, ip)
 	}
-	if opts.DNS != "" {
+
-		ip, err := netip.ParseAddr(opts.DNS)
+	tun, tnet, err := netstack.CreateNetTUN(clientIPs, nil, opts.MTU)
-		if err != nil {
-			return nil, nil, fmt.Errorf("parse DNS IP: %w", err)
-		}
-		dnsServers = append(dnsServers, ip)
-	}
-	tun, tnet, err := netstack.CreateNetTUN(clientIPs, dnsServers, opts.MTU)
 	if err != nil {
 		return nil, nil, fmt.Errorf("create netstack tun: %w", err)
 	}

stats.go

@@ -3,8 +3,6 @@ package main
 import (
 	"bufio"
 	"bytes"
-	"encoding/json"
-	"net/http"
 	"runtime"
 	"strconv"
 	"strings"
@@ -12,11 +10,12 @@ import (
 	"golang.zx2c4.com/wireguard/device"
 )
 
-func statsHandler(next http.Handler, dev *device.Device) http.Handler {
+func stats(dev *device.Device) func() (any, error) {
-	return http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
+	return func() (any, error) {
-		if r.URL.Host != "" || r.URL.Path != "/stats" {
+		var buf bytes.Buffer
-			next.ServeHTTP(rw, r)
+		if err := dev.IpcGetOperation(&buf); err != nil {
-			return
+			logger.Errorf("Get device config: %v", err)
+			return nil, err
 		}
 
 		stats := struct {
@@ -30,11 +29,6 @@ func statsHandler(next http.Handler, dev *device.Device) http.Handler {
 			NumGoroutine: runtime.NumGoroutine(),
 		}
 
-		var buf bytes.Buffer
-		if err := dev.IpcGetOperation(&buf); err != nil {
-			logger.Errorf("Get device config: %v", err)
-			rw.WriteHeader(http.StatusInternalServerError)
-		} else {
 		scanner := bufio.NewScanner(&buf)
 		for scanner.Scan() {
 			line := scanner.Text()
@@ -51,9 +45,6 @@ func statsHandler(next http.Handler, dev *device.Device) http.Handler {
 				stats.SentBytes, _ = strconv.ParseInt(strings.TrimPrefix(line, prefix), 10, 64)
 			}
 		}
-			resp, _ := json.MarshalIndent(stats, "", "  ")
+		return stats, nil
-			rw.Header().Set("Content-Type", "application/json")
-			_, _ = rw.Write(append(resp, '\n'))
 	}
-	})
 }