Expose DNS resolver functionality for extensions (#5421)

* refactor: change getDialAddr to return *types.Host instead of string

This simplifies the internal API by returning the structured Host type
directly rather than converting it to a string. The caller can
calls .String() on the returned Host when needed.

This change makes it easier for future code to access the IP and port
components separately without needing to parse the string representation.

* feat: add AddrResolver interface and implement ResolveAddr method in Dialer

* fix AddrResolver interface to return port as well

Author: szkiba

lib/netext/dialer.go

@@ -61,14 +61,32 @@ func (d *Dialer) DialContext(ctx context.Context, proto, addr string) (net.Conn,
 	if err != nil {
 		return nil, err
 	}
-	conn, err := d.Dialer.DialContext(ctx, proto, dialAddr)
+	conn, err := d.Dialer.DialContext(ctx, proto, dialAddr.String())
 	if err != nil {
 		return nil, err
 	}
 	conn = &Conn{conn, &d.BytesRead, &d.BytesWritten}
 	return conn, err
 }
 
+// ResolveAddr looks up the IP address for the given host and optionally port.
+// The address is expected in the form "host:port" or just "host".
+// It returns the resolved IP, and an error if any.
+func (d *Dialer) ResolveAddr(addr string) (net.IP, int, error) {
+	// Check if the address has a port, if not add a dummy port for parsing
+	if _, _, err := net.SplitHostPort(addr); err != nil {
+		// Address doesn't have a port, add a dummy one
+		addr = net.JoinHostPort(addr, "0")
+	}
+
+	remote, err := d.getDialAddr(addr)
+	if err != nil {
+		return nil, 0, err
+	}
+
+	return remote.IP, remote.Port, nil
+}
+
 // IOSamples returns samples for data send and received since it last call and zeros out.
 // It uses the provided time as the sample time and tags and builtinMetrics to build the samples.
 func (d *Dialer) IOSamples(
@@ -98,19 +116,19 @@ func (d *Dialer) IOSamples(
 	})
 }
 
-func (d *Dialer) getDialAddr(addr string) (string, error) {
+func (d *Dialer) getDialAddr(addr string) (*types.Host, error) {
 	remote, err := d.findRemote(addr)
 	if err != nil {
-		return "", err
+		return nil, err
 	}
 
 	for _, ipnet := range d.Blacklist {
 		if ipnet.Contains(remote.IP) {
-			return "", BlackListedIPError{ip: remote.IP, net: ipnet}
+			return nil, BlackListedIPError{ip: remote.IP, net: ipnet}
 		}
 	}
 
-	return remote.String(), nil
+	return remote, nil
 }
 
 func (d *Dialer) findRemote(addr string) (*types.Host, error) {

lib/netext/dialer_test.go

@@ -66,7 +66,7 @@ func TestDialerAddr(t *testing.T) {
 				require.EqualError(t, err, tc.expErr)
 			} else {
 				require.NoError(t, err)
-				require.Equal(t, tc.expAddress, addr)
+				require.Equal(t, tc.expAddress, addr.String())
 			}
 		})
 	}
@@ -103,7 +103,7 @@ func TestDialerAddrBlockHostnamesStar(t *testing.T) {
 				require.EqualError(t, err, tc.expErr)
 			} else {
 				require.NoError(t, err)
-				require.Equal(t, tc.expAddress, addr)
+				require.Equal(t, tc.expAddress, addr.String())
 			}
 		})
 	}
@@ -135,6 +135,80 @@ func BenchmarkDialerHosts(b *testing.B) {
 	}
 }
 
