feat: limit response body sizes

fasthttp allocates the body from the declared Content-Length before
reading it, so a server could exhaust memory by declaring a size it
never sends. The new max response body config bounds every response,
defaults to 10MiB and accepts sizes like 10MiB or 2MB. Scripts get the
same as a maxBodySize option on http().

internal/config/bytesize.go

@@ -0,0 +1,123 @@
+package config
+
+import (
+	"flag"
+	"fmt"
+	"math"
+	"strconv"
+	"strings"
+
+	"go.aykhans.me/sarin/internal/types"
+	"go.yaml.in/yaml/v4"
+)
+
+var (
+	_ flag.Value       = (*byteSize)(nil)
+	_ yaml.Unmarshaler = (*byteSize)(nil)
+)
+
+// byteSize is a size in bytes that reads and prints unit suffixes such as 10MiB.
+type byteSize uint64
+
+// byteSizeUnits run from the largest scale down, so formatting picks the largest unit
+// that fits. A bare "K" is out: the common libraries disagree on whether it counts
+// 1000 or 1024 bytes.
+var byteSizeUnits = []struct {
+	suffix string
+	scale  uint64
+}{
+	{"TiB", 1 << 40},
+	{"TB", 1000 * 1000 * 1000 * 1000},
+	{"GiB", 1 << 30},
+	{"GB", 1000 * 1000 * 1000},
+	{"MiB", 1 << 20},
+	{"MB", 1000 * 1000},
+	{"KiB", 1 << 10},
+	{"KB", 1000},
+	{"B", 1},
+}
+
+// parseByteSize reads a byte count, with or without a unit suffix.
+func parseByteSize(value string) (uint64, error) {
+	text := strings.TrimSpace(value)
+
+	// The suffix is split off by index, so no case mapping can shift the digits.
+	digitsEnd := len(text)
+	for digitsEnd > 0 && isASCIILetter(text[digitsEnd-1]) {
+		digitsEnd--
+	}
+
+	scale := uint64(1)
+	if suffix := text[digitsEnd:]; suffix != "" {
+		unitScale, ok := byteSizeScale(suffix)
+		if !ok {
+			return 0, types.ErrByteSizeInvalid
+		}
+		scale = unitScale
+	}
+
+	size, err := strconv.ParseUint(strings.TrimSpace(text[:digitsEnd]), 10, 64)
+	if err != nil {
+		return 0, types.ErrByteSizeInvalid
+	}
+	if size > math.MaxUint64/scale {
+		return 0, types.ErrByteSizeInvalid
+	}
+	return size * scale, nil
+}
+
+// byteSizeScale looks up a unit suffix, ignoring case.
+func byteSizeScale(suffix string) (uint64, bool) {
+	for _, unit := range byteSizeUnits {
+		if strings.EqualFold(suffix, unit.suffix) {
+			return unit.scale, true
+		}
+	}
+	return 0, false
+}
+
+func isASCIILetter(char byte) bool {
+	return (char >= 'a' && char <= 'z') || (char >= 'A' && char <= 'Z')
+}
+
+// formatByteSize renders a byte count so that parseByteSize reads back the same value.
+func formatByteSize(size uint64) string {
+	for _, unit := range byteSizeUnits {
+		if unit.scale > 1 && size >= unit.scale && size%unit.scale == 0 {
+			return strconv.FormatUint(size/unit.scale, 10) + unit.suffix
+		}
+	}
+	return strconv.FormatUint(size, 10)
+}
+
+func (size byteSize) String() string {
+	return formatByteSize(uint64(size))
+}
+
+// Set fills the size from a command line value.
+func (size *byteSize) Set(value string) error {
+	parsed, err := parseByteSize(value)
+	if err != nil {
+		return err
+	}
+	*size = byteSize(parsed)
+	return nil
+}
+
+// UnmarshalYAML accepts both a YAML integer and a string with a unit.
+func (size *byteSize) UnmarshalYAML(node *yaml.Node) error {
+	if node.Kind != yaml.ScalarNode {
+		return fmt.Errorf("expected a size, but got %s", node.Tag)
+	}
+
+	if node.Tag == "!!int" {
+		var parsed uint64
+		if err := node.Decode(&parsed); err != nil {
+			return types.ErrByteSizeInvalid
+		}
+		*size = byteSize(parsed)
+		return nil
+	}
+
+	return size.Set(node.Value)
+}

internal/config/cli.go

@@ -20,32 +20,33 @@ Simple usage:
 
 Flags:
   General Config:
