Skip to content

Commit

Permalink
Modularize input handler
Browse files Browse the repository at this point in the history
  • Loading branch information
nakabonne committed Sep 15, 2020
1 parent 6d7aae4 commit dca82f4
Show file tree
Hide file tree
Showing 4 changed files with 115 additions and 59 deletions.
1 change: 1 addition & 0 deletions attacker/attacker.go
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ type Options struct {
Method string
Body []byte
Header http.Header

Attacker Attacker
}

Expand Down
81 changes: 22 additions & 59 deletions gui/gui.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,9 +3,7 @@ package gui
import (
"context"
"fmt"
"net/http"
"net/url"
"strconv"
"time"

"github.com/mum4k/termdash"
Expand Down Expand Up @@ -65,18 +63,19 @@ func gridLayout(w *widgets) ([]container.Option, error) {
raw1 := grid.RowHeightPerc(60,
grid.ColWidthPerc(99, grid.Widget(w.latencyChart, container.Border(linestyle.Light), container.BorderTitle("Latency (ms)"))),
)
raw2 := grid.RowHeightPerc(36,
raw2 := grid.RowHeightPerc(34,
grid.ColWidthPerc(64,
grid.RowHeightPerc(31, grid.Widget(w.urlInput, container.Border(linestyle.None))),
grid.RowHeightPerc(31,
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.RowHeightPerc(31,
grid.ColWidthPerc(50, grid.Widget(w.methodInput, container.Border(linestyle.None))),
grid.ColWidthPerc(49, grid.Widget(w.bodyInput, container.Border(linestyle.None))),
grid.RowHeightPerc(29,
grid.ColWidthPerc(33, grid.Widget(w.methodInput, container.Border(linestyle.None))),
grid.ColWidthPerc(33, grid.Widget(w.headerInput, container.Border(linestyle.None))),
grid.ColWidthPerc(33, grid.Widget(w.bodyInput, container.Border(linestyle.None))),
),
grid.RowHeightPerc(6, grid.Widget(w.progressGauge, container.Border(linestyle.None))),
grid.RowHeightPerc(12, grid.Widget(w.progressGauge, container.Border(linestyle.None))),
),
grid.ColWidthPerc(35, grid.Widget(w.reportText, container.Border(linestyle.Light), container.BorderTitle("Report"))),
)
Expand Down Expand Up @@ -105,64 +104,28 @@ func keybinds(ctx context.Context, cancel context.CancelFunc, dr *drawer) func(*
}
}

func attack(ctx context.Context, dr *drawer) {
if dr.chartDrawing {
func attack(ctx context.Context, d *drawer) {
if d.chartDrawing {
return
}
var (
target string
rate int
duration time.Duration
method string
body string
err error
)
target = dr.widgets.urlInput.Read()
target := d.widgets.urlInput.Read()
if _, err := url.ParseRequestURI(target); err != nil {
dr.reportCh <- fmt.Sprintf("Bad URL: %v", err)
d.reportCh <- fmt.Sprintf("Bad URL: %v", err)
return
}
if s := dr.widgets.rateLimitInput.Read(); s == "" {
rate = attacker.DefaultRate
} else {
rate, err = strconv.Atoi(s)
if err != nil {
dr.reportCh <- fmt.Sprintf("Given rate limit %q isn't integer: %v", s, err)
return
}
}
if s := dr.widgets.durationInput.Read(); s == "" {
duration = attacker.DefaultDuration
} else {
duration, err = time.ParseDuration(s)
if err != nil {
dr.reportCh <- fmt.Sprintf("Unparseable duration %q: %v", s, err)
return
}
}
if method = dr.widgets.methodInput.Read(); method != "" {
if !validateMethod(method) {
dr.reportCh <- fmt.Sprintf("Given method %q isn't an HTTP request method", method)
return
}
opts, err := makeOptions(d)
if err != nil {
d.reportCh <- err.Error()
return
}
body = dr.widgets.bodyInput.Read()
requestNum := rate * int(duration/time.Second)
requestNum := opts.Rate * int(opts.Duration/time.Second)

// To pre-allocate, run redrawChart on a per-attack basis.
go dr.redrawChart(ctx, requestNum)
go dr.redrawGauge(ctx, requestNum)
go func(ctx context.Context, d *drawer, t string, r int, du time.Duration) {
metrics := attacker.Attack(ctx, t, d.chartCh, attacker.Options{Rate: r, Duration: du, Method: method, Body: []byte(body)})
go d.redrawChart(ctx, requestNum)
go d.redrawGauge(ctx, requestNum)
go func(ctx context.Context, d *drawer, t string, o attacker.Options) {
metrics := attacker.Attack(ctx, t, d.chartCh, o)
d.reportCh <- metrics.String()
d.chartCh <- &attacker.Result{End: true}
}(ctx, dr, target, rate, duration)
}

func validateMethod(method string) bool {
switch method {
case http.MethodGet, http.MethodHead, http.MethodPost, http.MethodPut, http.MethodPatch, http.MethodDelete, http.MethodConnect, http.MethodOptions, http.MethodTrace:
return true
}
return false
}(ctx, d, target, opts)
}
80 changes: 80 additions & 0 deletions gui/input_handler.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,80 @@
package gui

import (
"fmt"
"net/http"
"strconv"
"strings"
"time"

"github.com/nakabonne/ali/attacker"
)

// makeOptions gives back an options for attacker, with the input from UI.
func makeOptions(d *drawer) (attacker.Options, error) {
var (
rate int
duration time.Duration
method string
body string
header = make(http.Header)
err error
)

if s := d.widgets.rateLimitInput.Read(); s == "" {
rate = attacker.DefaultRate
} else {
rate, err = strconv.Atoi(s)
if err != nil {
return attacker.Options{}, fmt.Errorf("Given rate limit %q isn't integer: %w", s, err)
}
}

if s := d.widgets.durationInput.Read(); s == "" {
duration = attacker.DefaultDuration
} else {
duration, err = time.ParseDuration(s)
if err != nil {
return attacker.Options{}, fmt.Errorf("Unparseable duration %q: %w", s, err)
}
}

if method = d.widgets.methodInput.Read(); method != "" {
if !validateMethod(method) {
return attacker.Options{}, fmt.Errorf("Given method %q isn't an HTTP request method", method)
}
}

body = d.widgets.bodyInput.Read()

if s := d.widgets.headerInput.Read(); s != "" {
parts := strings.SplitN(s, ":", 2)
if len(parts) != 2 {
return attacker.Options{}, fmt.Errorf("Given header %q has a wrong format", s)
}
key, val := strings.TrimSpace(parts[0]), strings.TrimSpace(parts[1])
if key == "" || val == "" {
return attacker.Options{}, fmt.Errorf("Given header %q has a wrong format", s)
}
// Add key/value directly to the http.Header (map[string][]string).
// http.Header.Add() canonicalizes keys but the vegeta API is used
// to test systems that require case-sensitive headers.
header[key] = append(header[key], val)
}

return attacker.Options{
Rate: rate,
Duration: duration,
Method: method,
Body: []byte(body),
Header: header,
}, nil
}

func validateMethod(method string) bool {
switch method {
case http.MethodGet, http.MethodHead, http.MethodPost, http.MethodPut, http.MethodPatch, http.MethodDelete, http.MethodConnect, http.MethodOptions, http.MethodTrace:
return true
}
return false
}
12 changes: 12 additions & 0 deletions gui/widgets.go
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,8 @@ type widgets struct {
durationInput *textinput.TextInput
methodInput *textinput.TextInput
bodyInput *textinput.TextInput
headerInput *textinput.TextInput
timeoutInput *textinput.TextInput
latencyChart *linechart.LineChart
reportText *text.Text
progressGauge *gauge.Gauge
Expand Down Expand Up @@ -61,6 +63,14 @@ func newWidgets() (*widgets, error) {
if err != nil {
return nil, err
}
headerInput, err := newTextInput("Header: ", "", 30)
if err != nil {
return nil, err
}
timeoutInput, err := newTextInput("Timeout: ", "30s", 30)
if err != nil {
return nil, err
}
progressDonut, err := newGauge()
if err != nil {
return nil, err
Expand All @@ -71,6 +81,8 @@ func newWidgets() (*widgets, error) {
durationInput: durationInput,
methodInput: methodInput,
bodyInput: bodyInput,
headerInput: headerInput,
timeoutInput: timeoutInput,
latencyChart: latencyChart,
reportText: reportText,
progressGauge: progressDonut,
Expand Down

0 comments on commit dca82f4

Please sign in to comment.