Skip unroutable IP families before dialing

internal/dialer/dialer.go

@@ -5,18 +5,34 @@ import (
 	"errors"
 	"fmt"
 	"net"
+	"sync"
 	"syscall"
+	"time"
 
 	"github.com/rs/dnscache"
 )
 
 // ErrForbiddenRequest is returned when a request is made to a blocked IP.
 var ErrForbiddenRequest = errors.New("forbidden")
 
+// ErrNoUsableAddress is returned when DNS resolution succeeds but none of the
+// returned addresses can be used for the requested network.
+var ErrNoUsableAddress = errors.New("no usable address")
+
+type hostResolver interface {
+	LookupHost(ctx context.Context, host string) (addrs []string, err error)
+}
+
+type dialContextFunc func(ctx context.Context, network, address string) (net.Conn, error)
+
 // Dialer is a wrapper around net.Dialer that uses a dnscache.Resolver to cache DNS lookups.
 type Dialer struct {
 	net.Dialer
-	resolver *dnscache.Resolver
+	resolver      hostResolver
+	blockedIps    []net.IP
+	dialContext   dialContextFunc
+	ipv4Available func() bool
+	ipv6Available func() bool
 }
 
 // New creates a new Dialer.
@@ -25,7 +41,10 @@ func New(resolver *dnscache.Resolver, blockedIps []net.IP) *Dialer {
 		Dialer: net.Dialer{
 			Control: safeControl(blockedIps),
 		},
-		resolver: resolver,
+		resolver:      resolver,
+		blockedIps:    blockedIps,
+		ipv4Available: defaultIPv4Available,
+		ipv6Available: defaultIPv6Available,
 	}
 }
 
@@ -46,15 +65,129 @@ func (d *Dialer) DialContext(ctx context.Context, network, address string) (conn
 	if err != nil {
 		return nil, err
 	}
+	ips, err = d.usableIPs(network, host, ips)
+	if err != nil {
+		return nil, err
+	}
 	for _, ip := range ips {
-		conn, err = d.Dialer.DialContext(ctx, network, net.JoinHostPort(ip, port))
+		conn, err = d.dial(ctx, network, net.JoinHostPort(ip, port))
 		if err == nil {
 			break
 		}
 	}
 	return
 }
 