+func TestDialerResolveAddr(t *testing.T) {
+	t.Parallel()
+	dialer := NewDialer(net.Dialer{}, newResolver())
+	hosts, err := types.NewHosts(
+		map[string]types.Host{
+			"example.com":                {IP: net.ParseIP("3.4.5.6")},
+			"example.com:443":            {IP: net.ParseIP("3.4.5.6"), Port: 8443},
+			"example.com:8080":           {IP: net.ParseIP("3.4.5.6"), Port: 9090},
+			"example-deny-host.com":      {IP: net.ParseIP("8.9.10.11")},
+			"example-ipv6.com":           {IP: net.ParseIP("2001:db8::68")},
+			"example-ipv6.com:443":       {IP: net.ParseIP("2001:db8::68"), Port: 8443},
+			"example-ipv6-deny-host.com": {IP: net.ParseIP("::1")},
+		})
+	require.NoError(t, err)
+	dialer.Hosts = hosts
+	ipNet, err := lib.ParseCIDR("8.9.10.0/24")
+	require.NoError(t, err)
+
+	ipV6Net, err := lib.ParseCIDR("::1/24")
+	require.NoError(t, err)
+
+	dialer.Blacklist = []*lib.IPNet{ipNet, ipV6Net}
+
+	testCases := []struct {
+		name          string
+		address       string
+		expectedIP    string
+		expectedPort  int
+		expectedError string
+	}{
+		// IPv4 with port
+		{"IPv4_with_resolver", "example-resolver.com:80", "1.2.3.4", 80, ""},
+		{"IPv4_with_hosts_mapping", "example.com:80", "3.4.5.6", 80, ""},
+		{"IPv4_with_custom_port_in_hosts", "example.com:443", "3.4.5.6", 8443, ""},
+		{"IPv4_with_different_custom_port", "example.com:8080", "3.4.5.6", 9090, ""},
+		{"IPv4_direct_IP", "1.2.3.4:80", "1.2.3.4", 80, ""},
+		{"IPv4_blacklisted_via_resolver", "example-deny-resolver.com:80", "", 0, "IP (8.9.10.11) is in a blacklisted range (8.9.10.0/24)"},
+		{"IPv4_blacklisted_via_hosts", "example-deny-host.com:80", "", 0, "IP (8.9.10.11) is in a blacklisted range (8.9.10.0/24)"},
+		{"IPv4_non-existent_host", "no-such-host.com:80", "", 0, "lookup no-such-host.com: no such host"},
+
+		// IPv4 without port
+		{"IPv4_without_port", "example-resolver.com", "1.2.3.4", 0, ""},
+		{"IPv4_hosts_mapping_without_port", "example.com", "3.4.5.6", 0, ""},
+		{"IPv4_direct_IP_without_port", "1.2.3.4", "1.2.3.4", 0, ""},
+
+		// IPv6 with port
+		{"IPv6_with_hosts_mapping", "example-ipv6.com:443", "2001:db8::68", 8443, ""},
+		{"IPv6_direct_IP", "[2001:db8:aaaa:1::100]:443", "2001:db8:aaaa:1::100", 443, ""},
+		{"IPv6_blacklisted_via_resolver", "example-ipv6-deny-resolver.com:80", "", 0, "IP (::1) is in a blacklisted range (::/24)"},
+		{"IPv6_blacklisted_via_hosts", "example-ipv6-deny-host.com:80", "", 0, "IP (::1) is in a blacklisted range (::/24)"},
+
+		// IPv6 without port
+		{"IPv6_without_port", "example-ipv6.com", "2001:db8::68", 0, ""},
+		{"IPv6_direct_IP_without_port", "2001:db8:aaaa:1::100", "2001:db8:aaaa:1::100", 0, ""},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			ip, port, err := dialer.ResolveAddr(tc.address)
+
+			if tc.expectedError != "" {
+				require.EqualError(t, err, tc.expectedError)
+				require.Nil(t, ip)
+				require.Equal(t, 0, port)
+			} else {
+				require.NoError(t, err)
+				require.Equal(t, net.ParseIP(tc.expectedIP), ip)
+				require.Equal(t, tc.expectedPort, port)
+			}
+		})
+	}
+}
+
 func newResolver() *mockresolver.MockResolver {
 	return mockresolver.New(
 		map[string][]net.IP{

lib/vu_state.go

@@ -26,6 +26,16 @@ type TracerProvider interface {
 	Tracer(name string, options ...trace.TracerOption) trace.Tracer
 }
 
+// AddrResolver is an interface for DNS resolution.
+type AddrResolver interface {
+	// ResolveAddr looks up the IP address for the given host and optionally port.
+	// It uses the same DNS resolution logic as DialContext, respecting the
+	// DNS, Hosts, Blacklist IP, and Block hostnames options.
+	// The address can be in the form "host:port" or just "host".
+	// It returns the resolved IP, the port (0 if not specified), and an error if resolution fails.
+	ResolveAddr(addr string) (net.IP, int, error)
+}
+
 // State provides the volatile state for a VU.
 //
 // TODO: rename to VUState or, better yet, move to some other Go package outside
@@ -92,6 +102,16 @@ type State struct {
 	TestStatus *TestStatus
 }
 
+// GetAddrResolver returns the AddrResolver implementation or nil if not available.
+func (s *State) GetAddrResolver() AddrResolver {
+	resolver, ok := s.Dialer.(AddrResolver)
+	if !ok {
+		return nil
+	}
+
+	return resolver
+}
+
 // VUStateTags wraps the current VU's tags and ensures a thread-safe way to
 // access and modify them exists. This is necessary because the VU tags and
 // metadata can be modified from the JS scripts via the `vu.tags` API in the

lib/vu_state_test.go

@@ -0,0 +1,59 @@
+package lib
+
+import (
+	"context"
+	"errors"
+	"net"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+// mockDialerWithResolver is a mock that implements both DialContexter and AddrResolver
+type mockDialerWithResolver struct{}
+
+func (m *mockDialerWithResolver) DialContext(_ context.Context, _, _ string) (net.Conn, error) {
+	return nil, errors.ErrUnsupported
+}
+
+func (m *mockDialerWithResolver) ResolveAddr(_ string) (net.IP, int, error) {
+	return nil, 0, errors.ErrUnsupported
+}
+
+// mockDialerWithoutResolver is a mock that only implements DialContexter
+type mockDialerWithoutResolver struct{}
+
+func (m *mockDialerWithoutResolver) DialContext(_ context.Context, _, _ string) (net.Conn, error) {
+	return nil, errors.ErrUnsupported
+}
+
+func TestGetAddrResolver(t *testing.T) {
+	t.Parallel()
+
+	t.Run("returns_same_instance_when_Dialer_implements_AddrResolver", func(t *testing.T) {
+		t.Parallel()
+
+		mock := &mockDialerWithResolver{}
+
+		state := &State{
+			Dialer: mock,
+		}
+
+		resolver := state.GetAddrResolver()
+		require.NotNil(t, resolver)
+		require.Same(t, mock, resolver)
+	})
+
+	t.Run("returns_nil_when_Dialer_does_not_implement_AddrResolver", func(t *testing.T) {
+		t.Parallel()
+
+		mock := &mockDialerWithoutResolver{}
+
+		state := &State{
+			Dialer: mock,
+		}
+
+		resolver := state.GetAddrResolver()
+		require.Nil(t, resolver)
+	})
+}