From 07f3cd2ae5576043064757cf511a5c624e722a2d Mon Sep 17 00:00:00 2001
From: Skyxim <me@skyxim.dev>
Date: Fri, 10 Mar 2023 23:38:16 +0800
Subject: [PATCH] chore: exposure ipv6 wait time

---
 config/config.go         |  4 ++++
 dns/resolver.go          | 37 ++++++++++++++++++++++++-------------
 docs/config.yaml         |  2 +-
 hub/executor/executor.go |  1 +
 4 files changed, 30 insertions(+), 14 deletions(-)

diff --git a/config/config.go b/config/config.go
index 85b1fafa..9316f920 100644
--- a/config/config.go
+++ b/config/config.go
@@ -92,6 +92,7 @@ type DNS struct {
 	Enable                bool             `yaml:"enable"`
 	PreferH3              bool             `yaml:"prefer-h3"`
 	IPv6                  bool             `yaml:"ipv6"`
+	IPv6Timeout           uint             `yaml:"ipv6-timeout"`
 	NameServer            []dns.NameServer `yaml:"nameserver"`
 	Fallback              []dns.NameServer `yaml:"fallback"`
 	FallbackFilter        FallbackFilter   `yaml:"fallback-filter"`
@@ -171,6 +172,7 @@ type RawDNS struct {
 	Enable                bool              `yaml:"enable"`
 	PreferH3              bool              `yaml:"prefer-h3"`
 	IPv6                  bool              `yaml:"ipv6"`
+	IPv6Timeout           uint              `yaml:"ipv6-timeout"`
 	UseHosts              bool              `yaml:"use-hosts"`
 	NameServer            []string          `yaml:"nameserver"`
 	Fallback              []string          `yaml:"fallback"`
@@ -377,6 +379,7 @@ func UnmarshalRawConfig(buf []byte) (*RawConfig, error) {
 			Enable:       false,
 			IPv6:         false,
 			UseHosts:     true,
+			IPv6Timeout:  100,
 			EnhancedMode: C.DNSMapping,
 			FakeIPRange:  "198.18.0.1/16",
 			FallbackFilter: RawFallbackFilter{
@@ -1048,6 +1051,7 @@ func parseDNS(rawCfg *RawConfig, hosts *trie.DomainTrie[netip.Addr], rules []C.R
 		Enable:       cfg.Enable,
 		Listen:       cfg.Listen,
 		PreferH3:     cfg.PreferH3,
+		IPv6Timeout:  cfg.IPv6Timeout,
 		IPv6:         cfg.IPv6,
 		EnhancedMode: cfg.EnhancedMode,
 		FallbackFilter: FallbackFilter{
diff --git a/dns/resolver.go b/dns/resolver.go
index 59b1ee06..69725870 100644
--- a/dns/resolver.go
+++ b/dns/resolver.go
@@ -42,6 +42,7 @@ type geositePolicyRecord struct {
 
 type Resolver struct {
 	ipv6                  bool
+	ipv6Timeout           time.Duration
 	hosts                 *trie.DomainTrie[netip.Addr]
 	main                  []dnsClient
 	fallback              []dnsClient
@@ -91,14 +92,20 @@ func (r *Resolver) LookupIP(ctx context.Context, host string) (ips []netip.Addr,
 	}()
 
 	ips, err = r.lookupIP(ctx, host, D.TypeA)
-
+	var waitIPv6 *time.Timer
+	if r != nil {
+		waitIPv6 = time.NewTimer(r.ipv6Timeout)
+	} else {
+		waitIPv6 = time.NewTimer(100 * time.Millisecond)
+	}
+	defer waitIPv6.Stop()
 	select {
 	case ipv6s, open := <-ch:
 		if !open && err != nil {
 			return nil, resolver.ErrIPNotFound
 		}
 		ips = append(ips, ipv6s...)
-	case <-time.After(30 * time.Millisecond):
+	case <-waitIPv6.C:
 		// wait ipv6 result
 	}
 
@@ -419,6 +426,7 @@ type Config struct {
 	Default        []NameServer
 	ProxyServer    []NameServer
 	IPv6           bool
+	IPv6Timeout    uint
 	EnhancedMode   C.DNSMode
 	FallbackFilter FallbackFilter
 	Pool           *fakeip.Pool
@@ -428,15 +436,17 @@ type Config struct {
 
 func NewResolver(config Config) *Resolver {
 	defaultResolver := &Resolver{
-		main:     transform(config.Default, nil),
-		lruCache: cache.New(cache.WithSize[string, *D.Msg](4096), cache.WithStale[string, *D.Msg](true)),
+		main:        transform(config.Default, nil),
+		lruCache:    cache.New(cache.WithSize[string, *D.Msg](4096), cache.WithStale[string, *D.Msg](true)),
+		ipv6Timeout: time.Duration(config.IPv6Timeout) * time.Millisecond,
 	}
 
 	r := &Resolver{
-		ipv6:     config.IPv6,
-		main:     transform(config.Main, defaultResolver),
-		lruCache: cache.New(cache.WithSize[string, *D.Msg](4096), cache.WithStale[string, *D.Msg](true)),
-		hosts:    config.Hosts,
+		ipv6:        config.IPv6,
+		main:        transform(config.Main, defaultResolver),
+		lruCache:    cache.New(cache.WithSize[string, *D.Msg](4096), cache.WithStale[string, *D.Msg](true)),
+		hosts:       config.Hosts,
+		ipv6Timeout: time.Duration(config.IPv6Timeout) * time.Millisecond,
 	}
 
 	if len(config.Fallback) != 0 {
@@ -502,11 +512,12 @@ func NewResolver(config Config) *Resolver {
 
 func NewProxyServerHostResolver(old *Resolver) *Resolver {
 	r := &Resolver{
-		ipv6:     old.ipv6,
-		main:     old.proxyServer,
-		lruCache: old.lruCache,
-		hosts:    old.hosts,
-		policy:   old.policy,
+		ipv6:        old.ipv6,
+		main:        old.proxyServer,
+		lruCache:    old.lruCache,
+		hosts:       old.hosts,
+		policy:      old.policy,
+		ipv6Timeout: old.ipv6Timeout,
 	}
 	return r
 }
diff --git a/docs/config.yaml b/docs/config.yaml
index b0d4541d..b5af798b 100644
--- a/docs/config.yaml
+++ b/docs/config.yaml
@@ -166,7 +166,7 @@ dns:
   prefer-h3: true # 开启 DoH 支持 HTTP/3，将并发尝试
   listen: 0.0.0.0:53 # 开启 DNS 服务器监听
   # ipv6: false # false 将返回 AAAA 的空结果
-  
+  # ipv6-timeout: 300 # 单位：ms，内部双栈并发时，向上游查询 AAAA 时，等待 AAAA 的时间，默认 100ms
   # 用于解析 nameserver，fallback 以及其他DNS服务器配置的，DNS 服务域名
   # 只能使用纯 IP 地址，可使用加密 DNS
   default-nameserver:
diff --git a/hub/executor/executor.go b/hub/executor/executor.go
index d6ff7851..1bd22385 100644
--- a/hub/executor/executor.go
+++ b/hub/executor/executor.go
@@ -190,6 +190,7 @@ func updateDNS(c *config.DNS, generalIPv6 bool) {
 		Main:         c.NameServer,
 		Fallback:     c.Fallback,
 		IPv6:         c.IPv6 && generalIPv6,
+		IPv6Timeout:  c.IPv6Timeout,
 		EnhancedMode: c.EnhancedMode,
 		Pool:         c.FakeIPRange,
 		Hosts:        c.Hosts,