-    -h, -help                     Help for sarin
-    -v, -version                  Version for sarin
-    -s, -show-config   bool       Show the final config after parsing all sources (default %v)
-    -f, -config-file   []string   Path to the config file (local file / http URL)
-    -c, -concurrency   uint       Number of concurrent requests (default %d)
-    -r, -requests      uint       Number of total requests
-    -d, -duration      time       Maximum duration for the test (e.g. 30s, 1m, 5h)
-    -l, -log-level     string     Runtime log levels to emit, comma-separated (possible values: info, error) (default %s)
-    -w, -log-file      string     Write runtime logs to this file instead of the terminal/stderr
-    -p, -progress      string     Progress display (possible values: bar, none) (default '%v')
-    -o, -output        string     Output format (possible values: table, json, yaml, none) (default '%v')
-    -z, -dry-run       bool       Run without sending requests (default %v)
+    -h, -help                         Help for sarin
+    -v, -version                      Version for sarin
+    -s, -show-config       bool       Show the final config after parsing all sources (default %v)
+    -f, -config-file       []string   Path to the config file (local file / http URL)
+    -c, -concurrency       uint       Number of concurrent requests (default %d)
+    -r, -requests          uint       Number of total requests
+    -d, -duration          time       Maximum duration for the test (e.g. 30s, 1m, 5h)
+    -l, -log-level         string     Runtime log levels to emit, comma-separated (possible values: info, error) (default %s)
+    -w, -log-file          string     Write runtime logs to this file instead of the terminal/stderr
+    -p, -progress          string     Progress display (possible values: bar, none) (default '%v')
+    -o, -output            string     Output format (possible values: table, json, yaml, none) (default '%v')
+    -z, -dry-run           bool       Run without sending requests (default %v)
 
   Request Config:
-    -U, -url           string     Target URL for the request
-    -M, -method        []string   HTTP method for the request (default %s)
-    -B, -body          []string   Body for the request (e.g. "body text")
-    -P, -param         []string   URL parameter for the request (e.g. "key1=value1")
-    -H, -header        []string   Header for the request (e.g. "key1: value1")
-    -C, -cookie        []string   Cookie for the request (e.g. "key1=value1")
-    -X, -proxy         []string   Proxy for the request (e.g. "http://proxy.example.com:8080")
-    -V, -values        []string   List of values for templating (e.g. "key1=value1")
-    -T, -timeout       time       Timeout for the request (e.g. 400ms, 3s, 1m10s) (default %v)
-    -I, -insecure      bool       Skip SSL/TLS certificate verification (default %v)
-        -lua           []string   Lua script for request transformation (inline or @file/@url)
-        -js            []string   JavaScript script for request transformation (inline or @file/@url)`
+    -U, -url               string     Target URL for the request
+    -M, -method            []string   HTTP method for the request (default %s)
+    -B, -body              []string   Body for the request (e.g. "body text")
+    -P, -param             []string   URL parameter for the request (e.g. "key1=value1")
+    -H, -header            []string   Header for the request (e.g. "key1: value1")
+    -C, -cookie            []string   Cookie for the request (e.g. "key1=value1")
+    -X, -proxy             []string   Proxy for the request (e.g. "http://proxy.example.com:8080")
+    -V, -values            []string   List of values for templating (e.g. "key1=value1")
+    -T, -timeout           time       Timeout for the request (e.g. 400ms, 3s, 1m10s) (default %v)
+    -S, -max-response-body size       Largest response body to read (e.g. 10MiB, 2MB), 0 for no limit (default %s)
+    -I, -insecure          bool       Skip SSL/TLS certificate verification (default %v)
+        -lua               []string   Lua script for request transformation (inline or @file/@url)
+        -js                []string   JavaScript script for request transformation (inline or @file/@url)`
 
 var _ IParser = ConfigCLIParser{}
 
