Merge pull request #1666 from loadimpact/feature/block-hostname

Feature/block hostname

cmd/options.go

@@ -72,6 +72,8 @@ func optionFlagSet() *pflag.FlagSet {
 	flags.Duration("min-iteration-duration", 0, "minimum amount of time k6 will take executing a single iteration")
 	flags.BoolP("throw", "w", false, "throw warnings (like failed http requests) as errors")
 	flags.StringSlice("blacklist-ip", nil, "blacklist an `ip range` from being called")
+	flags.StringSlice("block-hostname", nil, "block a case-insensitive hostname `pattern`,"+
+		" with optional leading wildcard, from being called")
 
 	// The comment about system-tags also applies for summary-trend-stats. The default values
 	// are set in applyDefault().
@@ -187,6 +189,17 @@ func getOptions(flags *pflag.FlagSet) (lib.Options, error) {
 		opts.BlacklistIPs = append(opts.BlacklistIPs, net)
 	}
 
+	blockedHostnameStrings, err := flags.GetStringSlice("block-hostname")
+	if err != nil {
+		return opts, err
+	}
+	if flags.Changed("block-hostname") {
+		opts.BlockedHostnames, err = types.NewNullHostnameTrie(blockedHostnameStrings)
+		if err != nil {
+			return opts, err
+		}
+	}
+
 	if flags.Changed("summary-trend-stats") {
 		trendStats, errSts := flags.GetStringSlice("summary-trend-stats")
 		if errSts != nil {

js/runner.go

@@ -159,10 +159,11 @@ func (r *Runner) newVU(id int64, samplesOut chan<- stats.SampleContainer) (*VU,
 	}
 
 	dialer := &netext.Dialer{
-		Dialer:    r.BaseDialer,
-		Resolver:  r.Resolver,
-		Blacklist: r.Bundle.Options.BlacklistIPs,
-		Hosts:     r.Bundle.Options.Hosts,
+		Dialer:           r.BaseDialer,
+		Resolver:         r.Resolver,
+		Blacklist:        r.Bundle.Options.BlacklistIPs,
+		BlockedHostnames: r.Bundle.Options.BlockedHostnames.Trie,
+		Hosts:            r.Bundle.Options.Hosts,
 	}
 	tlsConfig := &tls.Config{
 		InsecureSkipVerify: r.Bundle.Options.InsecureSkipTLSVerify.Bool,

js/runner_test.go

@@ -885,6 +885,77 @@ func TestVUIntegrationBlacklistScript(t *testing.T) {
 	}
 }
 
+func TestVUIntegrationBlockHostnamesOption(t *testing.T) {
+	r1, err := getSimpleRunner(t, "/script.js", `
+					var http = require("k6/http");
+					exports.default = function() { http.get("https://k6.io/"); }
+				`)
+	require.NoError(t, err)
+
+	hostnames, err := types.NewNullHostnameTrie([]string{"*.io"})
+	require.NoError(t, err)
+	require.NoError(t, r1.SetOptions(lib.Options{
+		Throw:            null.BoolFrom(true),
+		BlockedHostnames: hostnames,
+	}))
+
+	r2, err := NewFromArchive(testutils.NewLogger(t), r1.MakeArchive(), lib.RuntimeOptions{})
+	if !assert.NoError(t, err) {
+		return
+	}
+
+	runners := map[string]*Runner{"Source": r1, "Archive": r2}
+
+	for name, r := range runners {
+		r := r
+		t.Run(name, func(t *testing.T) {
+			initVu, err := r.NewVU(1, make(chan stats.SampleContainer, 100))
+			require.NoError(t, err)
+			vu := initVu.Activate(&lib.VUActivationParams{RunContext: context.Background()})
+			err = vu.RunOnce()
+			require.Error(t, err)
+			assert.Contains(t, err.Error(), "hostname (k6.io) is in a blocked pattern (*.io)")
+		})
+	}
+}
+
+func TestVUIntegrationBlockHostnamesScript(t *testing.T) {
+	r1, err := getSimpleRunner(t, "/script.js", `
+					var http = require("k6/http");
+
+					exports.options = {
+						throw: true,
+						blockHostnames: ["*.io"],
+					};
+
+					exports.default = function() { http.get("https://k6.io/"); }
+				`)
+	if !assert.NoError(t, err) {
+		return
+	}
+
+	r2, err := NewFromArchive(testutils.NewLogger(t), r1.MakeArchive(), lib.RuntimeOptions{})
+	if !assert.NoError(t, err) {
+		return
+	}
+
+	runners := map[string]*Runner{"Source": r1, "Archive": r2}
+
+	for name, r := range runners {
+		r := r
+		t.Run(name, func(t *testing.T) {
+			initVu, err := r.NewVU(0, make(chan stats.SampleContainer, 100))
+			if !assert.NoError(t, err) {
+				return
+			}
+			vu := initVu.Activate(&lib.VUActivationParams{RunContext: context.Background()})
+			err = vu.RunOnce()
+			require.Error(t, err)
+			assert.Contains(t, err.Error(), "hostname (k6.io) is in a blocked pattern (*.io)")
+		})
+	}
+}
+
 func TestVUIntegrationHosts(t *testing.T) {
 	tb := httpmultibin.NewHTTPMultiBin(t)
 	defer tb.Cleanup()

lib/netext/dialer.go

@@ -33,6 +33,7 @@ import (
 
 	"github.com/loadimpact/k6/lib"
 	"github.com/loadimpact/k6/lib/metrics"
+	"github.com/loadimpact/k6/lib/types"
 	"github.com/loadimpact/k6/stats"
 )
 
@@ -47,9 +48,10 @@ type dnsResolver interface {
 type Dialer struct {
 	net.Dialer
 
-	Resolver  dnsResolver
-	Blacklist []*lib.IPNet
-	Hosts     map[string]*lib.HostAddress
+	Resolver         dnsResolver
+	Blacklist        []*lib.IPNet
+	BlockedHostnames *types.HostnameTrie
+	Hosts            map[string]*lib.HostAddress
 
 	BytesRead    int64
 	BytesWritten int64
@@ -77,6 +79,16 @@ func (b BlackListedIPError) Error() string {
 	return fmt.Sprintf("IP (%s) is in a blacklisted range (%s)", b.ip, b.net)
 }
 
+// BlockedHostError is returned when a given hostname is blocked
+type BlockedHostError struct {
+	hostname string
+	match    string
+}
+
+func (b BlockedHostError) Error() string {
+	return fmt.Sprintf("hostname (%s) is in a blocked pattern (%s)", b.hostname, b.match)
+}
+
 // DialContext wraps the net.Dialer.DialContext and handles the k6 specifics
 func (d *Dialer) DialContext(ctx context.Context, proto, addr string) (net.Conn, error) {
 	dialAddr, err := d.getDialAddr(addr)
@@ -163,12 +175,18 @@ func (d *Dialer) findRemote(addr string) (*lib.HostAddress, error) {
 		return nil, err
 	}
 
+	ip := net.ParseIP(host)
+	if d.BlockedHostnames != nil && ip == nil {
+		if match, blocked := d.BlockedHostnames.Contains(host); blocked {
+			return nil, BlockedHostError{hostname: host, match: match}
+		}
+	}
+
 	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)
 	}

lib/netext/dialer_test.go

@@ -25,6 +25,7 @@ import (
 	"testing"
 
 	"github.com/loadimpact/k6/lib"
+	"github.com/loadimpact/k6/lib/types"
 	"github.com/stretchr/testify/require"
 )
 
@@ -74,6 +75,42 @@ func TestDialerAddr(t *testing.T) {
 		{"[::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)"},
+		{"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 TestDialerAddrBlockHostnamesStar(t *testing.T) {
+	dialer := newDialerWithResolver(net.Dialer{}, newResolver())
+	dialer.Hosts = map[string]*lib.HostAddress{
+		"example.com": {IP: net.ParseIP("3.4.5.6")},
+	}
+
+	blocked, err := types.NewHostnameTrie([]string{"*"})
+	require.NoError(t, err)
+	dialer.BlockedHostnames = blocked
+	testCases := []struct {
+		address, expAddress, expErr string
+	}{
+		// IPv4
+		{"example.com:80", "", "hostname (example.com) is in a blocked pattern (*)"},
+		{"example.com:443", "", "hostname (example.com) is in a blocked pattern (*)"},
+		{"not.com:30", "", "hostname (not.com) is in a blocked pattern (*)"},
+		{"1.2.3.4:80", "1.2.3.4:80", ""},
 	}
 
 	for _, tc := range testCases {

lib/netext/httpext/error_codes.go

@@ -47,9 +47,10 @@ const (
 	defaultErrorCode          errCode = 1000
 	defaultNetNonTCPErrorCode errCode = 1010
 	// DNS errors
-	defaultDNSErrorCode    errCode = 1100
-	dnsNoSuchHostErrorCode errCode = 1101
-	blackListedIPErrorCode errCode = 1110
+	defaultDNSErrorCode      errCode = 1100
+	dnsNoSuchHostErrorCode   errCode = 1101
+	blackListedIPErrorCode   errCode = 1110
+	blockedHostnameErrorCode errCode = 1111
 	// tcp errors
 	defaultTCPErrorCode      errCode = 1200
 	tcpBrokenPipeErrorCode   errCode = 1201
@@ -91,6 +92,7 @@ const (
 	netUnknownErrnoErrorCodeMsg = "%s: unknown errno `%d` on %s with message `%s`"
 	dnsNoSuchHostErrorCodeMsg   = "lookup: no such host"
 	blackListedIPErrorCodeMsg   = "ip is blacklisted"
+	blockedHostnameErrorMsg     = "hostname is blocked"
 	http2GoAwayErrorCodeMsg     = "http2: received GoAway with http2 ErrCode %s"
 	http2StreamErrorCodeMsg     = "http2: stream error with http2 ErrCode %s"
 	http2ConnectionErrorCodeMsg = "http2: connection error with http2 ErrCode %s"
@@ -119,6 +121,8 @@ func errorCodeForError(err error) (errCode, string) {
 		}
 	case netext.BlackListedIPError:
 		return blackListedIPErrorCode, blackListedIPErrorCodeMsg
+	case netext.BlockedHostError:
+		return blockedHostnameErrorCode, blockedHostnameErrorMsg
 	case *http2.GoAwayError:
 		return unknownHTTP2GoAwayErrorCode + http2ErrCodeOffset(e.ErrCode),
 			fmt.Sprintf(http2GoAwayErrorCodeMsg, e.ErrCode)

lib/options.go

@@ -339,6 +339,9 @@ type Options struct {
 	// Blacklist IP ranges that tests may not contact. Mainly useful in hosted setups.
 	BlacklistIPs []*IPNet `json:"blacklistIPs" envconfig:"K6_BLACKLIST_IPS"`
 
+	// Block hostname patterns that tests may not contact.
+	BlockedHostnames types.NullHostnameTrie `json:"blockHostnames" envconfig:"K6_BLOCK_HOSTNAMES"`
+
 	// Hosts overrides dns entries for given hosts
 	Hosts map[string]*HostAddress `json:"hosts" envconfig:"K6_HOSTS"`
 
@@ -494,6 +497,9 @@ func (o Options) Apply(opts Options) Options {
 	if opts.BlacklistIPs != nil {
 		o.BlacklistIPs = opts.BlacklistIPs
 	}
+	if opts.BlockedHostnames.Valid {
+		o.BlockedHostnames = opts.BlockedHostnames
+	}
 	if opts.Hosts != nil {
 		o.Hosts = opts.Hosts
 	}

lib/options_test.go

@@ -314,6 +314,13 @@ func TestOptions(t *testing.T) {
 		assert.Equal(t, net.IPv4zero, opts.BlacklistIPs[0].IP)
 		assert.Equal(t, net.CIDRMask(1, 1), opts.BlacklistIPs[0].Mask)
 	})
+	t.Run("BlockedHostnames", func(t *testing.T) {
+		blockedHostnames, err := types.NewNullHostnameTrie([]string{"test.k6.io", "*valid.pattern"})
+		require.NoError(t, err)
+		opts := Options{}.Apply(Options{BlockedHostnames: blockedHostnames})
+		assert.NotNil(t, opts.BlockedHostnames)
+		assert.Equal(t, blockedHostnames, opts.BlockedHostnames)
+	})
 
 	t.Run("Hosts", func(t *testing.T) {
 		host, err := NewHostAddress(net.ParseIP("192.0.2.1"), "80")