Enable to specify timeout

nakabonne · Sep 15, 2020 · fe260b7 · fe260b7
1 parent 8965080
commit fe260b7
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 3 deletions.
diff --git a/attacker/attacker.go b/attacker/attacker.go
@@ -11,6 +11,7 @@ import (
 const (
 	DefaultRate     = 50
 	DefaultDuration = 10 * time.Second
+	DefaultTimeout  = 30 * time.Second
 	DefaultMethod   = http.MethodGet
 )
 
@@ -23,6 +24,7 @@ type Attacker interface {
 type Options struct {
 	Rate     int
 	Duration time.Duration
+	Timeout  time.Duration
 	Method   string
 	Body     []byte
 	Header   http.Header
@@ -50,11 +52,14 @@ func Attack(ctx context.Context, target string, resCh chan *Result, opts Options
 	if opts.Duration == 0 {
 		opts.Duration = DefaultDuration
 	}
+	if opts.Timeout == 0 {
+		opts.Timeout = DefaultTimeout
+	}
 	if opts.Method == "" {
 		opts.Method = DefaultMethod
 	}
 	if opts.Attacker == nil {
-		opts.Attacker = vegeta.NewAttacker()
+		opts.Attacker = vegeta.NewAttacker(vegeta.Timeout(opts.Timeout))
 	}
 
 	rate := vegeta.Rate{Freq: opts.Rate, Per: time.Second}

diff --git a/gui/gui.go b/gui/gui.go
@@ -67,8 +67,9 @@ func gridLayout(w *widgets) ([]container.Option, error) {
 		grid.ColWidthPerc(64,
 			grid.RowHeightPerc(29, grid.Widget(w.urlInput, container.Border(linestyle.None))),
 			grid.RowHeightPerc(29,
-				grid.ColWidthPerc(50, grid.Widget(w.rateLimitInput, container.Border(linestyle.None))),
-				grid.ColWidthPerc(49, grid.Widget(w.durationInput, container.Border(linestyle.None))),
+				grid.ColWidthPerc(33, grid.Widget(w.rateLimitInput, container.Border(linestyle.None))),
+				grid.ColWidthPerc(33, grid.Widget(w.durationInput, container.Border(linestyle.None))),
+				grid.ColWidthPerc(33, grid.Widget(w.timeoutInput, container.Border(linestyle.None))),
 			),
 			grid.RowHeightPerc(29,
 				grid.ColWidthPerc(33, grid.Widget(w.methodInput, container.Border(linestyle.None))),

diff --git a/gui/input_handler.go b/gui/input_handler.go
@@ -15,6 +15,7 @@ func makeOptions(w *widgets) (attacker.Options, error) {
 	var (
 		rate     int
 		duration time.Duration
+		timeout  time.Duration
 		method   string
 		body     string
 		header   = make(http.Header)
@@ -39,6 +40,15 @@ func makeOptions(w *widgets) (attacker.Options, error) {
 		}
 	}
 
+	if s := w.timeoutInput.Read(); s == "" {
+		timeout = attacker.DefaultTimeout
+	} else {
+		timeout, err = time.ParseDuration(s)
+		if err != nil {
+			return attacker.Options{}, fmt.Errorf("Unparseable timeout %q: %w", s, err)
+		}
+	}
+
 	if method = w.methodInput.Read(); method != "" {
 		if !validateMethod(method) {
 			return attacker.Options{}, fmt.Errorf("Given method %q isn't an HTTP request method", method)
@@ -65,6 +75,7 @@ func makeOptions(w *widgets) (attacker.Options, error) {
 	return attacker.Options{
 		Rate:     rate,
 		Duration: duration,
+		Timeout:  timeout,
 		Method:   method,
 		Body:     []byte(body),
 		Header:   header,