From 506028e15b2f5e515670cda3307e449044d138fa Mon Sep 17 00:00:00 2001 From: Joshua Gilman Date: Tue, 25 Aug 2026 12:51:16 -0700 Subject: [PATCH] feat(worker): bound intermediate native results --- internal/worker/child.go | 6 +- internal/worker/child_test.go | 39 ++++++++++- internal/worker/frame.go | 67 +++++++++++++------ internal/worker/frame_test.go | 51 +++++++++++++- internal/worker/parent.go | 9 ++- internal/worker/parent_test.go | 118 +++++++++++++++++++++++++++++++-- internal/worker/value.go | 14 ++-- internal/worker/value_test.go | 26 ++++++++ limits.go | 47 ++++++++----- limits_test.go | 6 ++ server.go | 15 +++-- server_test.go | 27 ++++++++ 12 files changed, 359 insertions(+), 66 deletions(-) diff --git a/internal/worker/child.go b/internal/worker/child.go index 6eaecef..f5d8c15 100644 --- a/internal/worker/child.go +++ b/internal/worker/child.go @@ -139,10 +139,10 @@ func executionLimits(limits childLimits) execution.Limits { // nativeForwarder sends one native_call and waits for the parent reply. func nativeForwarder(conn *childConn) execution.NativeCall { return func(id string, arguments map[string]any) (any, error) { - if err := validateBoundedValue(arguments, conn.limits); err != nil { - return nil, execution.ErrResourceLimit - } if err := conn.writeNativeCall(id, arguments); err != nil { + if errors.Is(err, errInvalidValue) || errors.Is(err, errFrameTooLarge) { + return nil, execution.ErrResourceLimit + } return nil, fmt.Errorf("%w: %w", errChildService, err) } frame, err := conn.read() diff --git a/internal/worker/child_test.go b/internal/worker/child_test.go index 2d2cd92..288ebd1 100644 --- a/internal/worker/child_test.go +++ b/internal/worker/child_test.go @@ -8,6 +8,8 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + + "github.com/meigma/codemode/internal/execution" ) // TestServeProbe proves one probe request is acknowledged and serve returns nil. @@ -222,6 +224,41 @@ func TestServeFinalWriteClassification(t *testing.T) { assert.False(t, ok) } +// TestNativeForwarderClassifiesBoundedWriteErrors proves oversized or nil +// argument maps fail as a resource limit before a native_call is written. +func TestNativeForwarderClassifiesBoundedWriteErrors(t *testing.T) { + tests := []struct { + // name identifies the rejected argument map. + name string + + // arguments are presented to the native forwarder. + arguments map[string]any + }{ + {name: "nil arguments", arguments: nil}, + {name: "oversized arguments", arguments: map[string]any{"org": "meigma"}}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + exec := validExecFrame() + exec.Limits.MaxValueBytes = 8 + parent, child := newBufferedExecPair(t, exec) + require.NoError(t, parent.writeExec(exec)) + _, err := child.read() + require.NoError(t, err) + written, ok := parent.r.(*bytes.Buffer) + require.True(t, ok) + require.Zero(t, written.Len()) + + _, err = nativeForwarder(child)("cap.lookup", tt.arguments) + + require.ErrorIs(t, err, execution.ErrResourceLimit) + require.NotErrorIs(t, err, errChildService) + assert.Zero(t, written.Len()) + }) + } +} + // TestServeEngineAbortSuppressesFinalError proves a real Engine native_abort // exits successfully without a terminal child frame. func TestServeEngineAbortSuppressesFinalError(t *testing.T) { @@ -250,7 +287,7 @@ func startServeExec(t *testing.T, exec execFrame) (*parentConn, <-chan error, fu childReader, parentWriter := io.Pipe() parentReader, childWriter := io.Pipe() - parent := newParentExecConn(parentWriter, parentReader, parentCap, childCap) + parent := newParentExecConn(parentWriter, parentReader, parentCap, childCap, exec.Limits.MaxValueBytes) done := make(chan error, 1) go func() { done <- serve(childReader, childWriter) diff --git a/internal/worker/frame.go b/internal/worker/frame.go index 607ea21..86ac1d9 100644 --- a/internal/worker/frame.go +++ b/internal/worker/frame.go @@ -207,6 +207,9 @@ type parentConn struct { // hasLimits reports whether exec limits have been stored. hasLimits bool + + // remainingIntermediate is the unused native-result value-body budget. + remainingIntermediate int } // childConn is the child side of one probe or execution exchange. @@ -316,9 +319,6 @@ func encodeExec(frame execFrame) ([]byte, error) { // encodeNativeCall encodes one child-originated native invocation. func encodeNativeCall(capabilityID string, arguments map[string]any) ([]byte, error) { - if capabilityID == "" { - return nil, errInvalidValue - } if arguments == nil { return nil, errInvalidValue } @@ -326,6 +326,14 @@ func encodeNativeCall(capabilityID string, arguments map[string]any) ([]byte, er if err != nil { return nil, err } + return encodeNativeCallBytes(capabilityID, encodedArgs) +} + +// encodeNativeCallBytes wraps one already-encoded argument body in a native_call frame. +func encodeNativeCallBytes(capabilityID string, encodedArgs []byte) ([]byte, error) { + if capabilityID == "" { + return nil, errInvalidValue + } encodedID, err := json.Marshal(capabilityID) if err != nil { return nil, errInvalidValue @@ -346,12 +354,17 @@ func encodeNativeResult(result any) ([]byte, error) { if err != nil { return nil, err } + return encodeNativeResultBytes(encoded), nil +} + +// encodeNativeResultBytes wraps one already-encoded value body in a native_result frame. +func encodeNativeResultBytes(encoded []byte) []byte { var buf bytes.Buffer buf.Grow(len(nativeResultPrefix) + len(encoded) + len(nativeResultSuffix)) buf.WriteString(nativeResultPrefix) buf.Write(encoded) buf.WriteString(nativeResultSuffix) - return buf.Bytes(), nil + return buf.Bytes() } // encodeNativeAbort encodes the payload-free parent abort frame. @@ -365,12 +378,17 @@ func encodeFinal(result any) ([]byte, error) { if err != nil { return nil, err } + return encodeFinalBytes(encoded), nil +} + +// encodeFinalBytes wraps one already-encoded value body in a final frame. +func encodeFinalBytes(encoded []byte) []byte { var buf bytes.Buffer buf.Grow(len(finalPrefix) + len(encoded) + len(finalSuffix)) buf.WriteString(finalPrefix) buf.Write(encoded) buf.WriteString(finalSuffix) - return buf.Bytes(), nil + return buf.Bytes() } // encodeFinalError encodes one child-owned terminal failure. @@ -459,14 +477,15 @@ func newChildProbeConn(r io.Reader, w io.Writer) *childConn { } // newParentExecConn constructs the parent side of an execution exchange. -func newParentExecConn(w io.Writer, r io.Reader, writeCap uint32, readCap uint32) *parentConn { +func newParentExecConn(w io.Writer, r io.Reader, writeCap uint32, readCap uint32, intermediateBudget int) *parentConn { return &parentConn{ - w: w, - r: r, - writeCap: writeCap, - readCap: readCap, - kind: connKindExec, - state: stateInit, + w: w, + r: r, + writeCap: writeCap, + readCap: readCap, + kind: connKindExec, + state: stateInit, + remainingIntermediate: intermediateBudget, } } @@ -527,16 +546,18 @@ func (c *parentConn) writeNativeResult(result any) error { if c == nil || c.kind != connKindExec || c.state != stateAwaitNative || !c.hasLimits { return errIllegalState } - if err := validateBoundedValue(result, c.limits); err != nil { - return err - } - payload, err := encodeNativeResult(result) + encoded, err := encodeBoundedValue(result, c.limits) if err != nil { return err } + if len(encoded) > c.remainingIntermediate { + return errFrameTooLarge + } + payload := encodeNativeResultBytes(encoded) if err := writeFrame(c.w, payload, c.writeCap); err != nil { return err } + c.remainingIntermediate -= len(encoded) c.state = stateReady return nil } @@ -654,10 +675,14 @@ func (c *childConn) writeNativeCall(capabilityID string, arguments map[string]an if c.nativeCalls >= c.limits.MaxNativeCalls { return errIllegalState } - if err := validateBoundedValue(arguments, c.limits); err != nil { + if arguments == nil { + return errInvalidValue + } + encodedArgs, err := encodeBoundedValue(arguments, c.limits) + if err != nil { return err } - payload, err := encodeNativeCall(capabilityID, arguments) + payload, err := encodeNativeCallBytes(capabilityID, encodedArgs) if err != nil { return err } @@ -674,13 +699,11 @@ func (c *childConn) writeFinal(result any) error { if c == nil || c.kind != connKindExec || c.state != stateReady || !c.hasLimits { return errIllegalState } - if err := validateBoundedValue(result, c.limits); err != nil { - return err - } - payload, err := encodeFinal(result) + encoded, err := encodeBoundedValue(result, c.limits) if err != nil { return err } + payload := encodeFinalBytes(encoded) if err := writeFrame(c.w, payload, c.writeCap); err != nil { return err } diff --git a/internal/worker/frame_test.go b/internal/worker/frame_test.go index 447dbe7..236ec5e 100644 --- a/internal/worker/frame_test.go +++ b/internal/worker/frame_test.go @@ -696,6 +696,53 @@ func TestFrameLimitsCheckedCaps(t *testing.T) { assert.Greater(t, expanded, unexpanded) } +// TestFrameNativeResultChargesValueBodyOnly proves the envelope is excluded from accounting. +func TestFrameNativeResultChargesValueBodyOnly(t *testing.T) { + exec := validExecFrame() + result := map[string]any{"zeta": int64(1), "alpha": true} + body, err := encodeNormalizedValue(result) + require.NoError(t, err) + frame, err := encodeNativeResult(result) + require.NoError(t, err) + assert.Equal(t, encodeNativeResultBytes(body), frame) + assert.Len(t, frame, len(nativeResultPrefix)+len(body)+len(nativeResultSuffix)) + assert.Greater(t, len(frame), len(body)) + + parent, child := newBufferedExecPair(t, exec) + parent.remainingIntermediate = len(body) + require.NoError(t, parent.writeExec(exec)) + _, err = child.read() + require.NoError(t, err) + require.NoError(t, child.writeNativeCall("cap.lookup", map[string]any{"org": "meigma"})) + _, err = parent.read() + require.NoError(t, err) + require.NoError(t, parent.writeNativeResult(result)) + assert.Zero(t, parent.remainingIntermediate) + reply, err := child.read() + require.NoError(t, err) + got, ok := reply.(nativeResultFrame) + require.True(t, ok) + assert.Equal(t, result, got.Result) + require.NoError(t, child.writeNativeCall("cap.lookup", map[string]any{"org": "meigma"})) + _, err = parent.read() + require.NoError(t, err) + err = parent.writeNativeResult(result) + require.ErrorIs(t, err, errFrameTooLarge) +} + +// TestWriteNativeCallRejectsNilArguments proves a nil map never encodes as JSON null. +func TestWriteNativeCallRejectsNilArguments(t *testing.T) { + exec := validExecFrame() + parent, child := newBufferedExecPair(t, exec) + require.NoError(t, parent.writeExec(exec)) + _, err := child.read() + require.NoError(t, err) + + err = child.writeNativeCall("cap.lookup", nil) + + require.ErrorIs(t, err, errInvalidValue) +} + // TestFrameLimitsRejectUint32Overflow proves checked addition never wraps the prefix. func TestFrameLimitsRejectUint32Overflow(t *testing.T) { manifest := validManifest() @@ -805,7 +852,7 @@ func newExecPair(t *testing.T, exec execFrame) (*parentConn, *childConn, func()) _ = childWriter.Close() } t.Cleanup(closePipes) - return newParentExecConn(parentWriter, parentReader, parentCap, childCap), + return newParentExecConn(parentWriter, parentReader, parentCap, childCap, exec.Limits.MaxValueBytes), newChildExecConn(childReader, childWriter, execCap, parentCap, childCap), closePipes } @@ -832,7 +879,7 @@ func newBufferedExecPair(t *testing.T, exec execFrame) (*parentConn, *childConn) var toChild bytes.Buffer var toParent bytes.Buffer - return newParentExecConn(&toChild, &toParent, parentCap, childCap), + return newParentExecConn(&toChild, &toParent, parentCap, childCap, exec.Limits.MaxValueBytes), newChildExecConn(&toChild, &toParent, execCap, parentCap, childCap) } diff --git a/internal/worker/parent.go b/internal/worker/parent.go index f4da869..38346c4 100644 --- a/internal/worker/parent.go +++ b/internal/worker/parent.go @@ -103,6 +103,11 @@ type Limits struct { // MaxValueBytes bounds type-preserving encoded value size. MaxValueBytes int + // MaxIntermediateValueBytes bounds cumulative successful native-result + // value-body bytes for one Execute call. It is independent of MaxValueBytes + // and is never projected to the child. + MaxIntermediateValueBytes int + // MaxConcurrentExecutions bounds live worker children. MaxConcurrentExecutions int } @@ -200,7 +205,7 @@ func NewRunner(bindings []execution.CapabilityBinding, limits Limits, dispatch D if dispatch == nil { return nil, errors.New("nil dispatch") } - if limits.MaxExecutionTime <= 0 || limits.MaxConcurrentExecutions <= 0 { + if limits.MaxExecutionTime <= 0 || limits.MaxConcurrentExecutions <= 0 || limits.MaxIntermediateValueBytes <= 0 { return nil, errInvalidLimits } child := childLimits{ @@ -669,7 +674,7 @@ func (r *Runner) runExecExchange( subject authz.Subject, source string, ) execOutcome { - conn := newParentExecConn(stdin, stdout, r.writeCap, r.readCap) + conn := newParentExecConn(stdin, stdout, r.writeCap, r.readCap, r.limits.MaxIntermediateValueBytes) frame := execFrame{ Type: frameTypeExec, Version: protocolVersion, diff --git a/internal/worker/parent_test.go b/internal/worker/parent_test.go index 0e76157..844ded6 100644 --- a/internal/worker/parent_test.go +++ b/internal/worker/parent_test.go @@ -34,13 +34,14 @@ const ( // testLimits returns positive budgets used by runner tests. func testLimits() Limits { return Limits{ - MaxSourceBytes: 64 * 1024, - MaxExecutionSteps: 1_000_000, - MaxExecutionTime: 5 * time.Second, - MaxNativeCalls: 100, - MaxValueDepth: 32, - MaxValueBytes: 1024 * 1024, - MaxConcurrentExecutions: 2, + MaxSourceBytes: 64 * 1024, + MaxExecutionSteps: 1_000_000, + MaxExecutionTime: 5 * time.Second, + MaxNativeCalls: 100, + MaxValueDepth: 32, + MaxValueBytes: 1024 * 1024, + MaxIntermediateValueBytes: 1024 * 1024, + MaxConcurrentExecutions: 2, } } @@ -239,6 +240,7 @@ func TestRunnerNewValidatesBudgetsAndManifest(t *testing.T) { {name: "zero time", bindings: nil, limits: withTime(valid, 0), dispatch: dispatch}, {name: "zero concurrency", bindings: nil, limits: withConcurrency(valid, 0), dispatch: dispatch}, {name: "zero source", bindings: nil, limits: withSource(valid, 0), dispatch: dispatch}, + {name: "zero intermediate", bindings: nil, limits: withIntermediate(valid, 0), dispatch: dispatch}, { name: "empty id", bindings: []execution.CapabilityBinding{withID(lookupBinding(), "")}, @@ -262,6 +264,102 @@ func TestRunnerNewValidatesBudgetsAndManifest(t *testing.T) { } } +// TestRunnerIntermediateBudget proves one Execute owns a fresh exact-body budget. +func TestRunnerIntermediateBudget(t *testing.T) { + const twoCallSource = "def main():\n records.lookup(value=\"a\")\n return records.lookup(value=\"b\")\n" + body, err := encodeNormalizedValue("xx") + require.NoError(t, err) + result := "xx" + + t.Run("below limit", func(t *testing.T) { + limits := testLimits() + limits.MaxIntermediateValueBytes = len(body)*2 + 1 + var calls int + runner := newTestRunner(t, limits, func(context.Context, authz.Subject, string, map[string]any) (any, error) { + calls++ + return result, nil + }) + + got, err := runner.Execute(context.Background(), authz.Subject{ID: "s"}, twoCallSource) + + require.NoError(t, err) + assert.Equal(t, result, got) + assert.Equal(t, 2, calls) + }) + + t.Run("inclusive edge", func(t *testing.T) { + limits := testLimits() + limits.MaxIntermediateValueBytes = len(body) * 2 + runner := newTestRunner(t, limits, func(context.Context, authz.Subject, string, map[string]any) (any, error) { + return result, nil + }) + + got, err := runner.Execute(context.Background(), authz.Subject{ID: "s"}, twoCallSource) + + require.NoError(t, err) + assert.Equal(t, result, got) + }) + + t.Run("over limit aborts before second delivery", func(t *testing.T) { + limits := testLimits() + limits.MaxIntermediateValueBytes = len(body)*2 - 1 + var calls int + runner := newTestRunner(t, limits, func(context.Context, authz.Subject, string, map[string]any) (any, error) { + calls++ + return result, nil + }) + + _, err := runner.Execute(context.Background(), authz.Subject{ID: "s"}, twoCallSource) + + require.ErrorIs(t, err, execution.ErrResourceLimit) + assert.Equal(t, 2, calls) + }) + + t.Run("fresh next execution", func(t *testing.T) { + limits := testLimits() + limits.MaxIntermediateValueBytes = len(body) + runner := newTestRunner(t, limits, func(context.Context, authz.Subject, string, map[string]any) (any, error) { + return result, nil + }) + + _, err := runner.Execute(context.Background(), authz.Subject{ID: "s"}, twoCallSource) + require.ErrorIs(t, err, execution.ErrResourceLimit) + + got, err := runner.Execute(context.Background(), authz.Subject{ID: "s"}, nativeSource) + require.NoError(t, err) + assert.Equal(t, result, got) + }) + + t.Run("concurrent executions stay isolated", func(t *testing.T) { + limits := testLimits() + limits.MaxIntermediateValueBytes = len(body) + limits.MaxConcurrentExecutions = 2 + runner := newTestRunner(t, limits, func(context.Context, authz.Subject, string, map[string]any) (any, error) { + return result, nil + }) + + var wg sync.WaitGroup + errCh := make(chan error, 2) + gotCh := make(chan any, 2) + for range 2 { + wg.Go(func() { + got, err := runner.Execute(context.Background(), authz.Subject{ID: "s"}, nativeSource) + errCh <- err + gotCh <- got + }) + } + wg.Wait() + close(errCh) + close(gotCh) + for err := range errCh { + require.NoError(t, err) + } + for got := range gotCh { + assert.Equal(t, result, got) + } + }) +} + // TestRunnerNewCopiesManifest proves caller slices are not retained. func TestRunnerNewCopiesManifest(t *testing.T) { capability := lookupBinding() @@ -669,6 +767,12 @@ func withSource(limits Limits, value int) Limits { return limits } +// withIntermediate returns a copy of limits with MaxIntermediateValueBytes replaced. +func withIntermediate(limits Limits, value int) Limits { + limits.MaxIntermediateValueBytes = value + return limits +} + // withID returns a copy of capability with ID replaced. func withID(capability execution.CapabilityBinding, id string) execution.CapabilityBinding { capability.ID = id diff --git a/internal/worker/value.go b/internal/worker/value.go index 0708822..79c8c10 100644 --- a/internal/worker/value.go +++ b/internal/worker/value.go @@ -181,15 +181,21 @@ func decodeNumber(number json.Number) (any, error) { // validateBoundedValue enforces the shared type, depth, and encoded-size limits. func validateBoundedValue(value any, limits childLimits) error { + _, err := encodeBoundedValue(value, limits) + return err +} + +// encodeBoundedValue validates one value and returns its canonical encoded body. +func encodeBoundedValue(value any, limits childLimits) ([]byte, error) { if err := binding.ValidateValue(value, limits.MaxValueDepth, limits.MaxValueBytes); err != nil { - return fmt.Errorf("%w: %w", errInvalidValue, err) + return nil, fmt.Errorf("%w: %w", errInvalidValue, err) } encoded, err := encodeNormalizedValue(value) if err != nil { - return err + return nil, err } if len(encoded) > limits.MaxValueBytes { - return errFrameTooLarge + return nil, errFrameTooLarge } - return nil + return encoded, nil } diff --git a/internal/worker/value_test.go b/internal/worker/value_test.go index c6504d3..8be8907 100644 --- a/internal/worker/value_test.go +++ b/internal/worker/value_test.go @@ -168,3 +168,29 @@ func TestValueCodecUsesBindingValidation(t *testing.T) { assert.Equal(t, value, decoded) require.NoError(t, binding.ValidateValue(decoded, 4, 1024)) } + +// TestValueCodecEncodesBoundedBodiesOnce proves validation reuses one encoded body. +func TestValueCodecEncodesBoundedBodiesOnce(t *testing.T) { + limits := childLimits{ + MaxSourceBytes: 256, + MaxExecutionSteps: 1, + MaxNativeCalls: 1, + MaxValueDepth: 8, + MaxValueBytes: 256, + } + value := map[string]any{"zeta": int64(1), "alpha": 1.0} + + encoded, err := encodeBoundedValue(value, limits) + require.NoError(t, err) + direct, err := encodeNormalizedValue(value) + require.NoError(t, err) + assert.Equal(t, direct, encoded) + assert.Equal(t, `{"alpha":1.0,"zeta":1}`, string(encoded)) + + require.NoError(t, validateBoundedValue(value, limits)) + tooSmall := limits + tooSmall.MaxValueBytes = 4 + _, err = encodeBoundedValue(value, tooSmall) + require.Error(t, err) + assert.True(t, errors.Is(err, errInvalidValue) || errors.Is(err, errFrameTooLarge)) +} diff --git a/limits.go b/limits.go index bc0e614..4736b8d 100644 --- a/limits.go +++ b/limits.go @@ -6,15 +6,16 @@ import ( ) const ( - defaultMaxSourceBytes = 64 * 1024 - defaultMaxExecutionSteps = 1_000_000 - defaultMaxExecutionTime = 5 * time.Second - defaultMaxNativeCalls = 100 - defaultMaxValueDepth = 32 - defaultMaxValueBytes = 1024 * 1024 - defaultMaxSearchQueryBytes = 256 - defaultMaxSearchResults = 20 - defaultMaxConcurrentExecutions = 8 + defaultMaxSourceBytes = 64 * 1024 + defaultMaxExecutionSteps = 1_000_000 + defaultMaxExecutionTime = 5 * time.Second + defaultMaxNativeCalls = 100 + defaultMaxValueDepth = 32 + defaultMaxValueBytes = 1024 * 1024 + defaultMaxIntermediateValueBytes = 8 * 1024 * 1024 + defaultMaxSearchQueryBytes = 256 + defaultMaxSearchResults = 20 + defaultMaxConcurrentExecutions = 8 ) // Limits bounds one execution and the model-facing catalog search surface. @@ -43,6 +44,13 @@ type Limits struct { // Size is measured by CodeMode's type-preserving JSON value encoder. MaxValueBytes int + // MaxIntermediateValueBytes is the maximum cumulative encoded size of + // successful parent-to-child native-result value bodies in one Execute + // call. Size is measured by CodeMode's type-preserving JSON value encoder + // and excludes frame envelopes, native-call arguments, failed handlers, + // and the final program value. The budget is independent of MaxValueBytes. + MaxIntermediateValueBytes int + // MaxSearchQueryBytes is the maximum capability-search query size in bytes. MaxSearchQueryBytes int @@ -58,15 +66,16 @@ type Limits struct { // DefaultLimits returns positive development defaults for every supported budget. func DefaultLimits() Limits { return Limits{ - MaxSourceBytes: defaultMaxSourceBytes, - MaxExecutionSteps: defaultMaxExecutionSteps, - MaxExecutionTime: defaultMaxExecutionTime, - MaxNativeCalls: defaultMaxNativeCalls, - MaxValueDepth: defaultMaxValueDepth, - MaxValueBytes: defaultMaxValueBytes, - MaxSearchQueryBytes: defaultMaxSearchQueryBytes, - MaxSearchResults: defaultMaxSearchResults, - MaxConcurrentExecutions: defaultMaxConcurrentExecutions, + MaxSourceBytes: defaultMaxSourceBytes, + MaxExecutionSteps: defaultMaxExecutionSteps, + MaxExecutionTime: defaultMaxExecutionTime, + MaxNativeCalls: defaultMaxNativeCalls, + MaxValueDepth: defaultMaxValueDepth, + MaxValueBytes: defaultMaxValueBytes, + MaxIntermediateValueBytes: defaultMaxIntermediateValueBytes, + MaxSearchQueryBytes: defaultMaxSearchQueryBytes, + MaxSearchResults: defaultMaxSearchResults, + MaxConcurrentExecutions: defaultMaxConcurrentExecutions, } } @@ -85,6 +94,8 @@ func (limits Limits) Validate() error { return fmt.Errorf("%w: MaxValueDepth must be positive", ErrInvalidRegistration) case limits.MaxValueBytes <= 0: return fmt.Errorf("%w: MaxValueBytes must be positive", ErrInvalidRegistration) + case limits.MaxIntermediateValueBytes <= 0: + return fmt.Errorf("%w: MaxIntermediateValueBytes must be positive", ErrInvalidRegistration) case limits.MaxSearchQueryBytes <= 0: return fmt.Errorf("%w: MaxSearchQueryBytes must be positive", ErrInvalidRegistration) case limits.MaxSearchResults <= 0: diff --git a/limits_test.go b/limits_test.go index af5b17a..01098fe 100644 --- a/limits_test.go +++ b/limits_test.go @@ -21,6 +21,7 @@ func TestDefaultLimitsAreValid(t *testing.T) { assert.Equal(t, uint64(100), limits.MaxNativeCalls) assert.Equal(t, 32, limits.MaxValueDepth) assert.Equal(t, 1_048_576, limits.MaxValueBytes) + assert.Equal(t, 8_388_608, limits.MaxIntermediateValueBytes) assert.Equal(t, 256, limits.MaxSearchQueryBytes) assert.Equal(t, 20, limits.MaxSearchResults) assert.Equal(t, 8, limits.MaxConcurrentExecutions) @@ -68,6 +69,11 @@ func TestLimitsRejectNonPositiveValues(t *testing.T) { mutate: func(limits *codemode.Limits) { limits.MaxValueBytes = 0 }, field: "MaxValueBytes", }, + { + name: "intermediate value bytes", + mutate: func(limits *codemode.Limits) { limits.MaxIntermediateValueBytes = 0 }, + field: "MaxIntermediateValueBytes", + }, { name: "search query bytes", mutate: func(limits *codemode.Limits) { limits.MaxSearchQueryBytes = 0 }, diff --git a/server.go b/server.go index 07fc957..70f6fd0 100644 --- a/server.go +++ b/server.go @@ -50,13 +50,14 @@ func newServer( runner, err := worker.NewRunner( capabilityBindings(capabilityCatalog), worker.Limits{ - MaxSourceBytes: limits.MaxSourceBytes, - MaxExecutionSteps: limits.MaxExecutionSteps, - MaxExecutionTime: limits.MaxExecutionTime, - MaxNativeCalls: limits.MaxNativeCalls, - MaxValueDepth: limits.MaxValueDepth, - MaxValueBytes: limits.MaxValueBytes, - MaxConcurrentExecutions: limits.MaxConcurrentExecutions, + MaxSourceBytes: limits.MaxSourceBytes, + MaxExecutionSteps: limits.MaxExecutionSteps, + MaxExecutionTime: limits.MaxExecutionTime, + MaxNativeCalls: limits.MaxNativeCalls, + MaxValueDepth: limits.MaxValueDepth, + MaxValueBytes: limits.MaxValueBytes, + MaxIntermediateValueBytes: limits.MaxIntermediateValueBytes, + MaxConcurrentExecutions: limits.MaxConcurrentExecutions, }, dispatch.dispatch, ) diff --git a/server_test.go b/server_test.go index ff33bfa..9fd4fb5 100644 --- a/server_test.go +++ b/server_test.go @@ -367,6 +367,33 @@ def main(): }) } +// TestServerExecuteEnforcesAggregateIntermediateLimit proves aggregate exhaustion +// projects only ErrResourceLimit and stays independent of MaxValueBytes. +func TestServerExecuteEnforcesAggregateIntermediateLimit(t *testing.T) { + const result = "xx" + body := `{"value":"` + result + `"}` + limits := codemode.DefaultLimits() + limits.MaxValueBytes = 64 + limits.MaxIntermediateValueBytes = len(body)*2 - 1 + var handlerCalls atomic.Int64 + capability := validBuilderCapability("cap.lookup", "records.lookup") + capability.Handler = func(_ context.Context, _ authz.Subject, _ builderInput) (builderOutput, error) { + handlerCalls.Add(1) + return builderOutput{Value: result}, nil + } + server := buildTestServer(t, authz.AllowAll(), limits, capability) + + _, err := server.Execute(t.Context(), authz.Subject{ID: "subject-1"}, ` +def main(): + records.lookup(value="first") + return records.lookup(value="second") +`) + + require.ErrorIs(t, err, codemode.ErrResourceLimit) + assert.Equal(t, codemode.ErrResourceLimit, err) + assert.Equal(t, int64(2), handlerCalls.Load()) +} + // TestServerExecuteReturnsOnlyMainResult proves top-level values and printed text do not escape. func TestServerExecuteReturnsOnlyMainResult(t *testing.T) { server := buildTestServer(