fix: close workspace preprocessing contract gaps

This commit is contained in:
2026-08-11 01:18:29 +02:00
parent 05a6e8cc2d
commit de27275bb1
15 changed files with 516 additions and 87 deletions
@@ -29,8 +29,11 @@ const (
maxAssumptions = 256
maxAssumptionBytes = 256
maxResult = 1 << 20
maxCandidate = 700 << 10
maxAnnotations = 16 << 20
// Requests carry up to 16 MiB of raw SQL or annotation bytes encoded as base64.
// This is deliberately independent from maxResult, which bounds child stdout.
maxRequest = 24 << 20
maxCandidate = 700 << 10
maxAnnotations = 16 << 20
)
var workspaceIDPattern = regexp.MustCompile(`^[a-z][a-z0-9-]{2,62}$`)
@@ -392,7 +395,7 @@ func makeInput(c Command) (inputEnvelope, string, error) {
if e != nil {
return env, "", e
}
if len(payload) > maxResult {
if len(payload) > maxRequest {
return env, "", errors.New("request exceeds limit")
}
return env, string(payload), nil
@@ -461,8 +464,8 @@ func Run(ctx context.Context, installation config.Installation, runner compose.R
}
payload := []byte(generated)
if stdin != nil {
payload, e = io.ReadAll(io.LimitReader(stdin, maxResult+1))
if e != nil || len(payload) > maxResult {
payload, e = io.ReadAll(io.LimitReader(stdin, maxRequest+1))
if e != nil || len(payload) > maxRequest {
return Result{}, errors.New("request exceeds limit")
}
if e = validateIngress(payload, env); e != nil {
@@ -83,9 +83,10 @@ func TestRunUsesBase64BasenameIngressAndNoTTY(t *testing.T) {
if err := os.WriteFile(sqlPath, []byte("select 1"), 0o600); err != nil {
t.Fatal(err)
}
out := filepath.Join(root, "candidate.yaml")
stdinCapture := filepath.Join(root, "stdin.json")
argsCapture := filepath.Join(root, "args.txt")
resultJSON := `{"schemaVersion":1,"status":"succeeded","code":"ok","workspaceId":"psd","workspaceRevision":"0123456789012345678901234567890123456789","descriptorBlob":"abcdefabcdefabcdefabcdefabcdefabcdefabcd","operation":"suggest-fks","runId":"0123456789abcdef0123456789abcdef","completedStages":[]}`
script := "#!/bin/sh\ncat >/dev/null\nprintf '%s' '" + resultJSON + "'\n"
script := "#!/bin/sh\nprintf '%s\n' \"$@\" > '" + argsCapture + "'\ncat > '" + stdinCapture + "'\nprintf '%s' '" + resultJSON + "'\n"
fake := filepath.Join(root, "docker")
if err := os.WriteFile(fake, []byte(script), 0o700); err != nil {
t.Fatal(err)
@@ -101,7 +102,20 @@ func TestRunUsesBase64BasenameIngressAndNoTTY(t *testing.T) {
if got.Code != "ok" {
t.Fatalf("result = %#v", got)
}
_ = out
stdinBytes, err := os.ReadFile(stdinCapture)
if err != nil {
t.Fatal(err)
}
if strings.Contains(string(stdinBytes), sqlPath) || !strings.Contains(string(stdinBytes), "contentBase64") || !strings.Contains(string(stdinBytes), "schema.sql") {
t.Fatalf("stdin envelope = %q; want basename/base64 without host path", stdinBytes)
}
argsBytes, err := os.ReadFile(argsCapture)
if err != nil {
t.Fatal(err)
}
if strings.Contains(string(argsBytes), sqlPath) || !strings.Contains(string(argsBytes), "--operation\nsuggest-fks") || !strings.Contains(string(argsBytes), "--workspace\npsd") {
t.Fatalf("child args = %q; want closed operation/workspace args", argsBytes)
}
}
func TestRunPublishesOnlyVerifiedCandidateExport(t *testing.T) {
@@ -295,3 +309,43 @@ func TestParseWorkspaceAcceptsAssumptionAndEnforcesLimits(t *testing.T) {
t.Fatal("accepted 33 SQL files")
}
}
func TestMakeInputAcceptsMaximumSingleSQLFile(t *testing.T) {
root, err := filepath.EvalSymlinks(t.TempDir())
if err != nil {
t.Fatal(err)
}
path := filepath.Join(root, "schema.sql")
if err := os.WriteFile(path, bytes.Repeat([]byte("x"), maxSQLFile), 0o600); err != nil {
t.Fatal(err)
}
command := SuggestFksRequest{WorkspaceID: "psd", FromSQL: []string{path}}
if _, payload, err := makeInput(command); err != nil {
t.Fatalf("makeInput() error = %v", err)
} else if len(payload) <= maxResult {
t.Fatalf("payload length = %d, want larger than result cap %d", len(payload), maxResult)
}
}
func TestMakeInputAcceptsMaximumAnnotationWithReviewDigestPair(t *testing.T) {
root, err := filepath.EvalSymlinks(t.TempDir())
if err != nil {
t.Fatal(err)
}
path := filepath.Join(root, "annotations.md")
contents := bytes.Repeat([]byte("a"), maxAnnotations)
if err := os.WriteFile(path, contents, 0o600); err != nil {
t.Fatal(err)
}
command := CheckSchemaRequest{WorkspaceID: "psd", Resume: strings.Repeat("0", 32), Annotations: path, ReviewedCandidates: DigestBytes([]byte("reviewed"))}
env, payload, err := makeInput(command)
if err != nil {
t.Fatal(err)
}
if env.Annotations == nil || env.ReviewedCandidates == "" || len(payload) <= maxResult {
t.Fatalf("annotation envelope = %#v payload=%d; want pair and request larger than result cap", env, len(payload))
}
if err := validateIngress([]byte(payload), env); err != nil {
t.Fatalf("validateIngress() = %v", err)
}
}