Add multi source IP support

cmd/options.go

@@ -93,6 +93,7 @@ func optionFlagSet() *pflag.FlagSet {
 	flags.StringSlice("tag", nil, "add a `tag` to be applied to all samples, as `[name]=[value]`")
 	flags.String("console-output", "", "redirects the console logging to the provided output file")
 	flags.Bool("discard-response-bodies", false, "Read but don't process or save HTTP response bodies")
+	flags.String("local-ips", "", "Client IP Ranges and/or CIDRs from which each VU will be making requests")
 	flags.String("dns", types.DefaultDNSConfig().String(), "DNS resolver configuration. Possible ttl values are: 'inf' "+
 		"for a persistent cache, '0' to disable the cache,\nor a positive duration, e.g. '1s', '1m', etc. "+
 		"Milliseconds are assumed if no unit is provided.\n"+
@@ -205,6 +206,17 @@ func getOptions(flags *pflag.FlagSet) (lib.Options, error) {
 		}
 	}
 
+	localIpsString, err := flags.GetString("local-ips")
+	if err != nil {
+		return opts, err
+	}
+	if flags.Changed("local-ips") {
+		err = opts.LocalIPs.UnmarshalText([]byte(localIpsString))
+		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

@@ -170,6 +170,10 @@ func (r *Runner) newVU(id int64, samplesOut chan<- stats.SampleContainer) (*VU,
 		BlockedHostnames: r.Bundle.Options.BlockedHostnames.Trie,
 		Hosts:            r.Bundle.Options.Hosts,
 	}
+	if r.Bundle.Options.LocalIPs.Valid {
+		dialer.Dialer.LocalAddr = &net.TCPAddr{IP: r.Bundle.Options.LocalIPs.Pool.GetIP(uint64(id - 1))}
+	}
+
 	tlsConfig := &tls.Config{
 		InsecureSkipVerify: r.Bundle.Options.InsecureSkipTLSVerify.Bool,
 		CipherSuites:       cipherSuites,
@@ -307,7 +311,6 @@ func (r *Runner) IsExecutable(name string) bool {
 
 func (r *Runner) SetOptions(opts lib.Options) error {
 	r.Bundle.Options = opts
-
 	r.RPSLimit = nil
 	if rps := opts.RPS; rps.Valid {
 		r.RPSLimit = rate.NewLimiter(rate.Limit(rps.Int64), 1)

lib/options.go

@@ -388,6 +388,9 @@ type Options struct {
 
 	// Redirect console logging to a file
 	ConsoleOutput null.String `json:"-" envconfig:"K6_CONSOLE_OUTPUT"`
+
+	// Specify client IP ranges and/or CIDR from which VUs will make requests
+	LocalIPs types.NullIPPool `json:"-" envconfig:"K6_LOCAL_IPS"`
 }
 
 // Returns the result of overwriting any fields with any that are set on the argument.
@@ -542,6 +545,9 @@ func (o Options) Apply(opts Options) Options {
 	if opts.ConsoleOutput.Valid {
 		o.ConsoleOutput = opts.ConsoleOutput
 	}
+	if opts.LocalIPs.Valid {
+		o.LocalIPs = opts.LocalIPs
+	}
 	if opts.DNS.TTL.Valid {
 		o.DNS.TTL = opts.DNS.TTL
 	}

lib/options_test.go

@@ -407,9 +407,21 @@ func TestOptions(t *testing.T) {
 		assert.True(t, opts.DiscardResponseBodies.Valid)
 		assert.True(t, opts.DiscardResponseBodies.Bool)
 	})
+	t.Run("ClientIPRanges", func(t *testing.T) {
+		clientIPRanges, err := types.NewIPPool("129.112.232.12,123.12.0.0/32")
+		require.NoError(t, err)
+		opts := Options{}.Apply(Options{LocalIPs: types.NullIPPool{Pool: clientIPRanges, Valid: true}})
+		assert.NotNil(t, opts.LocalIPs)
+	})
 }
 
 func TestOptionsEnv(t *testing.T) {
+	mustIPPool := func(s string) *types.IPPool {
+		p, err := types.NewIPPool(s)
+		require.NoError(t, err)
+		return p
+	}
+
 	testdata := map[struct{ Name, Key string }]map[string]interface{}{
 		{"Paused", "K6_PAUSED"}: {
 			"":      null.Bool{},
@@ -469,6 +481,11 @@ func TestOptionsEnv(t *testing.T) {
 			"":    null.String{},
 			"Hi!": null.StringFrom("Hi!"),
 		},
+		{"LocalIPs", "K6_LOCAL_IPS"}: {
+			"":                 types.NullIPPool{},
+			"192.168.220.2":    types.NullIPPool{Pool: mustIPPool("192.168.220.2"), Valid: true},
+			"192.168.220.2/24": types.NullIPPool{Pool: mustIPPool("192.168.220.0/24"), Valid: true},
+		},
 		{"Throw", "K6_THROW"}: {
 			"":      null.Bool{},
 			"true":  null.BoolFrom(true),

lib/types/ipblock.go

@@ -0,0 +1,201 @@
+/*
+ *
+ * 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 types
+
+import (
+	"errors"
+	"fmt"
+	"math/big"
+	"net"
+	"strings"
+)
+
+// ipBlock represents a continuous segment of IP addresses
+type ipBlock struct {
+	firstIP, count *big.Int
+	ipv6           bool
+}
+
+// ipPoolBlock is almost the same as ipBlock but is put in a IPPool and knows it's start id indest
+// of it's size/count
+type ipPoolBlock struct {
+	firstIP, startIndex *big.Int
+}
+
+// IPPool represent a slice of IPBlocks
+type IPPool struct {
+	list  []ipPoolBlock
+	count *big.Int
+}
+
+func getIPBlock(s string) (*ipBlock, error) {
+	switch {
+	case strings.Contains(s, "-"):
+		return ipBlockFromRange(s)
+	case strings.Contains(s, "/"):
+		return ipBlockFromCIDR(s)
+	default:
+		if net.ParseIP(s) == nil {
+			return nil, fmt.Errorf("%s is not a valid IP, IP range or CIDR", s)
+		}
+		return ipBlockFromRange(s + "-" + s)
+	}
+}
+
+func ipBlockFromRange(s string) (*ipBlock, error) {
+	ss := strings.SplitN(s, "-", 2)
+	ip0Str, ip1Str := strings.TrimSpace(ss[0]), strings.TrimSpace(ss[1])
+	ip0, ip1 := net.ParseIP(ip0Str), net.ParseIP(ip1Str)
+	if ip0 == nil || ip1 == nil {
+		return nil, errors.New("wrong IP range format: " + s)
+	}
+	if (ip0.To4() == nil) != (ip1.To4() == nil) { // XOR
+		return nil, errors.New("mixed IP range format: " + s)
+	}
+	block := ipBlockFromTwoIPs(ip0, ip1)
+
+	if block.count.Sign() <= 0 {
+		return nil, errors.New("negative IP range: " + s)
+	}
+	return block, nil
+}
+
+func ipBlockFromTwoIPs(ip0, ip1 net.IP) *ipBlock {
+	// This code doesn't do any checks on the validity of the arguments, that should be
+	// done before and/or after it is called
+	var block ipBlock
+	block.firstIP = new(big.Int)
+	block.count = new(big.Int)
+	block.ipv6 = ip0.To4() == nil
+	if block.ipv6 {
+		block.firstIP.SetBytes(ip0.To16())
+		block.count.SetBytes(ip1.To16())
+	} else {
+		block.firstIP.SetBytes(ip0.To4())
+		block.count.SetBytes(ip1.To4())
+	}
+	block.count.Sub(block.count, block.firstIP)
+	block.count.Add(block.count, big.NewInt(1))
+
+	return &block
+}
+
+func ipBlockFromCIDR(s string) (*ipBlock, error) {
+	_, pnet, err := net.ParseCIDR(s)
+	if err != nil {
+		return nil, fmt.Errorf("parseCIDR() failed parsing %s: %w", s, err)
+	}
+	ip0 := pnet.IP
+	// TODO: this is just to copy it, it will probably be better to copy the bytes ...
+	ip1 := net.ParseIP(ip0.String())
+	if ip1.To4() == nil {
+		ip1 = ip1.To16()
+	} else {
+		ip1 = ip1.To4()
+	}
+	for i := range ip1 {
+		ip1[i] |= (255 ^ pnet.Mask[i])
+	}
+	block := ipBlockFromTwoIPs(ip0, ip1)
+	// in the case of ipv4 if the network is bigger than 31 the first and last IP are reserved so we
+	// need to reduce the addresses by 2 and increment the first ip
+	if !block.ipv6 && big.NewInt(2).Cmp(block.count) < 0 {
+		block.count.Sub(block.count, big.NewInt(2))
+		block.firstIP.Add(block.firstIP, big.NewInt(1))
+	}
+	return block, nil
+}
+
+func (b ipPoolBlock) getIP(index *big.Int) net.IP {
+	// TODO implement walking ipv6 networks first
+	// that will probably require more math ... including knowing which is the next network and ...
+	// thinking about it - it looks like it's going to be kind of hard or badly defined
+	i := new(big.Int)
+	i.Add(b.firstIP, index)
+	// TODO use big.Int.FillBytes when golang 1.14 is no longer supported
+	return net.IP(i.Bytes())
+}
+
+// NewIPPool returns an IPPool slice from the provided string representation that should be comma
+// separated list of IPs, IP ranges(ip1-ip2) and CIDRs
+func NewIPPool(ranges string) (*IPPool, error) {
+	ss := strings.Split(strings.TrimSpace(ranges), ",")
+	pool := &IPPool{}
+	pool.list = make([]ipPoolBlock, len(ss))
+	pool.count = new(big.Int)
+	for i, bs := range ss {
+		r, err := getIPBlock(bs)
+		if err != nil {
+			return nil, err
+		}
+
+		pool.list[i] = ipPoolBlock{
+			firstIP:    r.firstIP,
+			startIndex: new(big.Int).Set(pool.count), // this is how many there are until now
+		}
+		pool.count.Add(pool.count, r.count)
+	}
+
+	// The list gets reversed here as later it is searched based on when the index we are looking is
+	// bigger than startIndex but it will be true always for the first block which is with
+	// startIndex 0. This can also be fixed by iterating in reverse but this seems better
+	for i := 0; i < len(pool.list)/2; i++ {
+		pool.list[i], pool.list[len(pool.list)/2-i] = pool.list[len(pool.list)/2-i], pool.list[i]
+	}
+	return pool, nil
+}
+
+// GetIP return an IP from a pool of IPBlock slice
+func (pool *IPPool) GetIP(index uint64) net.IP {
+	return pool.GetIPBig(new(big.Int).SetUint64(index))
+}
+
+// GetIPBig returns an IP from the pool with the provided index that is big.Int
+func (pool *IPPool) GetIPBig(index *big.Int) net.IP {
+	index = new(big.Int).Rem(index, pool.count)
+	for _, b := range pool.list {
+		if index.Cmp(b.startIndex) >= 0 {
+			return b.getIP(index.Sub(index, b.startIndex))
+		}
+	}
+	return nil
+}
+
+// NullIPPool is a nullable IPPool
+type NullIPPool struct {
+	Pool  *IPPool
+	Valid bool
+}
+
+// UnmarshalText converts text data to a valid NullIPPool
+func (n *NullIPPool) UnmarshalText(data []byte) error {
+	if len(data) == 0 {
+		*n = NullIPPool{}
+		return nil
+	}
+	var err error
+	n.Pool, err = NewIPPool(string(data))
+	if err != nil {
+		return err
+	}
+	n.Valid = true
+	return nil
+}