vk-turn-proxy/client/warp/internal/dns.go

package internal

import (
	"context"
	"fmt"
	"net"
	"net/netip"
	"time"

	"golang.zx2c4.com/wireguard/tun/netstack"
)

// TunnelDNSResolver implements a DNS resolver that uses the provided DNS servers
// either inside a MASQUE tunnel (if TunNet is set) or over the system network (if TunNet is nil).
type TunnelDNSResolver struct {
	// TunNet is the network stack for the tunnel you want to use for DNS resolution.
	// If nil, DNS queries are sent over the system network.
	TunNet *netstack.Net

	// DNSAddrs is the list of DNS servers to use for resolution.
	DNSAddrs []netip.Addr

	// Timeout is the timeout for DNS queries on a specific server before trying the next one.
	Timeout time.Duration
}

// Resolve performs a DNS lookup using the provided DNS resolvers.
// It tries each resolver in order until one succeeds, sending queries either through the tunnel
// or over the system network depending on TunNet.
//
// Parameters:
//   - ctx: context.Context - The context for the DNS lookup.
//   - name: string - The domain name to resolve.
//
// Returns:
//   - context.Context: The original context for the DNS lookup.
//   - net.IP: The resolved IP address.
//   - error: An error if the lookup fails.
func (r TunnelDNSResolver) Resolve(ctx context.Context, name string) (context.Context, net.IP, error) {
	if len(r.DNSAddrs) == 0 {
		return ctx, nil, fmt.Errorf("no DNS servers configured")
	}

	var queryCtx context.Context = ctx
	var cancel context.CancelFunc
	if r.Timeout > 0 {
		queryCtx, cancel = context.WithTimeout(ctx, r.Timeout)
		defer cancel()
	}

	type result struct {
		ip  net.IP
		err error
	}
	results := make(chan result, len(r.DNSAddrs))

	for _, dnsAddr := range r.DNSAddrs {
		dnsHost := net.JoinHostPort(dnsAddr.String(), "53")

		go func(dnsHost string) {
			var dialFunc func(context.Context, string, string) (net.Conn, error)
			if r.TunNet != nil {
				dialFunc = func(ctx context.Context, network, address string) (net.Conn, error) {
					return r.TunNet.DialContext(ctx, "udp", dnsHost)
				}
			} else {
				dialFunc = func(ctx context.Context, network, address string) (net.Conn, error) {
					return net.Dial("udp", dnsHost)
				}
			}

			resolver := &net.Resolver{
				PreferGo: true,
				Dial:     dialFunc,
			}
			ips, err := resolver.LookupIP(queryCtx, "ip", name)
			if err == nil && len(ips) > 0 {
				results <- result{ip: ips[0], err: nil}
			} else {
				results <- result{ip: nil, err: err}
			}
		}(dnsHost)
	}

	var lastErr error
	for i := 0; i < len(r.DNSAddrs); i++ {
		res := <-results
		if res.err == nil && res.ip != nil {
			if cancel != nil {
				cancel()
			}
			return ctx, res.ip, nil
		}
		lastErr = res.err
	}

	return ctx, nil, fmt.Errorf("all DNS servers failed: %v", lastErr)
}

// NewNetstackResolver returns a *net.Resolver that uses the tunnel network stack
// and provided DNS servers for DNS queries.
//
// Parameters:
//   - tunNet: *netstack.Net - The tunnel network stack.
//   - dnsAddrs: []netip.Addr - DNS server addresses.
//
// Returns:
//   - *net.Resolver - A resolver that routes queries through the tunnel.
func NewNetstackResolver(tunNet *netstack.Net, dnsAddrs []netip.Addr) *net.Resolver {
	return &net.Resolver{
		PreferGo: true,
		Dial: func(ctx context.Context, network, address string) (net.Conn, error) {
			if len(dnsAddrs) == 0 {
				return nil, fmt.Errorf("no DNS servers configured")
			}
			if len(dnsAddrs) > 1 {
				return raceDial(ctx, tunNet, dnsAddrs)
			}
			dnsHost := net.JoinHostPort(dnsAddrs[0].String(), "53")
			return tunNet.DialContext(ctx, "udp", dnsHost)
		},
	}
}

func raceDial(ctx context.Context, tunNet *netstack.Net, addrs []netip.Addr) (net.Conn, error) {
	type result struct {
		conn net.Conn
		err  error
	}
	resChan := make(chan result, len(addrs))
	childCtx, cancel := context.WithCancel(ctx)
	defer cancel()

	for _, addr := range addrs {
		go func(a netip.Addr) {
			dnsHost := net.JoinHostPort(a.String(), "53")
			conn, err := tunNet.DialContext(childCtx, "udp", dnsHost)
			if err == nil {
				select {
				case resChan <- result{conn: conn}:
				case <-ctx.Done():
					conn.Close()
				}
			} else {
				select {
				case resChan <- result{err: err}:
				case <-ctx.Done():
				}
			}
		}(addr)
	}

	var lastErr error
	for i := 0; i < len(addrs); i++ {
		res := <-resChan
		if res.err == nil {
			return res.conn, nil
		}
		lastErr = res.err
	}
	return nil, fmt.Errorf("all DNS race dials failed: %w", lastErr)
}

// NewStaticResolver returns a *net.Resolver that uses the provided DNS servers
// for lookups over the system network.
func NewStaticResolver(dnsAddrs []netip.Addr) *net.Resolver {
	return &net.Resolver{
		PreferGo: true,
		Dial: func(ctx context.Context, network, address string) (net.Conn, error) {
			if len(dnsAddrs) == 0 {
				return nil, fmt.Errorf("no DNS servers configured")
			}
			dnsHost := net.JoinHostPort(dnsAddrs[0].String(), "53")
			return net.Dial("udp", dnsHost)
		},
	}
}

// GetProxyResolver returns the appropriate *net.Resolver for proxy use
// based on the localDNS flag.
//
// Parameters:
//   - localDNS: bool - Whether to use the system network for DNS.
//   - tunNet: *netstack.Net - The tunnel network stack (if localDNS is false).
//   - dnsAddrs: []netip.Addr - DNS server addresses.
//   - timeout: time.Duration - Timeout for DNS queries.
//
// Returns:
//   - *net.Resolver - A resolver suitable for use with proxy connections.
func GetProxyResolver(localDNS bool, tunNet *netstack.Net, dnsAddrs []netip.Addr, timeout time.Duration) *net.Resolver {
	if localDNS {
		return NewStaticResolver(dnsAddrs)
	}
	return NewNetstackResolver(tunNet, dnsAddrs)
}