From a8176b0a14ef5d488ba7e300dfe81105297d7a44 Mon Sep 17 00:00:00 2001 From: ManuelReschke Date: Thu, 30 Jul 2026 20:22:51 +0200 Subject: [PATCH] feat: add SetResolver and SetHosts for custom DNS (#169) --- client.go | 114 +++++++++++++++ client_test.go | 276 ++++++++++++++++++++++++++++++++++++ client_wrapper.go | 12 ++ internal/http2/transport.go | 74 ++++++++-- transport.go | 23 ++- 5 files changed, 484 insertions(+), 15 deletions(-) diff --git a/client.go b/client.go index ece7dada..4ce6a22d 100644 --- a/client.go +++ b/client.go @@ -8,10 +8,12 @@ import ( "encoding/json" "encoding/xml" "errors" + "fmt" "io" "net" "net/http" "net/http/cookiejar" + "net/netip" urlpkg "net/url" "os" "reflect" @@ -1490,6 +1492,118 @@ func (c *Client) SetUnixSocket(file string) *Client { }) } +// SetResolver sets a custom DNS resolver used when dialing HTTP/1 and HTTP/2 +// connections. It is implemented via SetDial and a net.Dialer that uses r. +// If r is nil, the default resolver is used. +// +// Only valid for HTTP/1 and HTTP/2 (same limitation as SetDial). Calling +// SetDial, SetHosts, or SetUnixSocket replaces this dialer. +// +// For example, use a specific DNS server: +// +// r := &net.Resolver{ +// PreferGo: true, +// Dial: func(ctx context.Context, network, address string) (net.Conn, error) { +// var d net.Dialer +// return d.DialContext(ctx, "udp", "1.1.1.1:53") +// }, +// } +// client.SetResolver(r) +func (c *Client) SetResolver(r *net.Resolver) *Client { + return c.SetDial(func(ctx context.Context, network, addr string) (net.Conn, error) { + d := net.Dialer{Resolver: r} + return d.DialContext(ctx, network, addr) + }) +} + +// SetHosts configures a static hostname-to-IP mapping used when dialing HTTP/1 +// and HTTP/2 connections (like a hosts file). Hostnames not present in the map +// fail immediately with a "no such host" DNS error and do not consult the +// system resolver. This avoids long DNS timeouts when maintaining a custom +// host list (e.g. crawlers). +// +// Notes: +// - Keys are hostnames only (no port). Matching is case-insensitive and +// IDNA-normalized to match the dial address form used by the transport. +// - Values must be literal IP addresses (IPv4 or IPv6). Optional surrounding +// brackets on IPv6 (e.g. "[::1]") are accepted and normalized. Non-IP +// values never fall through to system DNS; dialing that host returns a +// clear error instead. +// - IP-literal request addresses (e.g. https://1.2.3.4/) skip the map and +// dial directly. Scoped IPv6 literals (e.g. https://[fe80::1%25eth0]/) +// are supported. +// - An empty or nil map makes every non-literal hostname fail closed. +// - Proxy routing is rejected while SetHosts is active because a proxy can +// resolve the destination remotely and bypass the static mapping. +// - Only valid for HTTP/1 and HTTP/2 (same limitation as SetDial). SetDialTLS +// still bypasses this dialer for HTTPS when set. Calling SetDial, +// SetResolver, or SetUnixSocket replaces this dialer. +// - The map is copied; later changes to the caller's map are ignored. +// +// For example: +// +// client.SetHosts(map[string]string{ +// "api.internal": "10.0.0.5", +// "db.internal": "10.0.0.6", +// "v6.internal": "::1", +// }) +func (c *Client) SetHosts(hosts map[string]string) *Client { + m := make(map[string]string, len(hosts)) + invalid := make(map[string]string) + for host, ipStr := range hosts { + key := hostsMapKey(host) + if key == "" { + continue + } + ip := net.ParseIP(strings.Trim(ipStr, "[]")) + if ip == nil { + invalid[key] = ipStr + continue + } + m[key] = ip.String() + } + c.SetDial(func(ctx context.Context, network, addr string) (net.Conn, error) { + host, port, err := net.SplitHostPort(addr) + if err != nil { + return nil, err + } + // IP literals need no DNS and are not subject to the hosts map. + if _, err := netip.ParseAddr(host); err == nil { + var d net.Dialer + return d.DialContext(ctx, network, addr) + } + key := hostsMapKey(host) + if raw, bad := invalid[key]; bad { + return nil, fmt.Errorf("req: SetHosts: invalid IP address %q for host %q", raw, host) + } + ip, ok := m[key] + if !ok { + return nil, &net.DNSError{ + Err: "no such host", + Name: host, + IsNotFound: true, + } + } + var d net.Dialer + return d.DialContext(ctx, network, net.JoinHostPort(ip, port)) + }) + c.Transport.rejectProxyWithSetHosts = true + return c +} + +// hostsMapKey normalizes a hostname for SetHosts lookup: trim, IDNA ToASCII, +// then lowercase so keys match dial addresses produced by the transport. +func hostsMapKey(host string) string { + host = strings.TrimSpace(host) + if host == "" { + return "" + } + if ascii, err := idnaASCII(host); err == nil { + host = ascii + } + return strings.ToLower(host) +} + // DisableHTTP3 disables the http3 protocol. func (c *Client) DisableHTTP3() *Client { c.Transport.DisableHTTP3() diff --git a/client_test.go b/client_test.go index c343b21d..b354ad86 100644 --- a/client_test.go +++ b/client_test.go @@ -108,6 +108,282 @@ func TestSetDialTLS(t *testing.T) { tests.AssertEqual(t, testErr, err) } +func TestSetResolver(t *testing.T) { + ts := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte("ok")) + })) + defer ts.Close() + + // Broken custom resolver: hostname lookups fail, but IP literals still work. + r := &net.Resolver{ + PreferGo: true, + Dial: func(ctx context.Context, network, address string) (net.Conn, error) { + return nil, errors.New("custom resolver down") + }, + } + c := C().SetResolver(r).EnableInsecureSkipVerify() + tests.AssertNotNil(t, c.DialContext) + + resp, err := c.R().Get(ts.URL) + assertSuccess(t, resp, err) + tests.AssertEqual(t, "ok", resp.String()) + + // Hostname must go through the custom resolver and fail. + _, err = c.SetTimeout(3 * time.Second).R().Get("https://example.invalid/") + tests.AssertNotNil(t, err) + tests.AssertErrorContains(t, err, "custom resolver down") + + // nil resolver uses the default resolver (still installs a dialer). + c2 := C().SetResolver(nil) + tests.AssertNotNil(t, c2.DialContext) +} + +func TestSetHosts(t *testing.T) { + ts := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte("hosts-ok")) + })) + defer ts.Close() + + ip, port, err := net.SplitHostPort(ts.Listener.Addr().String()) + tests.AssertNoError(t, err) + + c := C(). + EnableInsecureSkipVerify(). + EnableForceHTTP1(). + SetHosts(map[string]string{ + "Custom.Example": ip, // case-insensitive match + }) + + resp, err := c.R().Get(fmt.Sprintf("https://custom.example:%s/", port)) + assertSuccess(t, resp, err) + tests.AssertEqual(t, "hosts-ok", resp.String()) +} + +func TestSetHostsHTTP2(t *testing.T) { + ts := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte("h2-ok")) + })) + ts.EnableHTTP2 = true + ts.StartTLS() + defer ts.Close() + + ip, port, err := net.SplitHostPort(ts.Listener.Addr().String()) + tests.AssertNoError(t, err) + + c := C(). + EnableInsecureSkipVerify(). + EnableForceHTTP2(). + SetHosts(map[string]string{ + "h2.example": ip, + }) + + resp, err := c.R().Get(fmt.Sprintf("https://h2.example:%s/", port)) + assertSuccess(t, resp, err) + tests.AssertEqual(t, "h2-ok", resp.String()) + tests.AssertEqual(t, "HTTP/2.0", resp.Proto) +} + +func TestSetHostsUnknownHostFailsFast(t *testing.T) { + c := C(). + EnableInsecureSkipVerify(). + SetTimeout(5 * time.Second). + SetHosts(map[string]string{ + "known.example": "127.0.0.1", + }) + + start := time.Now() + _, err := c.R().Get("https://unknown.example/") + elapsed := time.Since(start) + + tests.AssertNotNil(t, err) + tests.AssertErrorContains(t, err, "no such host") + // Must not wait on system DNS timeouts. + if elapsed > 2*time.Second { + t.Fatalf("unknown host took too long: %v", elapsed) + } +} + +func TestSetHostsEmptyMapFailsFast(t *testing.T) { + for _, hosts := range []map[string]string{nil, {}} { + c := C().SetHosts(hosts).SetTimeout(2 * time.Second) + start := time.Now() + _, err := c.DialContext(context.Background(), "tcp", "any.example:80") + elapsed := time.Since(start) + tests.AssertNotNil(t, err) + tests.AssertErrorContains(t, err, "no such host") + if elapsed > time.Second { + t.Fatalf("empty hosts map dial took too long: %v (hosts=%v)", elapsed, hosts) + } + } +} + +func TestSetHostsIgnoresCallerMapMutation(t *testing.T) { + ts := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = w.Write([]byte("ok")) + })) + defer ts.Close() + + ip, port, err := net.SplitHostPort(ts.Listener.Addr().String()) + tests.AssertNoError(t, err) + + hosts := map[string]string{"static.example": ip} + c := C().EnableInsecureSkipVerify().EnableForceHTTP1().SetHosts(hosts) + // Mutating the original map after SetHosts must not affect the client. + delete(hosts, "static.example") + hosts["static.example"] = "0.0.0.0" + + resp, err := c.R().Get(fmt.Sprintf("https://static.example:%s/", port)) + assertSuccess(t, resp, err) + tests.AssertEqual(t, "ok", resp.String()) +} + +func TestSetHostsReplacedBySetDial(t *testing.T) { + called := false + c := C(). + SetTimeout(2 * time.Second). + SetHosts(map[string]string{"x.example": "127.0.0.1"}). + SetDial(func(ctx context.Context, network, addr string) (net.Conn, error) { + called = true + return nil, errors.New("custom dial") + }) + + _, err := c.R().Get("https://x.example/") + tests.AssertNotNil(t, err) + tests.AssertEqual(t, true, called) + tests.AssertErrorContains(t, err, "custom dial") +} + +func TestSetHostsIPLiteralPassthrough(t *testing.T) { + ts := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = w.Write([]byte("literal-ok")) + })) + defer ts.Close() + + // Map does not include the server IP; IP-literal URLs must still work. + c := C(). + EnableInsecureSkipVerify(). + SetHosts(map[string]string{ + "only.example": "10.0.0.1", + }) + + resp, err := c.R().Get(ts.URL) + assertSuccess(t, resp, err) + tests.AssertEqual(t, "literal-ok", resp.String()) +} + +func TestSetHostsScopedIPv6LiteralPassthrough(t *testing.T) { + c := C().SetHosts(nil) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + _, err := c.DialContext(ctx, "tcp", "[fe80::1%eth0]:443") + tests.AssertNotNil(t, err) + if !errors.Is(err, context.Canceled) { + t.Fatalf("scoped IPv6 literal was not dialed directly: %v", err) + } +} + +func TestSetHostsRejectsProxy(t *testing.T) { + proxyURL := "http://127.0.0.1:1" + cases := []struct { + name string + client *Client + }{ + { + name: "proxy configured before SetHosts", + client: C(). + SetProxyURL(proxyURL). + SetHosts(map[string]string{"allowed.example": "127.0.0.1"}), + }, + { + name: "proxy configured after SetHosts", + client: C(). + SetHosts(map[string]string{"allowed.example": "127.0.0.1"}). + SetProxyURL(proxyURL), + }, + } + for _, tt := range cases { + t.Run(tt.name, func(t *testing.T) { + for _, target := range []string{"http://allowed.example/", "http://unknown.example/"} { + _, err := tt.client.R().Get(target) + tests.AssertNotNil(t, err) + tests.AssertErrorContains(t, err, "SetHosts cannot be used with a proxy") + } + }) + } +} + +func TestSetHostsInvalidIPNoDNS(t *testing.T) { + c := C().SetHosts(map[string]string{ + "bad.example": "not-an-ip", + }) + + start := time.Now() + _, err := c.DialContext(context.Background(), "tcp", "bad.example:80") + elapsed := time.Since(start) + + tests.AssertNotNil(t, err) + tests.AssertErrorContains(t, err, "invalid IP address") + tests.AssertErrorContains(t, err, "not-an-ip") + if elapsed > time.Second { + t.Fatalf("invalid IP dial took too long (possible DNS fallback): %v", elapsed) + } +} + +func TestSetHostsIPv6BracketedValue(t *testing.T) { + // Bracketed values must normalize to a single JoinHostPort form, not [[::1]]:port. + ip := net.ParseIP(strings.Trim("[::1]", "[]")) + tests.AssertNotNil(t, ip) + tests.AssertEqual(t, "[::1]:443", net.JoinHostPort(ip.String(), "443")) + + // Dial with bracketed and plain map values must not produce address-parse errors. + // Connection may fail (no listener / no IPv6), but the address form must be valid. + for _, value := range []string{"[::1]", "::1"} { + c := C().SetHosts(map[string]string{"v6.example": value}) + _, err := c.DialContext(context.Background(), "tcp", "v6.example:1") + if err != nil { + msg := err.Error() + if strings.Contains(msg, "too many colons") || + strings.Contains(msg, "invalid IP") || + strings.Contains(msg, "[[") { + t.Fatalf("value %q produced bad address form: %v", value, err) + } + } + } + + // If IPv6 loopback is available, do a full connect as well. + ln, err := net.Listen("tcp", "[::1]:0") + if err != nil { + t.Logf("skipping live IPv6 connect: %v", err) + return + } + defer ln.Close() + _, port, err := net.SplitHostPort(ln.Addr().String()) + tests.AssertNoError(t, err) + go func() { + conn, err := ln.Accept() + if err == nil { + _ = conn.Close() + } + }() + c := C().SetHosts(map[string]string{"v6.example": "[::1]"}) + conn, err := c.DialContext(context.Background(), "tcp", net.JoinHostPort("v6.example", port)) + tests.AssertNoError(t, err) + _ = conn.Close() +} + +func TestHostsMapKeyIDNA(t *testing.T) { + // Non-ASCII host should normalize to the same key as its punycode form. + unicodeHost := "bücher.example" + asciiHost, err := idnaASCII(unicodeHost) + tests.AssertNoError(t, err) + tests.AssertEqual(t, hostsMapKey(unicodeHost), hostsMapKey(asciiHost)) + tests.AssertEqual(t, hostsMapKey(asciiHost), strings.ToLower(asciiHost)) +} + func TestSetFuncs(t *testing.T) { testErr := errors.New("test") marshalFunc := func(v any) ([]byte, error) { diff --git a/client_wrapper.go b/client_wrapper.go index cd4a9ca9..987ed04a 100644 --- a/client_wrapper.go +++ b/client_wrapper.go @@ -747,6 +747,18 @@ func SetUnixSocket(file string) *Client { return defaultClient.SetUnixSocket(file) } +// SetResolver is a global wrapper methods which delegated +// to the default client's Client.SetResolver. +func SetResolver(r *net.Resolver) *Client { + return defaultClient.SetResolver(r) +} + +// SetHosts is a global wrapper methods which delegated +// to the default client's Client.SetHosts. +func SetHosts(hosts map[string]string) *Client { + return defaultClient.SetHosts(hosts) +} + // SetTLSFingerprint is a global wrapper methods which delegated // to the default client's Client.SetTLSFingerprint. func SetTLSFingerprint(clientHelloID utls.ClientHelloID) *Client { diff --git a/internal/http2/transport.go b/internal/http2/transport.go index 25759d86..4caf07e5 100644 --- a/internal/http2/transport.go +++ b/internal/http2/transport.go @@ -606,16 +606,30 @@ func (tlsHandshakeTimeoutError) Timeout() bool { return true } func (tlsHandshakeTimeoutError) Temporary() bool { return true } func (tlsHandshakeTimeoutError) Error() string { return "net/http: TLS handshake timeout" } +// dialTCP dials a TCP connection, preferring a custom DialContext when set. +func (t *Transport) dialTCP(ctx context.Context, network, addr string) (net.Conn, error) { + if t.DialContext != nil { + c, err := t.DialContext(ctx, network, addr) + if c == nil && err == nil { + err = errors.New("net/http: Transport.DialContext hook returned (nil, nil)") + } + return c, err + } + return zeroDialer.DialContext(ctx, network, addr) +} + // dialTLSWithContext uses tls.Dialer, added in Go 1.15, to open a TLS -// connection. +// connection. When DialContext is set (e.g. via Client.SetDial / SetResolver / +// SetHosts), the custom dialer is used for the underlying TCP connection. func (t *Transport) dialTLSWithContext(ctx context.Context, network, addr string, cfg *tls.Config) (reqtls.Conn, error) { if t.TLSHandshakeContext != nil { - conn, err := zeroDialer.DialContext(ctx, network, addr) + conn, err := t.dialTCP(ctx, network, addr) if err != nil { return nil, err } var firstTLSHost string if firstTLSHost, _, err = net.SplitHostPort(addr); err != nil { + conn.Close() return nil, err } trace := httptrace.ContextClientTrace(ctx) @@ -653,17 +667,61 @@ func (t *Transport) dialTLSWithContext(ctx context.Context, network, addr string tlsCn := conn.(reqtls.Conn) return tlsCn, nil } - } else { - dialer := &tls.Dialer{ - Config: cfg, - } - conn, err := dialer.DialContext(ctx, network, addr) + } + + // Custom DialContext: dial TCP first, then perform TLS handshake so + // SetResolver / SetHosts / SetDial apply to HTTP/2 as well. + // Handshake timeout / trace handling matches HTTP/1 persistConn.addTLS. + if t.DialContext != nil { + conn, err := t.dialTCP(ctx, network, addr) if err != nil { return nil, err } - tlsCn := conn.(reqtls.Conn) + tlsCn := tls.Client(conn, cfg) + trace := httptrace.ContextClientTrace(ctx) + errc := make(chan error, 2) + var timer *time.Timer + if d := t.TLSHandshakeTimeout; d != 0 { + timer = time.AfterFunc(d, func() { + errc <- tlsHandshakeTimeoutError{} + }) + } + go func() { + if trace != nil && trace.TLSHandshakeStart != nil { + trace.TLSHandshakeStart() + } + err := tlsCn.HandshakeContext(ctx) + if timer != nil { + timer.Stop() + } + errc <- err + }() + if err := <-errc; err != nil { + conn.Close() + if err == (tlsHandshakeTimeoutError{}) { + // Wait for HandshakeContext to return after close. + <-errc + } + if trace != nil && trace.TLSHandshakeDone != nil { + trace.TLSHandshakeDone(tls.ConnectionState{}, err) + } + return nil, err + } + if trace != nil && trace.TLSHandshakeDone != nil { + trace.TLSHandshakeDone(tlsCn.ConnectionState(), nil) + } return tlsCn, nil } + + dialer := &tls.Dialer{ + Config: cfg, + } + conn, err := dialer.DialContext(ctx, network, addr) + if err != nil { + return nil, err + } + tlsCn := conn.(reqtls.Conn) + return tlsCn, nil } func (t *Transport) dialTLS(ctx context.Context) func(string, string, *tls.Config) (net.Conn, error) { diff --git a/transport.go b/transport.go index 75ef99f5..1fbf4e7a 100644 --- a/transport.go +++ b/transport.go @@ -129,6 +129,10 @@ type Transport struct { // Force using specific http version forceHttpVersion httpVersion + // rejectProxyWithSetHosts prevents proxy-side DNS from bypassing the + // fail-closed static host mapping installed by Client.SetHosts. + rejectProxyWithSetHosts bool + transport.Options t2 *h2internal.Transport // non-nil if http2 wired up @@ -482,6 +486,7 @@ func (t *Transport) SetProxy(proxy func(*http.Request) (*url.URL, error)) *Trans // earlier connection becomes idle before the later dial function completes. func (t *Transport) SetDial(fn func(ctx context.Context, network, addr string) (net.Conn, error)) *Transport { t.DialContext = fn + t.rejectProxyWithSetHosts = false return t } @@ -743,13 +748,14 @@ func (t *Transport) readBufferSize() int { // Clone returns a deep copy of t's exported fields. func (t *Transport) Clone() *Transport { tt := &Transport{ - Headers: t.Headers.Clone(), - Cookies: cloneSlice(t.Cookies), - Options: t.Options.Clone(), - disableAutoDecode: t.disableAutoDecode, - autoDecodeContentType: t.autoDecodeContentType, - forceHttpVersion: t.forceHttpVersion, - httpRoundTripWrappers: t.httpRoundTripWrappers, + Headers: t.Headers.Clone(), + Cookies: cloneSlice(t.Cookies), + Options: t.Options.Clone(), + disableAutoDecode: t.disableAutoDecode, + autoDecodeContentType: t.autoDecodeContentType, + forceHttpVersion: t.forceHttpVersion, + rejectProxyWithSetHosts: t.rejectProxyWithSetHosts, + httpRoundTripWrappers: t.httpRoundTripWrappers, } if len(tt.httpRoundTripWrappers) > 0 { // clone transport middleware fn := func(req *http.Request) (*http.Response, error) { @@ -1274,6 +1280,9 @@ func (t *Transport) connectMethodForRequest(treq *transportRequest) (cm connectM cm.targetAddr = canonicalAddr(treq.URL) if t.Proxy != nil { cm.proxyURL, err = t.Proxy(treq.Request) + if err == nil && cm.proxyURL != nil && t.rejectProxyWithSetHosts { + err = errors.New("req: SetHosts cannot be used with a proxy") + } } cm.onlyH1 = t.forceHttpVersion == h1 || requestRequiresHTTP1(treq.Request) return cm, err