+func (d *Dialer) dial(ctx context.Context, network, address string) (net.Conn, error) {
+	if d.dialContext != nil {
+		return d.dialContext(ctx, network, address)
+	}
+	return d.Dialer.DialContext(ctx, network, address)
+}
+
+func (d *Dialer) usableIPs(network, host string, ips []string) ([]string, error) {
+	usable := make([]string, 0, len(ips))
+	for _, ip := range ips {
+		parsedIP := net.ParseIP(ip)
+		if parsedIP == nil {
+			usable = append(usable, ip)
+			continue
+		}
+
+		if d.isBlockedIP(parsedIP) {
+			return nil, ErrForbiddenRequest
+		}
+
+		if !d.usableNetworkIP(network, parsedIP) {
+			continue
+		}
+
+		usable = append(usable, ip)
+	}
+
+	if len(usable) == 0 && len(ips) > 0 {
+		return nil, fmt.Errorf("%w: %s has no addresses usable on %s", ErrNoUsableAddress, host, network)
+	}
+
+	return usable, nil
+}
+
+func (d *Dialer) usableNetworkIP(network string, ip net.IP) bool {
+	isIPv4 := ip.To4() != nil
+	switch network {
+	case "tcp4":
+		return isIPv4
+	case "tcp6":
+		return !isIPv4
+	case "tcp":
+		return (isIPv4 && d.supportsIPv4()) || (!isIPv4 && d.supportsIPv6())
+	default:
+		return true
+	}
+}
+
+func (d *Dialer) supportsIPv4() bool {
+	if d.ipv4Available == nil {
+		return defaultIPv4Available()
+	}
+	return d.ipv4Available()
+}
+
+func (d *Dialer) supportsIPv6() bool {
+	if d.ipv6Available == nil {
+		return defaultIPv6Available()
+	}
+	return d.ipv6Available()
+}
+
+func (d *Dialer) isBlockedIP(ip net.IP) bool {
+	for _, blockedIP := range d.blockedIps {
+		if ip.Equal(blockedIP) {
+			return true
+		}
+	}
+	return false
+}
+
+var (
+	ipv4AvailabilityOnce sync.Once
+	ipv4Availability     bool
+	ipv6AvailabilityOnce sync.Once
+	ipv6Availability     bool
+)
+
+// defaultIPv4Available reports whether the host can route IPv4 traffic. It
+// probes once and caches the result, since IPv4 availability won't change
+// while the process is running.
+func defaultIPv4Available() bool {
+	ipv4AvailabilityOnce.Do(func() {
+		ipv4Availability = checkConnectivity("udp4", "192.0.2.1:80")
+	})
+	return ipv4Availability
+}
+
+// defaultIPv6Available reports whether the host can route IPv6 traffic. It
+// probes once and caches the result, since IPv6 availability won't change
+// while the process is running.
+func defaultIPv6Available() bool {
+	ipv6AvailabilityOnce.Do(func() {
+		ipv6Availability = checkConnectivity("udp6", "[2001:db8::1]:80")
+	})
+	return ipv6Availability
+}
+
+func checkConnectivity(network, address string) bool {
+	// A connected UDP dial sends no packets; it just asks the kernel to pick a
+	// route. That returns fast when the address family is disabled or has no
+	// route, as in single-stack Docker environments.
+	conn, err := net.DialTimeout(network, address, 100*time.Millisecond)
+	if err != nil {
+		return false
+	}
+	conn.Close()
+	return true
+}
+
 type control func(network, address string, conn syscall.RawConn) error
 
 func safeControl(blockedIps []net.IP) control {

internal/dialer/dialer_test.go

@@ -0,0 +1,180 @@
+package dialer
+
+import (
+	"context"
+	"io"
+	"net"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestDialContextSkipsIPv6WhenUnavailable(t *testing.T) {
+	dialer, dialed := testDialer([]string{"2001:db8::1", "192.0.2.1"}, false)
+
+	conn, err := dialer.DialContext(context.Background(), "tcp", "example.com:443")
+	require.NoError(t, err)
+	require.NotNil(t, conn)
+
+	assert.Equal(t, []string{"192.0.2.1:443"}, *dialed)
+}
+
+func TestDialContextUsesIPv6WhenAvailable(t *testing.T) {
+	dialer, dialed := testDialer([]string{"2001:db8::1", "192.0.2.1"}, true)
+
+	conn, err := dialer.DialContext(context.Background(), "tcp", "example.com:443")
+	require.NoError(t, err)
+	require.NotNil(t, conn)
+
+	assert.Equal(t, []string{"[2001:db8::1]:443"}, *dialed)
+}
+
+func TestDialContextSkipsIPv4WhenUnavailable(t *testing.T) {
+	dialer, dialed := testStackDialer([]string{"2001:db8::1", "192.0.2.1"}, false, true)
+
+	conn, err := dialer.DialContext(context.Background(), "tcp", "example.com:443")
+	require.NoError(t, err)
+	require.NotNil(t, conn)
+
+	assert.Equal(t, []string{"[2001:db8::1]:443"}, *dialed)
+}
+
+func TestDialContextReturnsErrorWhenNeitherFamilyIsAvailable(t *testing.T) {
+	dialer, dialed := testStackDialer([]string{"2001:db8::1", "192.0.2.1"}, false, false)
+
+	conn, err := dialer.DialContext(context.Background(), "tcp", "example.com:443")
+
+	assert.Nil(t, conn)
+	assert.ErrorIs(t, err, ErrNoUsableAddress)
+	assert.Empty(t, *dialed)
+}
+
+func TestDialContextFiltersByRequestedNetwork(t *testing.T) {
+	tests := []struct {
+		name    string
+		network string
+		want    []string
+	}{
+		{
+			name:    "tcp4",
+			network: "tcp4",
+			want:    []string{"192.0.2.1:443"},
+		},
+		{
+			name:    "tcp6",
+			network: "tcp6",
+			want:    []string{"[2001:db8::1]:443"},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			dialer, dialed := testDialer([]string{"2001:db8::1", "192.0.2.1"}, true)
+
+			conn, err := dialer.DialContext(context.Background(), tt.network, "example.com:443")
+			require.NoError(t, err)
+			require.NotNil(t, conn)
+
+			assert.Equal(t, tt.want, *dialed)
+		})
+	}
+}
+
+func TestDialContextReturnsErrorWhenNoResolvedAddressesAreUsable(t *testing.T) {
+	dialer, dialed := testDialer([]string{"2001:db8::1"}, false)
+
+	conn, err := dialer.DialContext(context.Background(), "tcp", "example.com:443")
+
+	assert.Nil(t, conn)
+	assert.ErrorIs(t, err, ErrNoUsableAddress)
+	assert.Empty(t, *dialed)
+}
+
+func TestDialContextRejectsBlockedIPsBeforeFilteringByNetwork(t *testing.T) {
+	dialer, dialed := testDialer([]string{"::1", "192.0.2.1"}, false)
+	dialer.blockedIps = []net.IP{net.ParseIP("::1")}
+
+	conn, err := dialer.DialContext(context.Background(), "tcp", "example.com:443")
+
+	assert.Nil(t, conn)
+	assert.ErrorIs(t, err, ErrForbiddenRequest)
+	assert.Empty(t, *dialed)
+}
+
+func testDialer(ips []string, ipv6Available bool) (*Dialer, *[]string) {
+	return testStackDialer(ips, true, ipv6Available)
+}
+
+func testStackDialer(ips []string, ipv4Available, ipv6Available bool) (*Dialer, *[]string) {
+	dialed := []string{}
+	dialer := &Dialer{
+		resolver: staticResolver{
+			ips: ips,
+		},
+		ipv4Available: func() bool {
+			return ipv4Available
+		},
+		ipv6Available: func() bool {
+			return ipv6Available
+		},
+		dialContext: func(ctx context.Context, network, address string) (net.Conn, error) {
+			dialed = append(dialed, address)
+			return fakeConn{}, nil
+		},
+	}
+	return dialer, &dialed
+}
+
+type staticResolver struct {
+	ips []string
+}
+
+func (r staticResolver) LookupHost(ctx context.Context, host string) ([]string, error) {
+	return r.ips, nil
+}
+
+type fakeConn struct{}
+
+func (fakeConn) Read(b []byte) (int, error) {
+	return 0, io.EOF
+}
+
+func (fakeConn) Write(b []byte) (int, error) {
+	return len(b), nil
+}
+
+func (fakeConn) Close() error {
+	return nil
+}
+
+func (fakeConn) LocalAddr() net.Addr {
+	return fakeAddr("local")
+}
+
+func (fakeConn) RemoteAddr() net.Addr {
+	return fakeAddr("remote")
+}
+
+func (fakeConn) SetDeadline(t time.Time) error {
+	return nil
+}
+
+func (fakeConn) SetReadDeadline(t time.Time) error {
+	return nil
+}
+
+func (fakeConn) SetWriteDeadline(t time.Time) error {
+	return nil
+}
+
+type fakeAddr string
+
+func (a fakeAddr) Network() string {
+	return string(a)
+}
+
+func (a fakeAddr) String() string {
+	return string(a)
+}