@@ -97,18 +98,19 @@ func (parser ConfigCLIParser) Parse() (*Config, error) {
 		dryRun       bool
 
 		// Request config
-		urlInput   string
-		methods    = stringSliceArg{}
-		bodies     = stringSliceArg{}
-		params     = stringSliceArg{}
-		headers    = stringSliceArg{}
-		cookies    = stringSliceArg{}
-		proxies    = stringSliceArg{}
-		values     = stringSliceArg{}
-		timeout    time.Duration
-		insecure   bool
-		luaScripts = stringSliceArg{}
-		jsScripts  = stringSliceArg{}
+		urlInput        string
+		methods         = stringSliceArg{}
+		bodies          = stringSliceArg{}
+		params          = stringSliceArg{}
+		headers         = stringSliceArg{}
+		cookies         = stringSliceArg{}
+		proxies         = stringSliceArg{}
+		values          = stringSliceArg{}
+		timeout         time.Duration
+		maxResponseBody byteSize
+		insecure        bool
+		luaScripts      = stringSliceArg{}
+		jsScripts       = stringSliceArg{}
 	)
 
 	{
@@ -174,6 +176,9 @@ func (parser ConfigCLIParser) Parse() (*Config, error) {
 		flagSet.DurationVar(&timeout, "timeout", 0, "Timeout for the request (e.g. 400ms, 15s, 1m10s)")
 		flagSet.DurationVar(&timeout, "T", 0, "Timeout for the request (e.g. 400ms, 15s, 1m10s)")
 
+		flagSet.Var(&maxResponseBody, "max-response-body", "Largest response body to read (e.g. 1048576, 10MiB, 2MB), 0 for no limit")
+		flagSet.Var(&maxResponseBody, "S", "Largest response body to read (e.g. 1048576, 10MiB, 2MB), 0 for no limit")
+
 		flagSet.BoolVar(&insecure, "insecure", false, "Skip SSL/TLS certificate verification")
 		flagSet.BoolVar(&insecure, "I", false, "Skip SSL/TLS certificate verification")
 
@@ -258,6 +263,8 @@ func (parser ConfigCLIParser) Parse() (*Config, error) {
 			config.Values = append(config.Values, values...)
 		case "timeout", "T":
 			config.Timeout = new(timeout)
+		case "max-response-body", "S":
+			config.MaxResponseBody = new(uint64(maxResponseBody))
 		case "insecure", "I":
 			config.Insecure = new(insecure)
 		case "lua":
@@ -286,6 +293,7 @@ func (parser ConfigCLIParser) PrintHelp() {
 
 		Defaults.Method,
 		Defaults.RequestTimeout,
+		formatByteSize(Defaults.MaxResponseBody),
 		Defaults.Insecure,
 	)
 }

internal/config/config.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"math"
 	"net/url"
 	"os"
 	"path/filepath"
@@ -27,27 +28,29 @@ import (
 )
 
 var Defaults = struct {
-	UserAgent      string
-	Method         string
-	RequestTimeout time.Duration
-	Concurrency    uint
-	ShowConfig     bool
-	Progress       ConfigProgressType
-	Insecure       bool
-	Output         ConfigOutputType
-	DryRun         bool
-	LogLevel       string
+	UserAgent       string
+	Method          string
+	RequestTimeout  time.Duration
+	MaxResponseBody uint64
+	Concurrency     uint
+	ShowConfig      bool
+	Progress        ConfigProgressType
+	Insecure        bool
+	Output          ConfigOutputType
+	DryRun          bool
+	LogLevel        string
 }{
-	UserAgent:      "Sarin/" + version.Version,
-	Method:         "GET",
-	RequestTimeout: time.Second * 10,
-	Concurrency:    1,
-	ShowConfig:     false,
-	Progress:       ConfigProgressTypeBar,
-	Insecure:       false,
-	Output:         ConfigOutputTypeTable,
-	DryRun:         false,
-	LogLevel:       "error",
+	UserAgent:       "Sarin/" + version.Version,
+	Method:          "GET",
+	RequestTimeout:  time.Second * 10,
+	MaxResponseBody: 10 << 20, // 10 MiB
+	Concurrency:     1,
+	ShowConfig:      false,
+	Progress:        ConfigProgressTypeBar,
+	Insecure:        false,
+	Output:          ConfigOutputTypeTable,
+	DryRun:          false,
+	LogLevel:        "error",
 }
 
 var (
@@ -56,6 +59,8 @@ var (
 	ValidLogLevels         = []string{"info", "error"}
 )
 
+const maxResponseBodyLimit = min(1<<40, math.MaxInt) // 1 TiB on 64 bit, 2 GiB - 1 on 32 bit
+
 var (
 	StyleYellow = lipgloss.NewStyle().Foreground(lipgloss.Color("220"))
 	StyleRed    = lipgloss.NewStyle().Foreground(lipgloss.Color("196"))
@@ -82,28 +87,29 @@ var (
 )
 
 type Config struct {
-	ShowConfig  *bool               `yaml:"showConfig,omitempty"`
-	Files       []types.ConfigFile  `yaml:"files,omitempty"`
-	Methods     []string            `yaml:"methods,omitempty"`
-	URL         *url.URL            `yaml:"url,omitempty"`
-	Timeout     *time.Duration      `yaml:"timeout,omitempty"`
-	Concurrency *uint               `yaml:"concurrency,omitempty"`
-	Requests    *uint64             `yaml:"requests,omitempty"`
-	Duration    *time.Duration      `yaml:"duration,omitempty"`
-	Progress    *ConfigProgressType `yaml:"progress,omitempty"`
-	Output      *ConfigOutputType   `yaml:"output,omitempty"`
-	Insecure    *bool               `yaml:"insecure,omitempty"`
-	DryRun      *bool               `yaml:"dryRun,omitempty"`
-	Params      types.Params        `yaml:"params,omitempty"`
-	Headers     types.Headers       `yaml:"headers,omitempty"`
-	Cookies     types.Cookies       `yaml:"cookies,omitempty"`
-	Bodies      []string            `yaml:"bodies,omitempty"`
-	Proxies     types.Proxies       `yaml:"proxies,omitempty"`
-	Values      []string            `yaml:"values,omitempty"`
-	Lua         []string            `yaml:"lua,omitempty"`
-	Js          []string            `yaml:"js,omitempty"`
-	LogLevel    *string             `yaml:"logLevel,omitempty"`
-	LogFile     *string             `yaml:"logFile,omitempty"`
+	ShowConfig      *bool               `yaml:"showConfig,omitempty"`
+	Files           []types.ConfigFile  `yaml:"files,omitempty"`
+	Methods         []string            `yaml:"methods,omitempty"`
+	URL             *url.URL            `yaml:"url,omitempty"`
+	Timeout         *time.Duration      `yaml:"timeout,omitempty"`
+	MaxResponseBody *uint64             `yaml:"maxResponseBody,omitempty"`
+	Concurrency     *uint               `yaml:"concurrency,omitempty"`
+	Requests        *uint64             `yaml:"requests,omitempty"`
+	Duration        *time.Duration      `yaml:"duration,omitempty"`
+	Progress        *ConfigProgressType `yaml:"progress,omitempty"`
+	Output          *ConfigOutputType   `yaml:"output,omitempty"`
+	Insecure        *bool               `yaml:"insecure,omitempty"`
+	DryRun          *bool               `yaml:"dryRun,omitempty"`
+	Params          types.Params        `yaml:"params,omitempty"`
+	Headers         types.Headers       `yaml:"headers,omitempty"`
+	Cookies         types.Cookies       `yaml:"cookies,omitempty"`
+	Bodies          []string            `yaml:"bodies,omitempty"`
+	Proxies         types.Proxies       `yaml:"proxies,omitempty"`
+	Values          []string            `yaml:"values,omitempty"`
+	Lua             []string            `yaml:"lua,omitempty"`
+	Js              []string            `yaml:"js,omitempty"`
+	LogLevel        *string             `yaml:"logLevel,omitempty"`
+	LogFile         *string             `yaml:"logFile,omitempty"`
 }
 
 func (config Config) MarshalYAML() (any, error) {
@@ -177,6 +183,9 @@ func (config Config) MarshalYAML() (any, error) {
 	if config.Timeout != nil {
 		addField(content, "timeout", toNode(*config.Timeout), "")
 	}
+	if config.MaxResponseBody != nil {
+		addField(content, "maxResponseBody", toNode(formatByteSize(*config.MaxResponseBody)), "")
+	}
 	if config.Concurrency != nil {
 		addField(content, "concurrency", toNode(*config.Concurrency), "")
 	}
@@ -301,6 +310,9 @@ func (config *Config) Merge(newConfig *Config) {
 	if newConfig.Timeout != nil {
 		config.Timeout = newConfig.Timeout
 	}
+	if newConfig.MaxResponseBody != nil {
+		config.MaxResponseBody = newConfig.MaxResponseBody
+	}
 	if newConfig.Concurrency != nil {
 		config.Concurrency = newConfig.Concurrency
 	}
@@ -379,6 +391,9 @@ func (config *Config) SetDefaults() {
 	if config.Timeout == nil {
 		config.Timeout = &Defaults.RequestTimeout
 	}
+	if config.MaxResponseBody == nil {
+		config.MaxResponseBody = &Defaults.MaxResponseBody
+	}
 	if config.Concurrency == nil {
 		config.Concurrency = new(Defaults.Concurrency)
 	}
@@ -454,6 +469,14 @@ func (config Config) Validate() error {
 		validationErrors = append(validationErrors, types.NewFieldValidationError("Timeout", "0", errors.New("timeout must be greater than 0")))
 	}
 
+	if config.MaxResponseBody != nil && *config.MaxResponseBody > maxResponseBodyLimit {
+		validationErrors = append(validationErrors, types.NewFieldValidationError(
+			"MaxResponseBody",
+			formatByteSize(*config.MaxResponseBody),
+			errors.New("max response body must not exceed "+formatByteSize(maxResponseBodyLimit)),
+		))
+	}
+
 	if config.ShowConfig == nil {
 		validationErrors = append(validationErrors, types.NewFieldValidationError("ShowConfig", "", errors.New("showConfig field is required")))
 	}

internal/config/env.go

@@ -195,6 +195,22 @@ func (parser ConfigENVParser) Parse() (*Config, error) {
 		}
 	}
 
+	if maxResponseBody := parser.getEnv("MAX_RESPONSE_BODY"); maxResponseBody != "" {
+		maxResponseBodyParsed, err := parseByteSize(maxResponseBody)
+		if err != nil {
+			fieldParseErrors = append(
+				fieldParseErrors,
+				types.NewFieldParseError(
+					parser.getFullEnvName("MAX_RESPONSE_BODY"),
+					maxResponseBody,
+					err,
+				),
+			)
+		} else {
+			config.MaxResponseBody = &maxResponseBodyParsed
+		}
+	}
+
 	if insecure := parser.getEnv("INSECURE"); insecure != "" {
 		insecureParsed, err := utilsParse.ParseString[bool](insecure)
 		if err != nil {

internal/config/file.go

@@ -192,28 +192,29 @@ func (kv *keyValuesField) unmarshalMapping(node *yaml.Node) error {
 }
 
 type configYAML struct {
-	ShowConfig   *bool              `yaml:"showConfig"`
-	ConfigFiles  stringOrSliceField `yaml:"configFile"`
-	Concurrency  *uint              `yaml:"concurrency"`
-	RequestCount *uint64            `yaml:"requests"`
-	Duration     *time.Duration     `yaml:"duration"`
-	LogLevel     *string            `yaml:"logLevel"`
-	LogFile      *string            `yaml:"logFile"`
-	Progress     *string            `yaml:"progress"`
-	Output       *string            `yaml:"output"`
-	DryRun       *bool              `yaml:"dryRun"`
-	URL          *string            `yaml:"url"`
-	Method       stringOrSliceField `yaml:"method"`
-	Bodies       stringOrSliceField `yaml:"body"`
-	Params       keyValuesField     `yaml:"params"`
-	Headers      keyValuesField     `yaml:"headers"`
-	Cookies      keyValuesField     `yaml:"cookies"`
-	Proxies      stringOrSliceField `yaml:"proxy"`
-	Values       stringOrSliceField `yaml:"values"`
-	Timeout      *time.Duration     `yaml:"timeout"`
-	Insecure     *bool              `yaml:"insecure"`
-	Lua          stringOrSliceField `yaml:"lua"`
-	Js           stringOrSliceField `yaml:"js"`
+	ShowConfig      *bool              `yaml:"showConfig"`
+	ConfigFiles     stringOrSliceField `yaml:"configFile"`
+	Concurrency     *uint              `yaml:"concurrency"`
+	RequestCount    *uint64            `yaml:"requests"`
+	Duration        *time.Duration     `yaml:"duration"`
+	LogLevel        *string            `yaml:"logLevel"`
+	LogFile         *string            `yaml:"logFile"`
+	Progress        *string            `yaml:"progress"`
+	Output          *string            `yaml:"output"`
+	DryRun          *bool              `yaml:"dryRun"`
+	URL             *string            `yaml:"url"`
+	Method          stringOrSliceField `yaml:"method"`
+	Bodies          stringOrSliceField `yaml:"body"`
+	Params          keyValuesField     `yaml:"params"`
+	Headers         keyValuesField     `yaml:"headers"`
+	Cookies         keyValuesField     `yaml:"cookies"`
+	Proxies         stringOrSliceField `yaml:"proxy"`
+	Values          stringOrSliceField `yaml:"values"`
+	Timeout         *time.Duration     `yaml:"timeout"`
+	MaxResponseBody *byteSize          `yaml:"maxResponseBody"`
+	Insecure        *bool              `yaml:"insecure"`
+	Lua             stringOrSliceField `yaml:"lua"`
+	Js              stringOrSliceField `yaml:"js"`
 }
 
 // ParseYAML parses YAML config file arguments into a Config object.
@@ -290,6 +291,9 @@ func (parser ConfigFileParser) ParseYAML(data []byte) (*Config, error) {
 
 	config.Values = append(config.Values, parsedData.Values...)
 	config.Timeout = parsedData.Timeout
+	if parsedData.MaxResponseBody != nil {
+		config.MaxResponseBody = new(uint64(*parsedData.MaxResponseBody))
+	}
 	config.Insecure = parsedData.Insecure
 	config.Lua = append(config.Lua, parsedData.Lua...)
 	config.Js = append(config.Js, parsedData.Js...)