diff --git a/request.go b/request.go index 16702ca..89e4830 100644 --- a/request.go +++ b/request.go @@ -19,7 +19,7 @@ const ( var log = logging.Logger("doh") -func doRequest(ctx context.Context, url string, m *dns.Msg) (*dns.Msg, error) { +func doRequest(ctx context.Context, client *http.Client, url string, m *dns.Msg) (*dns.Msg, error) { data, err := m.Pack() if err != nil { return nil, err @@ -35,7 +35,7 @@ func doRequest(ctx context.Context, url string, m *dns.Msg) (*dns.Msg, error) { req = req.WithContext(ctx) - resp, err := http.DefaultClient.Do(req) + resp, err := client.Do(req) if err != nil { return nil, err } @@ -62,13 +62,13 @@ func doRequest(ctx context.Context, url string, m *dns.Msg) (*dns.Msg, error) { return r, nil } -func doRequestA(ctx context.Context, url string, domain string) ([]net.IPAddr, uint32, error) { +func doRequestA(ctx context.Context, client *http.Client, url string, domain string) ([]net.IPAddr, uint32, error) { fqdn := dns.Fqdn(domain) m := new(dns.Msg) m.SetQuestion(fqdn, dns.TypeA) - r, err := doRequest(ctx, url, m) + r, err := doRequest(ctx, client, url, m) if err != nil { return nil, 0, err } @@ -90,13 +90,13 @@ func doRequestA(ctx context.Context, url string, domain string) ([]net.IPAddr, u return result, ttl, nil } -func doRequestAAAA(ctx context.Context, url string, domain string) ([]net.IPAddr, uint32, error) { +func doRequestAAAA(ctx context.Context, client *http.Client, url string, domain string) ([]net.IPAddr, uint32, error) { fqdn := dns.Fqdn(domain) m := new(dns.Msg) m.SetQuestion(fqdn, dns.TypeAAAA) - r, err := doRequest(ctx, url, m) + r, err := doRequest(ctx, client, url, m) if err != nil { return nil, 0, err } @@ -119,13 +119,13 @@ func doRequestAAAA(ctx context.Context, url string, domain string) ([]net.IPAddr return result, ttl, nil } -func doRequestTXT(ctx context.Context, url string, domain string) ([]string, uint32, error) { +func doRequestTXT(ctx context.Context, client *http.Client, url string, domain string) ([]string, uint32, error) { fqdn := dns.Fqdn(domain) m := new(dns.Msg) m.SetQuestion(fqdn, dns.TypeTXT) - r, err := doRequest(ctx, url, m) + r, err := doRequest(ctx, client, url, m) if err != nil { return nil, 0, err } diff --git a/resolver.go b/resolver.go index a9bfe36..7484a48 100644 --- a/resolver.go +++ b/resolver.go @@ -5,6 +5,7 @@ import ( "errors" "math" "net" + "net/http" "strings" "sync" "time" @@ -15,8 +16,9 @@ import ( ) type Resolver struct { - mx sync.Mutex - url string + mx sync.Mutex + url string + httpClient *http.Client // RR cache ipCache map[string]ipAddrEntry @@ -52,6 +54,19 @@ func WithCacheDisabled() Option { } } +// WithHTTPClient sets the http.Client used for DoH requests, letting callers +// control transport, timeouts, proxy, or TLS settings. When unset, the resolver +// uses http.DefaultClient. +func WithHTTPClient(client *http.Client) Option { + return func(tr *Resolver) error { + if client == nil { + return errors.New("http client must not be nil") + } + tr.httpClient = client + return nil + } +} + func NewResolver(url string, opts ...Option) (*Resolver, error) { if strings.HasPrefix(url, "http:") && !strings.HasPrefix(url, "http://localhost") && @@ -66,6 +81,7 @@ func NewResolver(url string, opts ...Option) (*Resolver, error) { r := &Resolver{ url: url, + httpClient: http.DefaultClient, ipCache: make(map[string]ipAddrEntry), txtCache: make(map[string]txtEntry), maxCacheTTL: time.Duration(math.MaxUint32) * time.Second, @@ -96,12 +112,12 @@ func (r *Resolver) LookupIPAddr(ctx context.Context, domain string) (result []ne resch := make(chan response, 2) go func() { - ip4, ttl, err := doRequestA(ctx, r.url, domain) + ip4, ttl, err := doRequestA(ctx, r.httpClient, r.url, domain) resch <- response{ip4, ttl, err} }() go func() { - ip6, ttl, err := doRequestAAAA(ctx, r.url, domain) + ip6, ttl, err := doRequestAAAA(ctx, r.httpClient, r.url, domain) resch <- response{ip6, ttl, err} }() @@ -129,7 +145,7 @@ func (r *Resolver) LookupTXT(ctx context.Context, domain string) ([]string, erro return result, nil } - result, ttl, err := doRequestTXT(ctx, r.url, domain) + result, ttl, err := doRequestTXT(ctx, r.httpClient, r.url, domain) if err != nil { return nil, err } diff --git a/resolver_test.go b/resolver_test.go index a58b1ff..b1d525c 100644 --- a/resolver_test.go +++ b/resolver_test.go @@ -6,6 +6,7 @@ import ( "net" "net/http" "net/http/httptest" + "sync/atomic" "testing" "time" @@ -216,6 +217,46 @@ func TestCleartextLocalhostEndpoint(t *testing.T) { } } +// countingRoundTripper counts the requests passing through it before delegating +// to the wrapped transport. +type countingRoundTripper struct { + rt http.RoundTripper + count atomic.Int64 +} + +func (c *countingRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) { + c.count.Add(1) + return c.rt.RoundTrip(req) +} + +func TestWithHTTPClient(t *testing.T) { + domain := "example.com" + resolver := mockDoHResolver(t, map[uint16]*dns.Msg{ + dns.TypeTXT: mockDNSAnswerTXT(dns.Fqdn(domain), []string{"dnslink=/ipns/example.com"}), + }) + defer resolver.Close() + + rt := &countingRoundTripper{rt: http.DefaultTransport} + r, err := NewResolver(resolver.URL, WithHTTPClient(&http.Client{Transport: rt})) + if err != nil { + t.Fatal("resolver cannot be initialised") + } + + if _, err := r.LookupTXT(context.Background(), domain); err != nil { + t.Fatal(err) + } + if rt.count.Load() == 0 { + t.Fatal("expected the custom http client to be used") + } +} + +func TestWithHTTPClientNil(t *testing.T) { + _, err := NewResolver("https://cloudflare-dns.com/dns-query", WithHTTPClient(nil)) + if err == nil { + t.Fatal("expected an error when passing a nil http client") + } +} + func sameIPs(a, b []net.IPAddr) bool { if len(a) != len(b) { return false