resolver: add DNS address preference wrapper

Author: Snawoot

main.go

@@ -197,6 +197,25 @@ type autocertCache struct {
 
 const envCacheEncKey = "DUMBPROXY_CACHE_ENC_KEY"
 
+type dnsPreferenceArg resolver.Preference
+
+func (a *dnsPreferenceArg) String() string {
+	return resolver.Preference(*a).String()
+}
+
+func (a *dnsPreferenceArg) Set(s string) error {
+	p, err := resolver.ParsePreference(s)
+	if err != nil {
+		return nil
+	}
+	*a = dnsPreferenceArg(p)
+	return nil
+}
+
+func (a *dnsPreferenceArg) Value() resolver.Preference {
+	return resolver.Preference(*a)
+}
+
 type bindSpec struct {
 	af      string
 	address string
@@ -241,6 +260,8 @@ type CLIArgs struct {
 	bwBurst                   int64
 	bwBuckets                 uint
 	bwSeparate                bool
+	dnsServers                []string
+	dnsPreferAddress          dnsPreferenceArg
 	dnsCacheTTL               time.Duration
 	dnsCacheNegTTL            time.Duration
 	dnsCacheTimeout           time.Duration
@@ -251,7 +272,6 @@ type CLIArgs struct {
 	jsProxyRouterInstances    int
 	proxyproto                bool
 	shutdownTimeout           time.Duration
-	nameservers               []string
 }
 
 func parse_args() CLIArgs {
@@ -277,6 +297,7 @@ func parse_args() CLIArgs {
 			address: ":8080",
 			af:      "tcp",
 		},
+		dnsPreferAddress: dnsPreferenceArg(resolver.PreferenceIPv4),
 	}
 	args.autocertCacheEncKey.Set(os.Getenv(envCacheEncKey))
 	flag.Func("bind-address", "HTTP proxy listen address. Set empty value to use systemd socket activation. (default \":8080\")", func(p string) error {
@@ -364,12 +385,13 @@ func parse_args() CLIArgs {
 	flag.BoolVar(&args.bwSeparate, "bw-limit-separate", false, "separate upload and download bandwidth limits")
 	flag.Func("dns-server", "nameserver specification (udp://..., tcp://..., https://..., tls://..., doh://..., dot://..., default://). Option can be used multiple times for parallel use of multiple nameservers. Empty string resets list", func(p string) error {
 		if p == "" {
-			args.nameservers = nil
+			args.dnsServers = nil
 		} else {
-			args.nameservers = append(args.nameservers, p)
+			args.dnsServers = append(args.dnsServers, p)
 		}
 		return nil
 	})
+	flag.Var(&args.dnsPreferAddress, "dns-prefer-address", "address resolution preference (none/ipv4/ipv6)")
 	flag.DurationVar(&args.dnsCacheTTL, "dns-cache-ttl", 0, "enable DNS cache with specified fixed TTL")
 	flag.DurationVar(&args.dnsCacheNegTTL, "dns-cache-neg-ttl", time.Second, "TTL for negative responses of DNS cache")
 	flag.DurationVar(&args.dnsCacheTimeout, "dns-cache-timeout", 5*time.Second, "timeout for shared resolves of DNS cache")
@@ -521,13 +543,14 @@ func run() int {
 	dialerRoot = dialer.NewFilterDialer(filterRoot.Access, dialerRoot) // must follow after resolving in chain
 
 	var nameResolver dialer.Resolver = net.DefaultResolver
-	if len(args.nameservers) > 0 {
-		nameResolver, err = resolver.FastFromURLs(args.nameservers...)
+	if len(args.dnsServers) > 0 {
+		nameResolver, err = resolver.FastFromURLs(args.dnsServers...)
 		if err != nil {
 			mainLogger.Critical("Failed to create name resolver: %v", err)
 			return 3
 		}
 	}
+	nameResolver = resolver.Prefer(nameResolver, args.dnsPreferAddress.Value())
 	if args.dnsCacheTTL > 0 {
 		cd := dialer.NewNameResolveCachingDialer(
 			dialerRoot,

resolver/prefer.go

@@ -0,0 +1,101 @@
+package resolver
+
+import (
+	"cmp"
+	"context"
+	"fmt"
+	"net/netip"
+	"slices"
+	"strings"
+)
+
+type Preference int
+
+const (
+	PreferenceNothing Preference = iota
+	PreferenceIPv4
+	PreferenceIPv6
+)
+
+func (p Preference) String() string {
+	switch p {
+	case PreferenceNothing:
+		return "none"
+	case PreferenceIPv4:
+		return "ipv4"
+	case PreferenceIPv6:
+		return "ipv6"
+	default:
+		return fmt.Sprintf("Preference(%d)", int(p))
+	}
+}
+
+func ParsePreference(p string) (Preference, error) {
+	var res Preference
+	switch lp := strings.ToLower(p); lp {
+	case "none", "nothing", "any", "anything":
+		res = PreferenceNothing
+	case "ipv4", "ip4", "v4", "4":
+		res = PreferenceIPv4
+	case "ipv6", "ip6", "v6", "6":
+		res = PreferenceIPv6
+	default:
+		return 0, fmt.Errorf("unknown preference specification %q", p)
+	}
+	return res, nil
+}
+
+func boolToInt(x bool) int {
+	if x {
+		return 0
+	}
+	return 1
+}
+
+type PreferIPv4 struct {
+	LookupNetIPer
+}
+
+func (p PreferIPv4) LookupNetIP(ctx context.Context, network, host string) ([]netip.Addr, error) {
+	addrs, err := p.LookupNetIPer.LookupNetIP(ctx, network, host)
+	if err != nil {
+		return nil, err
+	}
+	slices.SortStableFunc(addrs, func(a, b netip.Addr) int {
+		return cmp.Compare(
+			boolToInt(a.Unmap().Is4()),
+			boolToInt(b.Unmap().Is4()),
+		)
+	})
+	return addrs, nil
+}
+
+type PreferIPv6 struct {
+	LookupNetIPer
+}
+
+func (p PreferIPv6) LookupNetIP(ctx context.Context, network, host string) ([]netip.Addr, error) {
+	addrs, err := p.LookupNetIPer.LookupNetIP(ctx, network, host)
+	if err != nil {
+		return nil, err
+	}
+	slices.SortStableFunc(addrs, func(a, b netip.Addr) int {
+		return cmp.Compare(
+			boolToInt(a.Unmap().Is6()),
+			boolToInt(b.Unmap().Is6()),
+		)
+	})
+	return addrs, nil
+}
+
+func Prefer(resolver LookupNetIPer, p Preference) LookupNetIPer {
+	switch p {
+	case PreferenceNothing:
+		return resolver
+	case PreferenceIPv4:
+		return PreferIPv4{resolver}
+	case PreferenceIPv6:
+		return PreferIPv6{resolver}
+	}
+	panic("unknown address family preference")
+}