github.com/sandwich-go/boost@v1.3.29/httputil/dns/dns_imp.go (about)

     1  package dns
     2  
     3  import (
     4  	"context"
     5  	"github.com/sandwich-go/boost/z"
     6  	"net"
     7  	"strings"
     8  	"sync"
     9  	"time"
    10  )
    11  
    12  type dns struct {
    13  	Resolver
    14  	opts *Options
    15  }
    16  
    17  // New 创建 dns
    18  func New(opts ...Option) DNS {
    19  	cfg := NewOptions(opts...)
    20  	d := &dns{opts: cfg}
    21  	d.Resolver = d
    22  	return d
    23  }
    24  
    25  // GetDialContext 获取 Dial 函数
    26  func (d *dns) GetDialContext() func(ctx context.Context, network, address string) (net.Conn, error) {
    27  	return func(ctx context.Context, network, addr string) (net.Conn, error) {
    28  		sepIndex := strings.LastIndex(addr, ":")
    29  		ipAddrs, err := d.Resolver.LookupIPAddr(ctx, addr[:sepIndex])
    30  		if err != nil {
    31  			return nil, err
    32  		}
    33  		if len(ipAddrs) == 0 {
    34  			return nil, ErrNotFound
    35  		}
    36  		index := 0
    37  		if d.opts.GetPolicy() == PolicyRandom {
    38  			index = int(z.FastRand()) % len(ipAddrs)
    39  		}
    40  		addr = ipAddrs[index].String() + addr[sepIndex:]
    41  		return d.opts.GetDialer().DialContext(ctx, network, addr)
    42  	}
    43  }
    44  
    45  // LookupIPAddr 通过主机搜索 ip 地址列表
    46  func (d *dns) LookupIPAddr(ctx context.Context, host string) ([]net.IPAddr, error) {
    47  	var start = time.Now()
    48  	var cancel context.CancelFunc
    49  	if _, ok := ctx.Deadline(); !ok {
    50  		ctx, cancel = context.WithTimeout(ctx, d.opts.GetLookupTimeout())
    51  	}
    52  	defer func() {
    53  		if cancel != nil {
    54  			cancel()
    55  		}
    56  	}()
    57  
    58  	ipAddrs, err := d.opts.GetResolver().LookupIPAddr(ctx, host)
    59  	if err != nil {
    60  		return nil, err
    61  	}
    62  	if len(ipAddrs) == 0 {
    63  		return nil, ErrNotFound
    64  	}
    65  	// 成功则回调
    66  	if d.opts.GetOnLookup() != nil {
    67  		d.opts.GetOnLookup()(ctx, host, time.Since(start), ipAddrs)
    68  	}
    69  	return ipAddrs, nil
    70  }
    71  
    72  type cacheDns struct {
    73  	*dns
    74  	caches sync.Map
    75  }
    76  
    77  // NewCache 创建带缓存的 dns
    78  func NewCache(opts ...Option) CacheDNS {
    79  	dc := &cacheDns{dns: New(opts...).(*dns)}
    80  	dc.dns.Resolver = dc
    81  	return dc
    82  }
    83  
    84  // LookupIPAddr 通过主机搜索 ip 地址列表
    85  func (dc *cacheDns) LookupIPAddr(ctx context.Context, host string) ([]net.IPAddr, error) {
    86  	ipCache, exists := dc.Get(host)
    87  	if exists {
    88  		// 如果创建时间小于0，表示永久有效
    89  		// 如果在有效期内，直接返回
    90  		if ipCache.CreatedAt.IsZero() || ipCache.CreatedAt.Add(dc.opts.GetTTL()).After(time.Now()) {
    91  			return ipCache.IPAddrs, nil
    92  		}
    93  	}
    94  	ipAddrs, err := dc.dns.LookupIPAddr(ctx, host)
    95  	if err != nil {
    96  		return nil, err
    97  	}
    98  	if dc.opts.GetTTL() > 0 {
    99  		dc.Set(host, IpCache{IPAddrs: ipAddrs, CreatedAt: time.Now()})
   100  	} else {
   101  		dc.Set(host, IpCache{IPAddrs: ipAddrs})
   102  	}
   103  	return ipAddrs, nil
   104  }
   105  
   106  func (dc *cacheDns) Set(host string, cache IpCache) {
   107  	dc.caches.Store(host, cache)
   108  }
   109  
   110  func (dc *cacheDns) Remove(host string) {
   111  	dc.caches.Delete(host)
   112  }
   113  
   114  func (dc *cacheDns) Get(host string) (IpCache, bool) {
   115  	c, ok := dc.caches.Load(host)
   116  	if !ok {
   117  		return EmptyIpCache, false
   118  	}
   119  	return c.(IpCache), true
   120  }