From ecec5a9e4d4d25a737f4ec1e4352ccac95e5a244 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=BA=90=E6=96=87=E9=9B=A8?= <41315874+fumiama@users.noreply.github.com> Date: Sun, 21 Apr 2024 01:08:19 +0900 Subject: [PATCH] feat(dns): add lookup host fallback --- dns/dns.go | 42 ++++++++++++++++++++++++++++++++++++++++-- dns/dns_test.go | 10 +++++++--- http/http.go | 7 ++++++- http2/http2.go | 7 ++++++- 4 files changed, 59 insertions(+), 7 deletions(-) diff --git a/dns/dns.go b/dns/dns.go index f77050b..d5b27bc 100644 --- a/dns/dns.go +++ b/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{ diff --git a/dns/dns_test.go b/dns/dns_test.go index 8be8a7c..b17280c 100644 --- a/dns/dns_test.go +++ b/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") diff --git a/http/http.go b/http/http.go index 0ab4c25..0e28f38 100644 --- a/http/http.go +++ b/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 } diff --git a/http2/http2.go b/http2/http2.go index 6a49ea2..25d4f98 100644 --- a/http2/http2.go +++ b/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 }