diff --git a/client.go b/client.go index f7574bc..e475c7c 100644 --- a/client.go +++ b/client.go @@ -141,27 +141,23 @@ func NewWithOptions(options Options) (*Client, error) { client.proxyDialer = contextDialer } - // With a proxy, every query goes over TCP through it, so UDP pools would - // only dial resolvers directly. - if options.ConnectionPoolThreads > 1 && client.proxyDialer == nil { + // Pool exactly the resolvers that exchange sends through udpClient: UDP + // resolvers, and only without a proxy, since with one clientFor sends them + // over TCP. + if options.ConnectionPoolThreads > 1 { client.udpConnPool = mapsutil.SyncLockMap[string, *ConnPool]{ Map: make(mapsutil.Map[string, *ConnPool]), } for _, resolver := range client.resolvers { - resolverHost, resolverPort, err := net.SplitHostPort(resolver.String()) - if err != nil { - return nil, err - } - networkResolver := NetworkResolver{ - Protocol: UDP, - Port: resolverPort, - Host: resolverHost, + r, ok := resolver.(*NetworkResolver) + if !ok || client.clientFor(r.Protocol) != client.udpClient { + continue } - udpConnPool, err := NewConnPool(networkResolver, options.ConnectionPoolThreads) + udpConnPool, err := NewConnPool(*r, options.ConnectionPoolThreads) if err != nil { return nil, err } - _ = client.udpConnPool.Set(resolver.String(), udpConnPool) + _ = client.udpConnPool.Set(r.String(), udpConnPool) } } return &client, nil diff --git a/connpool_test.go b/connpool_test.go index 937dae1..5a677a1 100644 --- a/connpool_test.go +++ b/connpool_test.go @@ -181,3 +181,25 @@ func TestConnPoolRepeatedClose(t *testing.T) { t.Fatal("repeated pool close did not return") } } + +// TestNewWithOptionsPoolsOnlyUDPResolvers ensures NewWithOptions builds +// connection pools only for UDP resolvers, the only ones whose queries use +// them. It used to dial a DoH URL as a UDP address and fail. +func TestNewWithOptionsPoolsOnlyUDPResolvers(t *testing.T) { + udp := newFakeResolver(t, dns.RcodeSuccess) + doh := newFakeDoHResolver(t, dns.RcodeSuccess) + client, err := NewWithOptions(Options{ + BaseResolvers: []string{ + udp.Addr(), + "tcp:" + freeUDPAddr(t), + "dot:" + freeUDPAddr(t), + doh.ResolverString("post"), + }, + MaxRetries: 1, + ConnectionPoolThreads: 2, + }) + require.NoError(t, err) + t.Cleanup(client.Close) + require.Len(t, client.udpConnPool.Map, 1) + require.Contains(t, client.udpConnPool.Map, udp.Addr()) +}