From 137c59bb64df1f12a395c18e2ae9be2c7113e08a Mon Sep 17 00:00:00 2001 From: six-ddc Date: Sun, 27 Jun 2021 10:41:41 +0800 Subject: [PATCH] Support --key, --cert, --insecure for TLS, and read flags from ENV --- charts.go | 6 ++- main.go | 106 +++++++++++++++++++++++++++++++++++++++++++-------- requester.go | 26 +++++++++++++ 3 files changed, 120 insertions(+), 18 deletions(-) diff --git a/charts.go b/charts.go index 55c2d2a..33afe71 100644 --- a/charts.go +++ b/charts.go @@ -200,11 +200,13 @@ func (c *Charts) Handler(ctx *fasthttp.RequestCtx) { } } -func (c *Charts) Serve() { +func (c *Charts) Serve(open bool) { server := fasthttp.Server{ Handler: cors.DefaultHandler().CorsMiddleware(c.Handler), } - go openBrowser("http://" + c.ln.Addr().String()) + if open { + go openBrowser("http://" + c.ln.Addr().String()) + } _ = server.Serve(c.ln) } diff --git a/main.go b/main.go index 6442ab5..c7a0074 100644 --- a/main.go +++ b/main.go @@ -23,6 +23,9 @@ var ( headers = kingpin.Flag("header", "Custom HTTP headers").Short('H').PlaceHolder("K:V").Strings() host = kingpin.Flag("host", "Host header").String() contentType = kingpin.Flag("content", "Content-Type header").Short('T').String() + cert = kingpin.Flag("cert", "Path to the client's TLS Certificate").ExistingFile() + key = kingpin.Flag("key", "Path to the client's TLS Certificate Private Key").ExistingFile() + insecure = kingpin.Flag("insecure", "Controls whether a client verifies the server's certificate chain and host name").Short('k').Bool() chartsListenAddr = kingpin.Flag("listen", "Listen addr to serve Web UI").Default(":18888").String() timeout = kingpin.Flag("timeout", "Timeout for each http request").PlaceHolder("DURATION").Duration() @@ -31,27 +34,86 @@ var ( respReadTimeout = kingpin.Flag("resp-timeout", "Timeout for full response reading").PlaceHolder("DURATION").Duration() socks5 = kingpin.Flag("socks5", "Socks5 proxy").PlaceHolder("ip:port").String() + autoOpenBrowser = kingpin.Flag("auto-open-browser", "Specify whether auto open browser to show Web charts").Bool() + url = kingpin.Arg("url", "request url").Required().String() ) func errAndExit(msg string) { - fmt.Fprintln(os.Stderr, msg) + fmt.Fprintln(os.Stderr, "plow: "+msg) os.Exit(1) } -func main() { - kingpin.UsageTemplate(kingpin.CompactUsageTemplate).Version("1.0.0").Author("six-ddc@github") - kingpin.CommandLine.Help = `A high-performance HTTP benchmarking tool with real-time web UI and terminal displaying +var CompactUsageTemplate = `{{define "FormatCommand" -}} +{{if .FlagSummary}} {{.FlagSummary}}{{end -}} +{{range .Args}} {{if not .Required}}[{{end}}<{{.Name}}>{{if .Value|IsCumulative}} ...{{end}}{{if not .Required}}]{{end}}{{end -}} +{{end -}} + +{{define "FormatCommandList" -}} +{{range . -}} +{{if not .Hidden -}} +{{.Depth|Indent}}{{.Name}}{{if .Default}}*{{end}}{{template "FormatCommand" .}} +{{end -}} +{{template "FormatCommandList" .Commands -}} +{{end -}} +{{end -}} + +{{define "FormatUsage" -}} +{{template "FormatCommand" .}}{{if .Commands}} [ ...]{{end}} +{{if .Help}} +{{.Help|Wrap 0 -}} +{{end -}} + +{{end -}} + +{{if .Context.SelectedCommand -}} +{{T "usage:"}} {{.App.Name}} {{template "FormatUsage" .Context.SelectedCommand}} +{{else -}} +{{T "usage:"}} {{.App.Name}}{{template "FormatUsage" .App}} +{{end -}} +Examples: -Example: plow http://127.0.0.1:8080/ -c 20 -n 100000 plow https://httpbin.org/post -c 20 -d 5m --body @file.json -T 'application/json' -m POST + +{{if .Context.Flags -}} +{{T "Flags:"}} +{{.Context.Flags|FlagsToTwoColumns|FormatTwoColumns}} + Flags default values also read from env PLOW_SOME_FLAG, such as PLOW_TIMEOUT=5s equals to --timeout=5s + +{{end -}} +{{if .Context.Args -}} +{{T "Args:"}} +{{.Context.Args|ArgsToTwoColumns|FormatTwoColumns}} +{{end -}} +{{if .Context.SelectedCommand -}} +{{if .Context.SelectedCommand.Commands -}} +{{T "Commands:"}} + {{.Context.SelectedCommand}} +{{.Context.SelectedCommand.Commands|CommandsToTwoColumns|FormatTwoColumns}} +{{end -}} +{{else if .App.Commands -}} +{{T "Commands:"}} +{{.App.Commands|CommandsToTwoColumns|FormatTwoColumns}} +{{end -}} ` + +func main() { + kingpin.UsageTemplate(CompactUsageTemplate). + Version("1.0.0"). + Author("six-ddc@github"). + Resolver(kingpin.PrefixedEnvarResolver("PLOW_", ";")). + Help = `A high-performance HTTP benchmarking tool with real-time web UI and terminal displaying` kingpin.Parse() + if *requests >= 0 && *requests < int64(*concurrency) { errAndExit("requests must greater than or equal concurrency") return } + if (*cert != "" && *key == "") || (*cert == "" && *key != "") { + errAndExit("must specify cert and key at the same time") + return + } var err error var bodyBytes []byte @@ -76,11 +138,16 @@ Example: } clientOpt := ClientOpt{ - url: *url, - method: *method, - headers: *headers, - bodyBytes: bodyBytes, - bodyFile: bodyFile, + url: *url, + method: *method, + headers: *headers, + bodyBytes: bodyBytes, + bodyFile: bodyFile, + + certPath: *cert, + keyPath: *key, + insecure: *insecure, + maxConns: *concurrency, doTimeout: *timeout, readTimeout: *respReadTimeout, @@ -92,6 +159,13 @@ Example: host: *host, } + requester, err := NewRequester(*concurrency, *requests, *duration, &clientOpt) + if err != nil { + errAndExit(err.Error()) + return + } + + // description var desc string desc = fmt.Sprintf("Benchmarking %s", *url) if *requests > 0 { @@ -103,6 +177,7 @@ Example: desc += fmt.Sprintf(" using %d connection(s).", *concurrency) fmt.Println(desc) + // charts listener var ln net.Listener if *chartsListenAddr != "" { ln, err = net.Listen("tcp", *chartsListenAddr) @@ -114,25 +189,24 @@ Example: } fmt.Printf("\n") - requester, err := NewRequester(*concurrency, *requests, *duration, &clientOpt) - if err != nil { - errAndExit(err.Error()) - return - } + // do request go requester.Run() + // metrics collection report := NewStreamReport() go report.Collect(requester.RecordChan()) if ln != nil { + // serve charts data charts, err := NewCharts(ln, report.Charts, desc) if err != nil { errAndExit(err.Error()) return } - go charts.Serve() + go charts.Serve(*autoOpenBrowser) } + // terminal printer printer := NewPrinter(*requests, *duration) printer.PrintLoop(report.Snapshot, *interval, *seconds, report.Done()) } diff --git a/requester.go b/requester.go index d47151e..fd27dae 100644 --- a/requester.go +++ b/requester.go @@ -2,6 +2,7 @@ package main import ( "context" + "crypto/tls" "fmt" "github.com/valyala/fasthttp" "github.com/valyala/fasthttp/fasthttpproxy" @@ -112,6 +113,10 @@ type ClientOpt struct { bodyBytes []byte bodyFile string + certPath string + keyPath string + insecure bool + maxConns int doTimeout time.Duration readTimeout time.Duration @@ -156,6 +161,21 @@ func addMissingPort(addr string, isTLS bool) string { return net.JoinHostPort(addr, strconv.Itoa(port)) } +func buildTLSConfig(opt *ClientOpt) (*tls.Config, error) { + var certs []tls.Certificate + if opt.certPath != "" && opt.keyPath != "" { + c, err := tls.LoadX509KeyPair(opt.certPath, opt.keyPath) + if err != nil { + return nil, err + } + certs = append(certs, c) + } + return &tls.Config{ + InsecureSkipVerify: opt.insecure, + Certificates: certs, + }, nil +} + func buildRequestClient(opt *ClientOpt, r *int64, w *int64) (*fasthttp.HostClient, *fasthttp.RequestHeader, error) { u, err := url2.Parse(opt.url) if err != nil { @@ -180,6 +200,12 @@ func buildRequestClient(opt *ClientOpt, r *int64, w *int64) (*fasthttp.HostClien } httpClient.Dial = ThroughputInterceptorDial(httpClient.Dial, r, w) + tlsConfig, err := buildTLSConfig(opt) + if err != nil { + return nil, nil, err + } + httpClient.TLSConfig = tlsConfig + var requestHeader fasthttp.RequestHeader if opt.contentType != "" { requestHeader.SetContentType(opt.contentType)