Add HTTP size limits

This commit is contained in:
2026-07-05 00:17:07 +00:00
parent a16f66cbc7
commit f7d821067f
15 changed files with 609 additions and 46 deletions

View File

@@ -74,11 +74,14 @@ type renderConfig struct {
type serveConfig struct {
configPath string
addr string
promptDir string
profileDir string
schemaDir string
artifactRoot string
addr string
promptDir string
profileDir string
schemaDir string
artifactRoot string
maxRequestBytes int64
maxArtifactBytes int64
maxResponseBytes int64
}
type commonCommandSettings struct {
@@ -87,6 +90,9 @@ type commonCommandSettings struct {
schemaDir string
serverAddr string
artifactRoot string
maxRequestBytes int64
maxArtifactBytes int64
maxResponseBytes int64
defaultRenderFormat renderformat.PreparedRunOutputFormat
}
@@ -204,7 +210,7 @@ func serveCommand(args []string, stderr io.Writer) int {
return ExitRuntimeError
}
artifactReader, err := artifactadapter.NewRestrictedCompositeReader(cfg.artifactRoot)
artifactReader, err := artifactadapter.NewRestrictedCompositeReaderWithLimit(cfg.artifactRoot, cfg.maxArtifactBytes)
if err != nil {
fmt.Fprintf(stderr, "artifact root error: %v\n", err)
return ExitRuntimeError
@@ -212,7 +218,10 @@ func serveCommand(args []string, stderr io.Writer) int {
runner := newRunnerWithArtifactReader(cfg.promptDir, cfg.profileDir, cfg.schemaDir, llmClient, artifactReader)
h := httpadapter.NewHandler(runner)
h := httpadapter.NewHandlerWithOptions(runner, httpadapter.HandlerOptions{
MaxRequestBytes: cfg.maxRequestBytes,
MaxResponseBytes: cfg.maxResponseBytes,
})
srv := &http.Server{
Addr: cfg.addr,
Handler: h,
@@ -290,6 +299,9 @@ func parseServeArgs(args []string) (*serveConfig, error) {
fs.StringVar(&cfg.profileDir, "profile-dir", "", "directory containing execution profile YAML files")
fs.StringVar(&cfg.schemaDir, "schema-dir", "", "base directory for validation schemas")
fs.StringVar(&cfg.artifactRoot, "artifact-root", "", "base directory for HTTP file input artifacts")
fs.Int64Var(&cfg.maxRequestBytes, "max-request-bytes", 0, "maximum HTTP request body bytes; 0 disables the limit")
fs.Int64Var(&cfg.maxArtifactBytes, "max-artifact-bytes", 0, "maximum HTTP file artifact bytes; 0 disables the limit")
fs.Int64Var(&cfg.maxResponseBytes, "max-response-bytes", 0, "maximum HTTP response body bytes; 0 disables the limit")
if err := fs.Parse(args); err != nil {
return nil, err
@@ -299,11 +311,14 @@ func parseServeArgs(args []string) (*serveConfig, error) {
}
settings, err := resolveCommonSettings(fs, cfg.configPath, appconfig.CLIOverrides{
PromptDir: cfg.promptDirIfSet(fs),
ProfileDir: cfg.profileDirIfSet(fs),
SchemaDir: cfg.schemaDirIfSet(fs),
ServerAddr: cfg.addrIfSet(fs),
ArtifactRoot: cfg.artifactRootIfSet(fs),
PromptDir: cfg.promptDirIfSet(fs),
ProfileDir: cfg.profileDirIfSet(fs),
SchemaDir: cfg.schemaDirIfSet(fs),
ServerAddr: cfg.addrIfSet(fs),
ArtifactRoot: cfg.artifactRootIfSet(fs),
MaxRequestBytes: cfg.maxRequestBytesIfSet(fs),
MaxArtifactBytes: cfg.maxArtifactBytesIfSet(fs),
MaxResponseBytes: cfg.maxResponseBytesIfSet(fs),
})
if err != nil {
return nil, err
@@ -314,6 +329,9 @@ func parseServeArgs(args []string) (*serveConfig, error) {
cfg.schemaDir = settings.schemaDir
cfg.addr = settings.serverAddr
cfg.artifactRoot = settings.artifactRoot
cfg.maxRequestBytes = settings.maxRequestBytes
cfg.maxArtifactBytes = settings.maxArtifactBytes
cfg.maxResponseBytes = settings.maxResponseBytes
if err := validateRequiredLibraryDirs(cfg.promptDir); err != nil {
return nil, err
@@ -450,6 +468,27 @@ func (c *serveConfig) artifactRootIfSet(fs *flag.FlagSet) string {
return ""
}
func (c *serveConfig) maxRequestBytesIfSet(fs *flag.FlagSet) *int64 {
if flagWasSet(fs, "max-request-bytes") {
return &c.maxRequestBytes
}
return nil
}
func (c *serveConfig) maxArtifactBytesIfSet(fs *flag.FlagSet) *int64 {
if flagWasSet(fs, "max-artifact-bytes") {
return &c.maxArtifactBytes
}
return nil
}
func (c *serveConfig) maxResponseBytesIfSet(fs *flag.FlagSet) *int64 {
if flagWasSet(fs, "max-response-bytes") {
return &c.maxResponseBytes
}
return nil
}
func registerConfigPathFlag(fs *flag.FlagSet, target *string) {
fs.StringVar(
target,
@@ -488,6 +527,9 @@ func resolveCommonSettings(fs *flag.FlagSet, configPath string, overrides appcon
schemaDir: settings.SchemaDir,
serverAddr: settings.ServerAddr,
artifactRoot: settings.ArtifactRoot,
maxRequestBytes: settings.MaxRequestBytes,
maxArtifactBytes: settings.MaxArtifactBytes,
maxResponseBytes: settings.MaxResponseBytes,
defaultRenderFormat: settings.DefaultRenderFormat,
}, nil
}

View File

@@ -436,12 +436,18 @@ schema_dir: ./from-config/schemas
server:
addr: 127.0.0.1:9000
artifact_root: ./from-config/artifacts
max_request_bytes: 1024
max_artifact_bytes: 2048
max_response_bytes: 4096
`)
cfg, err := parseServeArgs([]string{
"--config", configPath,
"--addr", ":7777",
"--artifact-root", "./from-cli/artifacts",
"--max-request-bytes", "0",
"--max-artifact-bytes", "8192",
"--max-response-bytes", "16384",
})
if err != nil {
t.Fatalf("expected valid args, got %v", err)
@@ -462,6 +468,15 @@ server:
if cfg.artifactRoot != filepath.Clean("./from-cli/artifacts") {
t.Fatalf("expected CLI artifact root override, got %q", cfg.artifactRoot)
}
if cfg.maxRequestBytes != 0 {
t.Fatalf("expected CLI max request bytes override, got %d", cfg.maxRequestBytes)
}
if cfg.maxArtifactBytes != 8192 {
t.Fatalf("expected CLI max artifact bytes override, got %d", cfg.maxArtifactBytes)
}
if cfg.maxResponseBytes != 16384 {
t.Fatalf("expected CLI max response bytes override, got %d", cfg.maxResponseBytes)
}
}
func TestParseServeArgsWithConfigProvidesRequiredDirectoriesAndAddr(t *testing.T) {
@@ -472,6 +487,9 @@ schema_dir: ./from-config/schemas
server:
addr: 127.0.0.1:9000
artifact_root: ./from-config/artifacts
max_request_bytes: 1024
max_artifact_bytes: 2048
max_response_bytes: 4096
`)
cfg, err := parseServeArgs([]string{
@@ -496,6 +514,72 @@ server:
if cfg.artifactRoot != filepath.Clean("./from-config/artifacts") {
t.Fatalf("expected artifact root from config, got %q", cfg.artifactRoot)
}
if cfg.maxRequestBytes != 1024 {
t.Fatalf("expected max request bytes from config, got %d", cfg.maxRequestBytes)
}
if cfg.maxArtifactBytes != 2048 {
t.Fatalf("expected max artifact bytes from config, got %d", cfg.maxArtifactBytes)
}
if cfg.maxResponseBytes != 4096 {
t.Fatalf("expected max response bytes from config, got %d", cfg.maxResponseBytes)
}
}
func TestParseServeArgsRejectsNegativeSizeLimits(t *testing.T) {
tests := []struct {
name string
flag string
}{
{name: "request", flag: "--max-request-bytes"},
{name: "artifact", flag: "--max-artifact-bytes"},
{name: "response", flag: "--max-response-bytes"},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
_, err := parseServeArgs([]string{
"--prompt-dir", "./prompts",
tc.flag, "-1",
})
if err == nil {
t.Fatal("expected negative size limit error")
}
})
}
}
func TestRunAndRenderRejectServeSizeLimitFlags(t *testing.T) {
for _, tc := range []struct {
name string
parse func([]string) error
}{
{
name: "run",
parse: func(args []string) error {
_, err := parseRunArgs(args)
return err
},
},
{
name: "render",
parse: func(args []string) error {
_, err := parseRenderArgs(args)
return err
},
},
} {
t.Run(tc.name, func(t *testing.T) {
err := tc.parse([]string{
"--prompt-dir", "./prompts",
"--prompt", "p",
"--input", "a=b",
"--max-request-bytes", "1024",
})
if err == nil {
t.Fatal("expected unsupported flag error")
}
})
}
}
func TestRunAndRenderBuildEquivalentRuntimeOverrideRequestsForSharedFlags(t *testing.T) {