Add blocking by hostname patterns (--block-hostname)

cmd/options.go

@@ -63,6 +63,7 @@ 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 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().
@@ -151,6 +152,19 @@ 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 len(blockedHostnameStrings) > 0 {
+		opts.BlockedHostnames = &lib.HostnameTrie{}
+	}
+	for _, s := range blockedHostnameStrings {
+		if insertErr := opts.BlockedHostnames.Insert(s); insertErr != nil {
+			return opts, errors.Wrap(insertErr, "block-hostname")
+		}
+	}
+
 	if flags.Changed("summary-trend-stats") {
 		trendStats, errSts := flags.GetStringSlice("summary-trend-stats")
 		if errSts != nil {

js/runner.go

@@ -154,10 +154,11 @@ func (r *Runner) newVU(samplesOut chan<- stats.SampleContainer) (*VU, error) {
 	}
 
 	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,
+		Hosts:            r.Bundle.Options.Hosts,
 	}
 	tlsConfig := &tls.Config{
 		InsecureSkipVerify: r.Bundle.Options.InsecureSkipTLSVerify.Bool,

js/runner_test.go

@@ -818,6 +818,81 @@ func TestVUIntegrationBlacklistScript(t *testing.T) {
 	}
 }
 
+func TestVUIntegrationBlockHostnamesOption(t *testing.T) {
+	r1, err := getSimpleRunner("/script.js", `
+					import http from "k6/http";
+					export default function() { http.get("https://k6.io/"); }
+				`)
+	if !assert.NoError(t, err) {
+		return
+	}
+
+	hostnames := lib.HostnameTrie{}
+	if insertErr := hostnames.Insert("*.io"); !assert.NoError(t, insertErr) {
+		return
+	}
+	require.NoError(t, r1.SetOptions(lib.Options{
+		Throw:            null.BoolFrom(true),
+		BlockedHostnames: &hostnames,
+	}))
+
+	r2, err := NewFromArchive(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) {
+			vu, err := r.NewVU(make(chan stats.SampleContainer, 100))
+			if !assert.NoError(t, err) {
+				return
+			}
+			err = vu.RunOnce(context.Background())
+			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("/script.js", `
+					import http from "k6/http";
+
+					export let options = {
+						throw: true,
+						blockHostnames: ["*.io"],
+					};
+
+					export default function() { http.get("https://k6.io/"); }
+				`)
+	if !assert.NoError(t, err) {
+		return
+	}
+
+	r2, err := NewFromArchive(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) {
+			vu, err := r.NewVU(make(chan stats.SampleContainer, 100))
+			if !assert.NoError(t, err) {
+				return
+			}
+			err = vu.RunOnce(context.Background())
+			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

@@ -40,9 +40,10 @@ import (
 type Dialer struct {
 	net.Dialer
 
-	Resolver  *dnscache.Resolver
-	Blacklist []*lib.IPNet
-	Hosts     map[string]net.IP
+	Resolver         *dnscache.Resolver
+	Blacklist        []*lib.IPNet
+	BlockedHostnames *lib.HostnameTrie
+	Hosts            map[string]net.IP
 
 	BytesRead    int64
 	BytesWritten int64
@@ -66,11 +67,27 @@ 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) {
 	delimiter := strings.LastIndex(addr, ":")
 	host := addr[:delimiter]
 
+	if d.BlockedHostnames != nil {
+		if blocked, match := d.BlockedHostnames.Contains(host); blocked {
+			return nil, BlockedHostError{hostname: host, match: match}
+		}
+	}
+
 	// lookup for domain defined in Hosts option before trying to resolve DNS.
 	ip, ok := d.Hosts[host]
 	if !ok {

lib/netext/httpext/error_codes.go

@@ -46,9 +46,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
@@ -90,6 +91,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"
@@ -118,6 +120,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

@@ -26,6 +26,8 @@ import (
 	"fmt"
 	"net"
 	"reflect"
+	"regexp"
+	"strings"
 
 	"github.com/loadimpact/k6/lib/scheduler"
 	"github.com/loadimpact/k6/lib/types"
@@ -187,6 +189,116 @@ func ParseCIDR(s string) (*IPNet, error) {
 	return &parsedIPNet, nil
 }
 
+// HostnameTrie is a tree-structured list of hostname matches with support
+// for wildcards exclusively at the start of the pattern. Items may only
+// be inserted and searched. Internationalized hostnames are valid.
+type HostnameTrie struct {
+	children []*HostnameTrie
+	r        rune
+	terminal bool // end of a valid match
+}
+
+// Regex description of hostname pattern to enforce blocks by. Global var
+// to avoid compilation penalty at runtime.
+// Matches against strings composed entirely of letters, numbers, or '.'s
+// with an optional wildcard at the start.
+var legalHostnamePattern *regexp.Regexp = regexp.MustCompile("^\\*?(\\pL|[0-9\\.])*")
+
+func legalHostname(s string) error {
+	if len(legalHostnamePattern.FindString(s)) != len(s) {
+		return errors.Errorf("invalid hostname pattern %s", s)
+	}
+	return nil
+}
+
+// UnmarshalJSON forms a HostnameTrie from the provided hostname pattern
+// list.
+func (t *HostnameTrie) UnmarshalJSON(data []byte) error {
+	m := make([]string, 0)
+	if err := json.Unmarshal(data, &m); err != nil {
+		return err
+	}
+	for _, h := range m {
+		if insertErr := t.Insert(h); insertErr != nil {
+			return insertErr
+		}
+	}
+	return nil
+}
+
+// UnmarshalText forms a HostnameTrie from a comma-delimited list
+// of hostname patterns.
+func (t *HostnameTrie) UnmarshalText(b []byte) error {
+	for _, s := range strings.Split(string(b), ",") {
+		if err := t.Insert(s); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+// Insert a hostname pattern into the given HostnameTrie. Returns an error
+// if hostname pattern is illegal.
+func (t *HostnameTrie) Insert(s string) error {
+	if len(s) == 0 {
+		return nil
+	}
+
+	if err := legalHostname(s); err != nil {
+		return err
+	}
+
+	rStr := []rune(s) // need to iterate by runes for intl' names
+	last := len(rStr) - 1
+	for _, c := range t.children {
+		if c.r == rStr[last] {
+			return c.Insert(string(rStr[:last]))
+		}
+	}
+
+	n := &HostnameTrie{nil, rStr[last], len(rStr) == 1}
+	t.children = append(t.children, n)
+	return n.Insert(string(rStr[:last]))
+}
+
+// Contains returns whether s matches a pattern in the HostnameTrie
+// along with the matching pattern, if one was found.
+func (t *HostnameTrie) Contains(s string) (bool, string) {
+	for _, c := range t.children {
+		if b, m := c.childContains(s, ""); b {
+			return b, m
+		}
+	}
+	return false, ""
+}
+
+// recursively traverse HostnameTrie children searching for a match.
+func (t *HostnameTrie) childContains(s string, match string) (bool, string) {
+	if len(s) == 0 {
+		return false, ""
+	}
+
+	rStr := []rune(s)
+	last := len(rStr) - 1
+
+	switch {
+	case t.r == '*': // wildcard encounters validate the string
+		return true, string(t.r) + match
+	case t.r != rStr[last]:
+		return false, ""
+	case len(s) == 1:
+		return t.terminal, string(t.r) + match
+	default:
+		for _, c := range t.children {
+			if b, m := c.childContains(string(rStr[:last]), string(t.r)+match); b {
+				return b, m
+			}
+		}
+	}
+
+	return false, ""
+}
+
 type Options struct {
 	// Should the test start in a paused state?
 	Paused null.Bool `json:"paused" envconfig:"K6_PAUSED"`
@@ -242,6 +354,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 *HostnameTrie `json:"blockHostnames" envconfig:"K6_BLOCK_HOSTNAMES"`
+
 	// Hosts overrides dns entries for given hosts
 	Hosts map[string]net.IP `json:"hosts" envconfig:"K6_HOSTS"`
 
@@ -389,6 +504,9 @@ func (o Options) Apply(opts Options) Options {
 	if opts.BlacklistIPs != nil {
 		o.BlacklistIPs = opts.BlacklistIPs
 	}
+	if opts.BlockedHostnames != nil {
+		o.BlockedHostnames = opts.BlockedHostnames
+	}
 	if opts.Hosts != nil {
 		o.Hosts = opts.Hosts
 	}

lib/options_test.go

@@ -318,6 +318,18 @@ 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) {
+		hostnames := HostnameTrie{}
+		err := hostnames.Insert("*")
+		assert.Nil(t, err)
+		opts := Options{}.Apply(Options{
+			BlockedHostnames: &hostnames,
+		})
+		assert.NotNil(t, opts.BlockedHostnames)
+		assert.NotEmpty(t, opts.BlockedHostnames)
+		matches, _ := opts.BlockedHostnames.Contains("loadimpact.com")
+		assert.True(t, matches)
+	})
 
 	t.Run("Hosts", func(t *testing.T) {
 		opts := Options{}.Apply(Options{Hosts: map[string]net.IP{