Merge pull request #54 from markruler/support-local-address

Support local address option

attacker/attacker.go

@@ -3,6 +3,7 @@ package attacker
 import (
 	"context"
 	"math"
+	"net"
 	"net/http"
 	"time"
 
@@ -39,6 +40,7 @@ type Options struct {
 	KeepAlive   bool
 	Connections int
 	HTTP2       bool
+	LocalAddr   net.IPAddr
 
 	Attacker Attacker
 }
@@ -72,6 +74,9 @@ func Attack(ctx context.Context, target string, resCh chan *Result, metricsCh ch
 	if opts.Connections == 0 {
 		opts.Connections = DefaultConnections
 	}
+	if opts.LocalAddr.IP == nil {
+		opts.LocalAddr = vegeta.DefaultLocalAddr
+	}
 	if opts.Attacker == nil {
 		opts.Attacker = vegeta.NewAttacker(
 			vegeta.Timeout(opts.Timeout),
@@ -81,6 +86,7 @@ func Attack(ctx context.Context, target string, resCh chan *Result, metricsCh ch
 			vegeta.Connections(opts.Connections),
 			vegeta.KeepAlive(opts.KeepAlive),
 			vegeta.HTTP2(opts.HTTP2),
+			vegeta.LocalAddr(opts.LocalAddr),
 		)
 	}
 

main.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"io"
 	"io/ioutil"
+	"net"
 	"net/http"
 	"net/url"
 	"os"
@@ -30,18 +31,19 @@ var (
 )
 
 type cli struct {
-	rate        int
-	duration    time.Duration
-	timeout     time.Duration
-	method      string
-	headers     []string
-	body        string
-	bodyFile    string
-	maxBody     int64
-	workers     uint64
-	maxWorkers  uint64
-	connections int
-	http2       bool
+	rate         int
+	duration     time.Duration
+	timeout      time.Duration
+	method       string
+	headers      []string
+	body         string
+	bodyFile     string
+	maxBody      int64
+	workers      uint64
+	maxWorkers   uint64
+	connections  int
+	http2        bool
+	localAddress string
 
 	debug     bool
 	version   bool
@@ -78,6 +80,7 @@ func parseFlags(stdout, stderr io.Writer) (*cli, error) {
 	flagSet.Uint64VarP(&c.maxWorkers, "max-workers", "W", attacker.DefaultMaxWorkers, "Amount of maximum workers to spawn.")
 	flagSet.IntVarP(&c.connections, "connections", "c", attacker.DefaultConnections, "Amount of maximum open idle connections per target host")
 	flagSet.BoolVar(&c.http2, "http2", true, "Issue HTTP/2 requests to servers which support it. (default true)")
+	flagSet.StringVarP(&c.localAddress, "local-address", "L", "0.0.0.0", "Local IP address.")
 	flagSet.Usage = c.usage
 	if err := flagSet.Parse(os.Args[1:]); err != nil {
 		if !errors.Is(err, flag.ErrHelp) {
@@ -169,6 +172,8 @@ func (c *cli) makeOptions() (*attacker.Options, error) {
 		body = b
 	}
 
+	localAddr := net.IPAddr{IP: net.ParseIP(c.localAddress)}
+
 	return &attacker.Options{
 		Rate:        c.rate,
 		Duration:    c.duration,
@@ -182,6 +187,7 @@ func (c *cli) makeOptions() (*attacker.Options, error) {
 		MaxWorkers:  c.maxWorkers,
 		Connections: c.connections,
 		HTTP2:       c.http2,
+		LocalAddr:   localAddr,
 	}, nil
 }
 

main_test.go

@@ -46,19 +46,20 @@ func TestParseFlags(t *testing.T) {
 		{
 			name: "with default options",
 			want: &cli{
-				rate:        50,
-				duration:    time.Second * 10,
-				timeout:     time.Second * 30,
-				method:      "GET",
-				headers:     []string{},
-				maxBody:     -1,
-				keepAlive:   true,
-				workers:     10,
-				maxWorkers:  math.MaxUint64,
-				connections: 10000,
-				stdout:      new(bytes.Buffer),
-				stderr:      new(bytes.Buffer),
-				http2:       true,
+				rate:         50,
+				duration:     time.Second * 10,
+				timeout:      time.Second * 30,
+				method:       "GET",
+				headers:      []string{},
+				maxBody:      -1,
+				keepAlive:    true,
+				workers:      10,
+				maxWorkers:   math.MaxUint64,
+				connections:  10000,
+				stdout:       new(bytes.Buffer),
+				stderr:       new(bytes.Buffer),
+				http2:        true,
+				localAddress: "0.0.0.0",
 			},
 			wantErr: false,
 		},