Add optional port to host mappings and fix IPv6 (#1489)

Add optional port to host mappings.

That way people can map a name under different ports.
This is very convenient when you run k6 against a fleet of docker
containers to test the same service but want to keep the port mapping.

This also fixes the IPv6 support in k6.

Signed-off-by: David Calavera <[email protected]>

js/runner_test.go

@@ -908,8 +908,8 @@ func TestVUIntegrationHosts(t *testing.T) {
 
 	r1.SetOptions(lib.Options{
 		Throw: null.BoolFrom(true),
-		Hosts: map[string]net.IP{
-			"test.loadimpact.com": net.ParseIP("127.0.0.1"),
+		Hosts: map[string]*lib.HostAddress{
+			"test.loadimpact.com": {IP: net.ParseIP("127.0.0.1")},
 		},
 	})
 

lib/netext/dialer.go

@@ -24,35 +24,46 @@ import (
 	"context"
 	"fmt"
 	"net"
-	"strings"
+	"strconv"
 	"sync/atomic"
 	"time"
 
+	"github.com/pkg/errors"
 	"github.com/viki-org/dnscache"
 
 	"github.com/loadimpact/k6/lib"
 	"github.com/loadimpact/k6/lib/metrics"
 	"github.com/loadimpact/k6/stats"
 )
 
+// dnsResolver is an interface that fetches dns information
+// about a given address.
+type dnsResolver interface {
+	FetchOne(address string) (net.IP, error)
+}
+
 // Dialer wraps net.Dialer and provides k6 specific functionality -
 // tracing, blacklists and DNS cache and aliases.
 type Dialer struct {
 	net.Dialer
 
-	Resolver  *dnscache.Resolver
+	Resolver  dnsResolver
 	Blacklist []*lib.IPNet
-	Hosts     map[string]net.IP
+	Hosts     map[string]*lib.HostAddress
 
 	BytesRead    int64
 	BytesWritten int64
 }
 
 // NewDialer constructs a new Dialer and initializes its cache.
 func NewDialer(dialer net.Dialer) *Dialer {
+	return newDialerWithResolver(dialer, dnscache.New(0))
+}
+
+func newDialerWithResolver(dialer net.Dialer, resolver dnsResolver) *Dialer {
 	return &Dialer{
 		Dialer:   dialer,
-		Resolver: dnscache.New(0),
+		Resolver: resolver,
 	}
 }
 
@@ -68,29 +79,11 @@ func (b BlackListedIPError) Error() string {
 
 // DialContext wraps the net.Dialer.DialContext and handles the k6 specifics
 func (d *Dialer) DialContext(ctx context.Context, proto, addr string) (net.Conn, error) {
-	delimiter := strings.LastIndex(addr, ":")
-	host := addr[:delimiter]
-
-	// lookup for domain defined in Hosts option before trying to resolve DNS.
-	ip, ok := d.Hosts[host]
-	if !ok {
-		var err error
-		ip, err = d.Resolver.FetchOne(host)
-		if err != nil {
-			return nil, err
-		}
-	}
-
-	for _, ipnet := range d.Blacklist {
-		if (*net.IPNet)(ipnet).Contains(ip) {
-			return nil, BlackListedIPError{ip: ip, net: ipnet}
-		}
-	}
-	ipStr := ip.String()
-	if strings.ContainsRune(ipStr, ':') {
-		ipStr = "[" + ipStr + "]"
+	dialAddr, err := d.getDialAddr(addr)
+	if err != nil {
+		return nil, err
 	}
-	conn, err := d.Dialer.DialContext(ctx, proto, ipStr+":"+addr[delimiter+1:])
+	conn, err := d.Dialer.DialContext(ctx, proto, dialAddr)
 	if err != nil {
 		return nil, err
 	}
@@ -149,6 +142,77 @@ func (d *Dialer) GetTrail(
 	}
 }
 
+func (d *Dialer) getDialAddr(addr string) (string, error) {
+	remote, err := d.findRemote(addr)
+	if err != nil {
+		return "", err
+	}
+
+	for _, ipnet := range d.Blacklist {
+		if ipnet.Contains(remote.IP) {
+			return "", BlackListedIPError{ip: remote.IP, net: ipnet}
+		}
+	}
+
+	return remote.String(), nil
+}
+
+func (d *Dialer) findRemote(addr string) (*lib.HostAddress, error) {
+	host, port, err := net.SplitHostPort(addr)
+	if err != nil {
+		return nil, err
+	}
+
+	remote, err := d.getConfiguredHost(addr, host, port)
+	if err != nil || remote != nil {
+		return remote, err
+	}
+
+	ip := net.ParseIP(host)
+	if ip != nil {
+		return lib.NewHostAddress(ip, port)
+	}
+
+	return d.fetchRemoteFromResolver(host, port)
+}
+
+func (d *Dialer) fetchRemoteFromResolver(host, port string) (*lib.HostAddress, error) {
+	ip, err := d.Resolver.FetchOne(host)
+	if err != nil {
+		return nil, err
+	}
+
+	if ip == nil {
+		return nil, errors.Errorf("lookup %s: no such host", host)
+	}
+
+	return lib.NewHostAddress(ip, port)
+}
+
+func (d *Dialer) getConfiguredHost(addr, host, port string) (*lib.HostAddress, error) {
+	if remote, ok := d.Hosts[addr]; ok {
+		return remote, nil
+	}
+
+	if remote, ok := d.Hosts[host]; ok {
+		if remote.Port != 0 || port == "" {
+			return remote, nil
+		}
+
+		newPort, err := strconv.Atoi(port)
+		if err != nil {
+			return nil, err
+		}
+
+		newRemote := *remote
+		newRemote.Port = newPort
+
+		return &newRemote, nil
+	}
+
+	return nil, nil
+}
+
 // NetTrail contains information about the exchanged data size and length of a
 // series of connections from a particular netext.Dialer
 type NetTrail struct {

lib/netext/dialer_test.go

@@ -0,0 +1,103 @@
+/*
+ *
+ * k6 - a next-generation load testing tool
+ * Copyright (C) 2020 Load Impact
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as
+ * published by the Free Software Foundation, either version 3 of the
+ * License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ *
+ */
+
+package netext
+
+import (
+	"net"
+	"testing"
+
+	"github.com/loadimpact/k6/lib"
+	"github.com/stretchr/testify/require"
+)
+
+type testResolver struct {
+	hosts map[string]net.IP
+}
+
+func (r testResolver) FetchOne(host string) (net.IP, error) { return r.hosts[host], nil }
+
+func TestDialerAddr(t *testing.T) {
+	dialer := newDialerWithResolver(net.Dialer{}, newResolver())
+	dialer.Hosts = map[string]*lib.HostAddress{
+		"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")},
+	}
+
+	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 {
+		address, expAddress, expErr string
+	}{
+		// IPv4
+		{"example-resolver.com:80", "1.2.3.4:80", ""},
+		{"example.com:80", "3.4.5.6:80", ""},
+		{"example.com:443", "3.4.5.6:8443", ""},
+		{"example.com:8080", "3.4.5.6:9090", ""},
+		{"1.2.3.4:80", "1.2.3.4:80", ""},
+		{"1.2.3.4", "", "address 1.2.3.4: missing port in address"},
+		{"example-deny-resolver.com:80", "", "IP (8.9.10.11) is in a blacklisted range (8.9.10.0/24)"},
+		{"example-deny-host.com:80", "", "IP (8.9.10.11) is in a blacklisted range (8.9.10.0/24)"},
+		{"no-such-host.com:80", "", "lookup no-such-host.com: no such host"},
+
+		// IPv6
+		{"example-ipv6.com:443", "[2001:db8::68]:8443", ""},
+		{"[2001:db8:aaaa:1::100]:443", "[2001:db8:aaaa:1::100]:443", ""},
+		{"[::1.2.3.4]", "", "address [::1.2.3.4]: missing port in address"},
+		{"example-ipv6-deny-resolver.com:80", "", "IP (::1) is in a blacklisted range (::/24)"},
+		{"example-ipv6-deny-host.com:80", "", "IP (::1) is in a blacklisted range (::/24)"},
+	}
+
+	for _, tc := range testCases {
+		tc := tc
+
+		t.Run(tc.address, func(t *testing.T) {
+			addr, err := dialer.getDialAddr(tc.address)
+
+			if tc.expErr != "" {
+				require.EqualError(t, err, tc.expErr)
+			} else {
+				require.NoError(t, err)
+				require.Equal(t, tc.expAddress, addr)
+			}
+		})
+	}
+}
+
+func newResolver() testResolver {
+	return testResolver{
+		hosts: map[string]net.IP{
+			"example-resolver.com":           net.ParseIP("1.2.3.4"),
+			"example-deny-resolver.com":      net.ParseIP("8.9.10.11"),
+			"example-ipv6-deny-resolver.com": net.ParseIP("::1"),
+		},
+	}
+}