github.com/Asutorufa/yuhaiin@v0.3.6-0.20240502055049-7984da7023a0/pkg/net/netapi/dns.go (about)

     1  package netapi
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"io"
     7  	"net"
     8  
     9  	"github.com/Asutorufa/yuhaiin/pkg/net/dialer"
    10  	"golang.org/x/net/dns/dnsmessage"
    11  )
    12  
    13  type LookupIPOption struct {
    14  	AAAA bool
    15  	A    bool
    16  }
    17  
    18  type ForceFakeIP struct{}
    19  
    20  type Resolver interface {
    21  	LookupIP(ctx context.Context, domain string, opts ...func(*LookupIPOption)) ([]net.IP, error)
    22  	Raw(ctx context.Context, req dnsmessage.Question) (dnsmessage.Message, error)
    23  	io.Closer
    24  }
    25  
    26  var _ Resolver = (*ErrorResolver)(nil)
    27  
    28  type ErrorResolver func(domain string) error
    29  
    30  func (e ErrorResolver) LookupIP(_ context.Context, domain string, opts ...func(*LookupIPOption)) ([]net.IP, error) {
    31  	return nil, e(domain)
    32  }
    33  func (e ErrorResolver) Close() error { return nil }
    34  func (e ErrorResolver) Raw(_ context.Context, req dnsmessage.Question) (dnsmessage.Message, error) {
    35  	return dnsmessage.Message{}, e(req.Name.String())
    36  }
    37  
    38  var InternetResolver Resolver = NewSystemResolver("8.8.8.8:53", "1.1.1.1:53", "223.5.5.5:53", "114.114.114.114:53")
    39  
    40  var Bootstrap Resolver = InternetResolver
    41  
    42  type SystemResolver struct {
    43  	resolver *net.Resolver
    44  }
    45  
    46  func NewSystemResolver(host ...string) *SystemResolver {
    47  	return &SystemResolver{
    48  		resolver: &net.Resolver{
    49  			Dial: func(ctx context.Context, network, address string) (net.Conn, error) {
    50  				for _, h := range host {
    51  					conn, err := dialer.DialContext(ctx, network, h)
    52  					if err == nil {
    53  						return conn, nil
    54  					}
    55  				}
    56  				return nil, fmt.Errorf("system dailer failed")
    57  			},
    58  		},
    59  	}
    60  }
    61  
    62  func (d *SystemResolver) LookupIP(ctx context.Context, domain string, opts ...func(*LookupIPOption)) ([]net.IP, error) {
    63  	network := "ip"
    64  
    65  	opt := &LookupIPOption{
    66  		A: true,
    67  	}
    68  
    69  	for _, o := range opts {
    70  		o(opt)
    71  	}
    72  
    73  	if opt.AAAA && !opt.A {
    74  		network = "ip6"
    75  	}
    76  
    77  	if opt.A && !opt.AAAA {
    78  		network = "ip4"
    79  	}
    80  
    81  	return d.resolver.LookupIP(ctx, network, domain)
    82  }
    83  func (d *SystemResolver) Raw(context.Context, dnsmessage.Question) (dnsmessage.Message, error) {
    84  	return dnsmessage.Message{}, fmt.Errorf("system dns not support")
    85  }
    86  func (d *SystemResolver) Close() error { return nil }
    87  
    88  type DNSErrCode struct {
    89  	code dnsmessage.RCode
    90  }
    91  
    92  func NewDNSErrCode(code dnsmessage.RCode) *DNSErrCode {
    93  	return &DNSErrCode{
    94  		code: code,
    95  	}
    96  }
    97  
    98  func (d *DNSErrCode) Code() dnsmessage.RCode {
    99  	return d.code
   100  }
   101  
   102  func (d DNSErrCode) Error() string {
   103  	return d.code.String()
   104  }
   105  
   106  func (d *DNSErrCode) As(err any) bool {
   107  	dd, ok := err.(*DNSErrCode)
   108  	if ok {
   109  		dd.code = d.code
   110  	}
   111  
   112  	return ok
   113  }