github.com/xxf098/lite-proxy@v0.15.1-0.20230422081941-12c69f323218/dns/client.go (about) 1 package dns 2 3 import ( 4 "context" 5 "fmt" 6 "net" 7 8 "github.com/xxf098/lite-proxy/transport/dialer" 9 10 D "github.com/miekg/dns" 11 ) 12 13 type client struct { 14 *D.Client 15 r *Resolver 16 port string 17 host string 18 } 19 20 func (c *client) Exchange(m *D.Msg) (msg *D.Msg, err error) { 21 return c.ExchangeContext(context.Background(), m) 22 } 23 24 func (c *client) ExchangeContext(ctx context.Context, m *D.Msg) (msg *D.Msg, err error) { 25 var ip net.IP 26 if c.r == nil { 27 // a default ip dns 28 ip = net.ParseIP(c.host) 29 } else { 30 var err error 31 if ip, err = c.r.ResolveIP(c.host); err != nil { 32 return nil, fmt.Errorf("use default dns resolve failed: %w", err) 33 } 34 } 35 36 d, err := dialer.Dialer() 37 if err != nil { 38 return nil, err 39 } 40 41 c.Client.Dialer = d 42 43 // miekg/dns ExchangeContext doesn't respond to context cancel. 44 // this is a workaround 45 type result struct { 46 msg *D.Msg 47 err error 48 } 49 ch := make(chan result, 1) 50 go func() { 51 msg, _, err := c.Client.Exchange(m, net.JoinHostPort(ip.String(), c.port)) 52 ch <- result{msg, err} 53 }() 54 55 select { 56 case <-ctx.Done(): 57 return nil, ctx.Err() 58 case ret := <-ch: 59 return ret.msg, ret.err 60 } 61 }