feat(dns): add lookup host fallback

dns/dns.go

@@ -28,6 +28,12 @@ type dnsstat struct {
 type DNSList struct {
 	sync.RWMutex
 	m map[string][]*dnsstat
+	b map[string][]string
+}
+
+type DNSConfig struct {
+	Servers   map[string][]string // Servers map[dot.com]ip:ports
+	Fallbacks map[string][]string // Fallbacks map[domain]ips
 }
 
 // hasrecord no lock, use under lock
@@ -40,11 +46,21 @@ func hasrecord(lst []*dnsstat, a string) bool {
 	return false
 }
 
-func (ds *DNSList) Add(m map[string][]string) {
+// hasrecord no lock, use under lock
+func hasfallback(lst []string, a string) bool {
+	for _, addr := range lst {
+		if addr == a {
+			return true
+		}
+	}
+	return false
+}
+
+func (ds *DNSList) Add(c *DNSConfig) {
 	ds.Lock()
 	defer ds.Unlock()
 	addList := map[string][]*dnsstat{}
-	for host, addrs := range m {
+	for host, addrs := range c.Servers {
 		for _, addr := range addrs {
 			if !hasrecord(ds.m[host], addr) && !hasrecord(addList[host], addr) {
 				addList[host] = append(addList[host], &dnsstat{addr, true})
@@ -54,6 +70,26 @@ func (ds *DNSList) Add(m map[string][]string) {
 	for host, addrs := range addList {
 		ds.m[host] = append(ds.m[host], addrs...)
 	}
+	addListFallback := map[string][]string{}
+	for host, addrs := range c.Fallbacks {
+		for _, addr := range addrs {
+			if !hasfallback(ds.b[host], addr) && !hasfallback(addListFallback[host], addr) {
+				addListFallback[host] = append(addListFallback[host], addr)
+			}
+		}
+	}
+	for host, addrs := range addListFallback {
+		ds.b[host] = append(ds.b[host], addrs...)
+	}
+}
+
+func (ds *DNSList) LookupHostFallback(ctx context.Context, host string) ([]string, error) {
+	ds.RLock()
+	defer ds.RUnlock()
+	if addrs, ok := ds.b[host]; ok {
+		return addrs, nil
+	}
+	return net.DefaultResolver.LookupHost(ctx, host)
 }
 
 func (ds *DNSList) DialContext(ctx context.Context, dialer *net.Dialer, firstFragmentLen uint8) (tlsConn *tls.Conn, err error) {
@@ -124,6 +160,7 @@ var IPv6Servers = DNSList{
 			{"[2620:fe::fe:10]:853", true},
 		},
 	},
+	b: map[string][]string{},
 }
 
 var IPv4Servers = DNSList{
@@ -149,6 +186,7 @@ var IPv4Servers = DNSList{
 			{"149.112.112.10:853", true},
 		},
 	},
+	b: map[string][]string{},
 }
 
 var DefaultResolver = &net.Resolver{

dns/dns_test.go

@@ -29,7 +29,7 @@ func TestDNS(t *testing.T) {
 		IPv6Servers.test()
 	}
 	IPv4Servers.test()
-	for i := 0; i < 100; i++ {
+	for i := 0; i < 10; i++ {
 		addrs, err := DefaultResolver.LookupHost(context.TODO(), "huggingface.co")
 		if err != nil {
 			t.Fatal(err)
@@ -53,12 +53,16 @@ func TestBadDNS(t *testing.T) {
 		IPv6Servers = DNSList{
 			m: map[string][]*dnsstat{},
 		}
-		IPv6Servers.Add(map[string][]string{"test.bad.host": {"169.254.122.111"}})
+		IPv6Servers.Add(&DNSConfig{
+			Servers: map[string][]string{"test.bad.host": {"169.254.122.111"}},
+		})
 	} else {
 		IPv4Servers = DNSList{
 			m: map[string][]*dnsstat{},
 		}
-		IPv4Servers.Add(map[string][]string{"test.bad.host": {"169.254.122.111:853"}})
+		IPv4Servers.Add(&DNSConfig{
+			Servers: map[string][]string{"test.bad.host": {"169.254.122.111:853"}},
+		})
 	}
 	for i := 0; i < 10; i++ {
 		addrs, err := DefaultResolver.LookupHost(context.TODO(), "api.mangacopy.com")

http/http.go

@@ -14,6 +14,7 @@ import (
 
 	"github.com/fumiama/terasu"
 	"github.com/fumiama/terasu/dns"
+	"github.com/fumiama/terasu/ip"
 )
 
 var (
@@ -50,7 +51,11 @@ var DefaultClient = http.Client{
 			if len(addrs) == 0 {
 				addrs, err = dns.DefaultResolver.LookupHost(ctx, host)
 				if err != nil {
-					addrs, err = net.DefaultResolver.LookupHost(ctx, host)
+					if ip.IsIPv6Available.Get() {
+						addrs, err = dns.IPv6Servers.LookupHostFallback(ctx, host)
+					} else {
+						addrs, err = dns.IPv4Servers.LookupHostFallback(ctx, host)
+					}
 					if err != nil {
 						return nil, err
 					}

http2/http2.go

@@ -15,6 +15,7 @@ import (
 
 	"github.com/fumiama/terasu"
 	"github.com/fumiama/terasu/dns"
+	"github.com/fumiama/terasu/ip"
 )
 
 var (
@@ -50,7 +51,11 @@ var DefaultClient = http.Client{
 			if len(addrs) == 0 {
 				addrs, err = dns.DefaultResolver.LookupHost(ctx, host)
 				if err != nil {
-					addrs, err = net.DefaultResolver.LookupHost(ctx, host)
+					if ip.IsIPv6Available.Get() {
+						addrs, err = dns.IPv6Servers.LookupHostFallback(ctx, host)
+					} else {
+						addrs, err = dns.IPv4Servers.LookupHostFallback(ctx, host)
+					}
 					if err != nil {
 						return nil, err
 					}