103 lines
3.2 KiB
Go
103 lines
3.2 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"flag"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"os"
|
|
"os/signal"
|
|
"strings"
|
|
"syscall"
|
|
"time"
|
|
|
|
"proxy-pool/internal/loadgen"
|
|
)
|
|
|
|
type loadRun func(context.Context, loadgen.Options) (loadgen.Report, error)
|
|
|
|
func main() {
|
|
ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
|
|
defer stop()
|
|
os.Exit(execute(ctx, os.Args[1:], loadgen.Run, os.Stdout, os.Stderr))
|
|
}
|
|
|
|
func execute(ctx context.Context, args []string, run loadRun, stdout, stderr io.Writer) int {
|
|
flags := flag.NewFlagSet("proxy-loadgen", flag.ContinueOnError)
|
|
flags.SetOutput(stderr)
|
|
targetURL := flags.String("target", "", "HTTP or HTTPS target URL")
|
|
proxyURL := flags.String("proxy", "", "optional HTTP or HTTPS forward proxy URL")
|
|
method := flags.String("method", http.MethodGet, "HTTP method")
|
|
body := flags.String("body", "", "UTF-8 request body")
|
|
requests := flags.Int("requests", 0, "fixed request count; mutually exclusive with -duration")
|
|
duration := flags.Duration("duration", 0, "time-boxed workload duration")
|
|
rate := flags.Int("rate", 0, "maximum request starts per second for -duration; zero saturates workers")
|
|
concurrency := flags.Int("concurrency", 64, "maximum concurrent requests")
|
|
timeout := flags.Duration("timeout", 10*time.Second, "per-request timeout")
|
|
var headers headerValues
|
|
flags.Var(&headers, "header", "repeatable HTTP header in Name: Value form")
|
|
if err := flags.Parse(args); err != nil {
|
|
if errors.Is(err, flag.ErrHelp) {
|
|
return 0
|
|
}
|
|
return 2
|
|
}
|
|
if flags.NArg() != 0 || ctx == nil || run == nil {
|
|
_, _ = fmt.Fprintln(stderr, "proxy-loadgen: target and a bounded workload are required")
|
|
return 2
|
|
}
|
|
parsedHeaders, err := headers.Header()
|
|
if err != nil {
|
|
_, _ = fmt.Fprintf(stderr, "proxy-loadgen: %v\n", err)
|
|
return 2
|
|
}
|
|
report, err := run(ctx, loadgen.Options{
|
|
TargetURL: *targetURL, ProxyURL: *proxyURL, Method: *method, Headers: parsedHeaders, RequestBody: []byte(*body),
|
|
Requests: *requests, Duration: *duration, Rate: *rate, Concurrency: *concurrency, RequestTimeout: *timeout,
|
|
})
|
|
if err != nil {
|
|
if errors.Is(err, loadgen.ErrInvalidOptions) {
|
|
_, _ = fmt.Fprintln(stderr, "proxy-loadgen: -target, positive -concurrency/-timeout, and exactly one bounded workload are required")
|
|
return 2
|
|
}
|
|
_, _ = fmt.Fprintf(stderr, "proxy-loadgen: %v\n", err)
|
|
return 1
|
|
}
|
|
if err := json.NewEncoder(stdout).Encode(report); err != nil {
|
|
_, _ = fmt.Fprintf(stderr, "proxy-loadgen: encode report: %v\n", err)
|
|
return 1
|
|
}
|
|
return 0
|
|
}
|
|
|
|
type headerValues []string
|
|
|
|
func (values *headerValues) String() string {
|
|
return strings.Join(*values, ",")
|
|
}
|
|
|
|
func (values *headerValues) Set(value string) error {
|
|
if strings.TrimSpace(value) != value || value == "" {
|
|
return errors.New("invalid header")
|
|
}
|
|
*values = append(*values, value)
|
|
return nil
|
|
}
|
|
|
|
func (values headerValues) Header() (http.Header, error) {
|
|
result := make(http.Header, len(values))
|
|
for _, value := range values {
|
|
name, content, found := strings.Cut(value, ":")
|
|
name = strings.TrimSpace(name)
|
|
content = strings.TrimSpace(content)
|
|
if !found || name == "" || strings.ContainsAny(name, "\r\n") || strings.ContainsAny(content, "\r\n") {
|
|
return nil, errors.New("invalid header")
|
|
}
|
|
result.Add(name, content)
|
|
}
|
|
return result, nil
|
|
}
|