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  }