Add HTTP size limits
This commit is contained in:
@@ -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
|
||||
}
|
||||
|
||||
@@ -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) {
|
||||
|
||||
Reference in New Issue
Block a user