diff --git a/cmd/options.go b/cmd/options.go index 0276b4daaab..4dd3dbe6896 100644 --- a/cmd/options.go +++ b/cmd/options.go @@ -93,6 +93,8 @@ 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, "+ + "e.g. '192.168.220.1,192.168.0.10-192.168.0.25', 'fd:1::0/120', etc.") 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 +207,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 { diff --git a/js/runner.go b/js/runner.go index 370ebf17790..78629479ba7 100644 --- a/js/runner.go +++ b/js/runner.go @@ -170,6 +170,14 @@ 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 { + var ipIndex uint64 + if id > 0 { + ipIndex = uint64(id - 1) + } + dialer.Dialer.LocalAddr = &net.TCPAddr{IP: r.Bundle.Options.LocalIPs.Pool.GetIP(ipIndex)} + } + tlsConfig := &tls.Config{ InsecureSkipVerify: r.Bundle.Options.InsecureSkipTLSVerify.Bool, CipherSuites: cipherSuites, @@ -307,7 +315,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) diff --git a/lib/options.go b/lib/options.go index 6dfec6ed4c7..89deb6d2c07 100644 --- a/lib/options.go +++ b/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 } diff --git a/lib/options_test.go b/lib/options_test.go index a7eb55b4147..72d7aacca0c 100644 --- a/lib/options_test.go +++ b/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), diff --git a/lib/types/ipblock.go b/lib/types/ipblock.go new file mode 100644 index 00000000000..0c3f18b8f45 --- /dev/null +++ b/lib/types/ipblock.go @@ -0,0 +1,200 @@ +/* + * + * 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 . + * + */ + +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 similar to ipBlock but instead of knowing its count/size it knows the first index +// from which it starts in an IPPool +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) + ip0, ip1 := net.ParseIP(ss[0]), net.ParseIP(ss[1]) + 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(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)-1-i] = pool.list[len(pool.list)-1-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 +} diff --git a/lib/types/ipblock_test.go b/lib/types/ipblock_test.go new file mode 100644 index 00000000000..dcbddaa02af --- /dev/null +++ b/lib/types/ipblock_test.go @@ -0,0 +1,156 @@ +/* + * + * 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 . + * + */ + +package types + +import ( + "math/big" + "net" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +//nolint:gochecknoglobals +var max64 = new(big.Int).Exp(big.NewInt(2), big.NewInt(64), nil) + +func get128BigInt(hi, lo int64) *big.Int { + h := big.NewInt(hi) + h.Mul(h, max64) + return h.Add(h, big.NewInt(lo)) +} + +func TestIpBlock(t *testing.T) { + testdata := map[string]struct { + count *big.Int + firstIP, lastIP net.IP + }{ + "192.168.0.101": {new(big.Int).SetInt64(1), net.ParseIP("192.168.0.101"), net.ParseIP("192.168.0.101")}, + + "192.168.0.101-192.168.0.200": {new(big.Int).SetInt64(100), net.ParseIP("192.168.0.101"), net.ParseIP("192.168.0.200")}, + "192.168.0.100-192.168.0.200": {new(big.Int).SetInt64(101), net.ParseIP("192.168.0.100"), net.ParseIP("192.168.0.200")}, + "fd00:1:1:0::0-fd00:1:1:ff::3ff": {get128BigInt(255, 1024), net.ParseIP("fd00:1:1:0::0"), net.ParseIP("fd00:1:1:ff::3ff")}, + "fd00:1:1:2::1-fd00:1:1:ff::3ff": {get128BigInt(253, 1023), net.ParseIP("fd00:1:1:2::1"), net.ParseIP("fd00:1:1:ff::3ff")}, + + "192.168.0.0/16": {get128BigInt(0, 65534), net.ParseIP("192.168.0.1"), net.ParseIP("192.168.255.254")}, + "192.168.0.1/16": {get128BigInt(0, 65534), net.ParseIP("192.168.0.1"), net.ParseIP("192.168.255.254")}, + "192.168.0.10/16": {get128BigInt(0, 65534), net.ParseIP("192.168.0.1"), net.ParseIP("192.168.255.254")}, + "192.168.0.10/31": {get128BigInt(0, 2), net.ParseIP("192.168.0.10"), net.ParseIP("192.168.0.11")}, + "192.168.0.10/32": {get128BigInt(0, 1), net.ParseIP("192.168.0.10"), net.ParseIP("192.168.0.10")}, + "fd00::0/120": {get128BigInt(0, 256), net.ParseIP("fd00::0"), net.ParseIP("fd00::ff")}, + "fd00::1/120": {get128BigInt(0, 256), net.ParseIP("fd00::0"), net.ParseIP("fd00::ff")}, + "fd00::3/120": {get128BigInt(0, 256), net.ParseIP("fd00::0"), net.ParseIP("fd00::ff")}, + "fd00::0/112": {get128BigInt(0, 65536), net.ParseIP("fd00::0"), net.ParseIP("fd00::ffff")}, + "fd00::1/112": {get128BigInt(0, 65536), net.ParseIP("fd00::0"), net.ParseIP("fd00::ffff")}, + "fd00::2/112": {get128BigInt(0, 65536), net.ParseIP("fd00::0"), net.ParseIP("fd00::ffff")}, + } + for name, data := range testdata { + name, data := name, data + t.Run(name, func(t *testing.T) { + b, err := getIPBlock(name) + require.NoError(t, err) + assert.Equal(t, data.count, b.count) + pb := ipPoolBlock{firstIP: b.firstIP} + idx := big.NewInt(0) + assert.Equal(t, data.firstIP.To16(), pb.getIP(idx).To16()) + idx.Sub(idx.Add(idx, b.count), big.NewInt(1)) + assert.Equal(t, data.lastIP.To16(), pb.getIP(idx).To16()) + }) + } +} + +func TestIPPool(t *testing.T) { + testdata := map[string]struct { + count *big.Int + queries map[uint64]net.IP + }{ + "192.168.0.101": { + count: new(big.Int).SetInt64(1), + queries: map[uint64]net.IP{0: net.ParseIP("192.168.0.101"), 12: net.ParseIP("192.168.0.101")}, + }, + "192.168.0.101,192.168.0.102": { + count: new(big.Int).SetInt64(2), + queries: map[uint64]net.IP{ + 0: net.ParseIP("192.168.0.101"), + 1: net.ParseIP("192.168.0.102"), + 12: net.ParseIP("192.168.0.101"), + 13: net.ParseIP("192.168.0.102"), + }, + }, + "192.168.0.101-192.168.0.105,fd00::2/112": { + count: new(big.Int).SetInt64(65541), + queries: map[uint64]net.IP{ + 0: net.ParseIP("192.168.0.101"), + 1: net.ParseIP("192.168.0.102"), + 5: net.ParseIP("fd00::0"), + 6: net.ParseIP("fd00::1"), + 65541: net.ParseIP("192.168.0.101"), + }, + }, + + "192.168.0.101,192.168.0.102,192.168.0.103,192.168.0.104,192.168.0.105-192.168.0.105,fd00::2/112": { + count: new(big.Int).SetInt64(65541), + queries: map[uint64]net.IP{ + 0: net.ParseIP("192.168.0.101"), + 1: net.ParseIP("192.168.0.102"), + 2: net.ParseIP("192.168.0.103"), + 3: net.ParseIP("192.168.0.104"), + 4: net.ParseIP("192.168.0.105"), + 5: net.ParseIP("fd00::0"), + 6: net.ParseIP("fd00::1"), + 65541: net.ParseIP("192.168.0.101"), + }, + }, + } + for name, data := range testdata { + name, data := name, data + t.Run(name, func(t *testing.T) { + p, err := NewIPPool(name) + require.NoError(t, err) + assert.Equal(t, data.count, p.count) + for q, a := range data.queries { + assert.Equal(t, a.To16(), p.GetIP(q).To16(), "index %d", q) + } + }) + } +} + +func TestIpBlockError(t *testing.T) { + testdata := map[string]string{ + "whatever": "not a valid IP", + "192.168.0.1012": "not a valid IP", + "192.168.0.10/244": "invalid CIDR", + "fd00::0/244": "invalid CIDR", + "192.168.0.101-192.168.0.102/32": "wrong IP range format", + "192.168.0.101-fd00::1": "mixed IP range format", + "fd00::1-192.168.0.101": "mixed IP range format", + "192.168.0.100-192.168.0.2": "negative IP range", + "fd00:1:1:0::0-fd00:1:0:ff::3ff": "negative IP range", + } + for name, data := range testdata { + name, data := name, data + t.Run(name, func(t *testing.T) { + _, err := getIPBlock(name) + require.Error(t, err) + require.Contains(t, err.Error(), data) + }) + } +}