diff --git a/.mockery.yml b/.mockery.yml new file mode 100644 index 0000000..13b0083 --- /dev/null +++ b/.mockery.yml @@ -0,0 +1,16 @@ +all: false +dir: '{{.InterfaceDir}}/mocks' +filename: authorizer.go +formatter: goimports +generate: true +include-auto-generated: false +log-level: info +pkgname: mocks +require-template-schema-exists: true +structname: 'Mock{{.InterfaceName}}' +template: testify +template-schema: '{{.Template}}.schema.json' +packages: + github.com/meigma/codemode/authz: + interfaces: + Authorizer: diff --git a/authz/mocks/authorizer.go b/authz/mocks/authorizer.go new file mode 100644 index 0000000..ad7e8a8 --- /dev/null +++ b/authz/mocks/authorizer.go @@ -0,0 +1,96 @@ +// Code generated by mockery; DO NOT EDIT. +// github.com/vektra/mockery +// template: testify + +package mocks + +import ( + "context" + + "github.com/meigma/codemode/authz" + mock "github.com/stretchr/testify/mock" +) + +// NewMockAuthorizer creates a new instance of MockAuthorizer. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewMockAuthorizer(t interface { + mock.TestingT + Cleanup(func()) +}) *MockAuthorizer { + mock := &MockAuthorizer{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} + +// MockAuthorizer is an autogenerated mock type for the Authorizer type +type MockAuthorizer struct { + mock.Mock +} + +type MockAuthorizer_Expecter struct { + mock *mock.Mock +} + +func (_m *MockAuthorizer) EXPECT() *MockAuthorizer_Expecter { + return &MockAuthorizer_Expecter{mock: &_m.Mock} +} + +// Authorize provides a mock function for the type MockAuthorizer +func (_mock *MockAuthorizer) Authorize(context1 context.Context, authorizationInput authz.AuthorizationInput) error { + ret := _mock.Called(context1, authorizationInput) + + if len(ret) == 0 { + panic("no return value specified for Authorize") + } + + var r0 error + if returnFunc, ok := ret.Get(0).(func(context.Context, authz.AuthorizationInput) error); ok { + r0 = returnFunc(context1, authorizationInput) + } else { + r0 = ret.Error(0) + } + return r0 +} + +// MockAuthorizer_Authorize_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Authorize' +type MockAuthorizer_Authorize_Call struct { + *mock.Call +} + +// Authorize is a helper method to define mock.On call +// - context1 context.Context +// - authorizationInput authz.AuthorizationInput +func (_e *MockAuthorizer_Expecter) Authorize(context1 any, authorizationInput any) *MockAuthorizer_Authorize_Call { + return &MockAuthorizer_Authorize_Call{Call: _e.mock.On("Authorize", context1, authorizationInput)} +} + +func (_c *MockAuthorizer_Authorize_Call) Run(run func(context1 context.Context, authorizationInput authz.AuthorizationInput)) *MockAuthorizer_Authorize_Call { + _c.Call.Run(func(args mock.Arguments) { + var arg0 context.Context + if args[0] != nil { + arg0 = args[0].(context.Context) + } + var arg1 authz.AuthorizationInput + if args[1] != nil { + arg1 = args[1].(authz.AuthorizationInput) + } + run( + arg0, + arg1, + ) + }) + return _c +} + +func (_c *MockAuthorizer_Authorize_Call) Return(err error) *MockAuthorizer_Authorize_Call { + _c.Call.Return(err) + return _c +} + +func (_c *MockAuthorizer_Authorize_Call) RunAndReturn(run func(context1 context.Context, authorizationInput authz.AuthorizationInput) error) *MockAuthorizer_Authorize_Call { + _c.Call.Return(run) + return _c +} diff --git a/authz/mocks/doc.go b/authz/mocks/doc.go new file mode 100644 index 0000000..1f0745a --- /dev/null +++ b/authz/mocks/doc.go @@ -0,0 +1,2 @@ +// Package mocks contains generated test doubles for authorization ports. +package mocks diff --git a/builder.go b/builder.go new file mode 100644 index 0000000..1b057f7 --- /dev/null +++ b/builder.go @@ -0,0 +1,149 @@ +package codemode + +import ( + "context" + "fmt" + "reflect" + "slices" + + "github.com/meigma/codemode/authz" + "github.com/meigma/codemode/internal/binding" + "github.com/meigma/codemode/internal/catalog" +) + +// Options configures one immutable CodeMode server build. +type Options struct { + // Authorizer decides whether each validated native capability call may dispatch. + Authorizer authz.Authorizer + + // DisabledCapabilities lists stable capability IDs removed from every live server surface. + DisabledCapabilities []CapabilityID + + // Limits contains positive execution, conversion, and discovery budgets. + Limits Limits +} + +// Builder collects capability registrations for one immutable Server. +// +// A Builder is single-threaded and one-shot. Its first Build call closes registration even when +// validation fails. Construct another Builder to change configuration or capability visibility. +type Builder struct { + // authorizer is the required authorization port retained by the eventual server. + authorizer authz.Authorizer + + // disabledCapabilities is a private copy of the static deployment filter. + disabledCapabilities []string + + // limits is the value-copied server budget configuration. + limits Limits + + // registrations contains validated, precompiled capability registrations until Build. + registrations []catalog.Registration + + // built reports whether Build has closed this builder. + built bool +} + +// New creates a mutable one-shot Builder and copies caller-owned option slices. +func New(options Options) *Builder { + disabled := make([]string, len(options.DisabledCapabilities)) + for index, id := range options.DisabledCapabilities { + disabled[index] = string(id) + } + return &Builder{ + authorizer: options.Authorizer, + disabledCapabilities: disabled, + limits: options.Limits, + } +} + +// Register compiles and retains one typed capability without erasing its binding contract first. +func Register[Input, Output any](builder *Builder, capability Capability[Input, Output]) error { + if builder == nil { + return fmt.Errorf("%w: nil builder", ErrInvalidRegistration) + } + if builder.built { + return fmt.Errorf("%w: builder is already closed", ErrInvalidRegistration) + } + if capability.Handler == nil { + return fmt.Errorf("%w: capability handler must not be nil", ErrInvalidRegistration) + } + + plan, err := binding.CompileFor[Input, Output]() + if err != nil { + return fmt.Errorf("%w: capability %q: %w", ErrInvalidRegistration, capability.Name, err) + } + handler := capability.Handler + registration := catalog.Registration{ + ID: string(capability.ID), + Name: string(capability.Name), + Summary: capability.Summary, + Description: capability.Description, + Plan: plan, + Invoke: func(ctx context.Context, subject authz.Subject, input any) (any, error) { + typed, ok := input.(Input) + if !ok { + return nil, catalog.ErrInputTypeMismatch + } + return handler(ctx, subject, typed) + }, + } + if err := catalog.ValidateRegistration(registration); err != nil { + return fmt.Errorf("%w: %w", ErrInvalidRegistration, err) + } + for _, existing := range builder.registrations { + if existing.ID == registration.ID { + return fmt.Errorf("%w: duplicate capability ID %q", ErrInvalidRegistration, registration.ID) + } + if existing.Name == registration.Name { + return fmt.Errorf("%w: duplicate capability name %q", ErrInvalidRegistration, registration.Name) + } + } + builder.registrations = append(builder.registrations, registration) + return nil +} + +// Build closes the Builder and returns an immutable concurrency-safe Server after full validation. +func (builder *Builder) Build() (*Server, error) { + if builder == nil { + return nil, fmt.Errorf("%w: nil builder", ErrInvalidRegistration) + } + if builder.built { + return nil, fmt.Errorf("%w: builder is already closed", ErrInvalidRegistration) + } + builder.built = true + registrations := slices.Clone(builder.registrations) + builder.registrations = nil + + if isNilAuthorizer(builder.authorizer) { + return nil, fmt.Errorf("%w: Authorizer must not be nil", ErrInvalidRegistration) + } + if err := builder.limits.Validate(); err != nil { + return nil, err + } + capabilityCatalog, err := catalog.Build(registrations, catalog.Options{ + DisabledCapabilities: slices.Clone(builder.disabledCapabilities), + MaxSearchQueryBytes: builder.limits.MaxSearchQueryBytes, + MaxSearchResults: builder.limits.MaxSearchResults, + }) + if err != nil { + return nil, fmt.Errorf("%w: %w", ErrInvalidRegistration, err) + } + server, err := newServer(capabilityCatalog, builder.authorizer, builder.limits) + if err != nil { + return nil, fmt.Errorf("%w: %w", ErrInvalidRegistration, err) + } + return server, nil +} + +// isNilAuthorizer reports whether an authorization interface is nil or contains a typed nil value. +func isNilAuthorizer(authorizer authz.Authorizer) bool { + if authorizer == nil { + return true + } + value := reflect.ValueOf(authorizer) + kind := value.Kind() + nilable := kind == reflect.Chan || kind == reflect.Func || kind == reflect.Map || + kind == reflect.Pointer || kind == reflect.Slice + return nilable && value.IsNil() +} diff --git a/builder_test.go b/builder_test.go new file mode 100644 index 0000000..33a491a --- /dev/null +++ b/builder_test.go @@ -0,0 +1,191 @@ +package codemode_test + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/meigma/codemode" + "github.com/meigma/codemode/authz" +) + +// builderInput is the representative builder test input. +type builderInput struct { + // Value is one required capability argument. + Value string `json:"value"` +} + +// builderOutput is the representative builder test output. +type builderOutput struct { + // Value is one handler result field. + Value string `json:"value"` +} + +// invalidBuilderInput contains a type outside the restricted binder contract. +type invalidBuilderInput struct { + // Count is intentionally unsupported as a required integer. + Count int64 `json:"count"` +} + +// nilPolicy is a typed-nil authorization implementation used to test Build validation. +type nilPolicy struct{} + +// Authorize permits an invocation when the receiver is non-nil. +func (*nilPolicy) Authorize(context.Context, authz.AuthorizationInput) error { + return nil +} + +// TestBuilderBuildsOnceAndClosesRegistration proves the mutable lifecycle ends at the first Build call. +func TestBuilderBuildsOnceAndClosesRegistration(t *testing.T) { + builder := codemode.New(codemode.Options{ + Authorizer: authz.AllowAll(), + Limits: codemode.DefaultLimits(), + }) + require.NoError(t, codemode.Register(builder, validBuilderCapability("cap.one", "records.one"))) + + server, err := builder.Build() + + require.NoError(t, err) + require.NotNil(t, server) + _, err = builder.Build() + require.ErrorIs(t, err, codemode.ErrInvalidRegistration) + err = codemode.Register(builder, validBuilderCapability("cap.two", "records.two")) + require.ErrorIs(t, err, codemode.ErrInvalidRegistration) +} + +// TestBuilderCopiesStaticFilteringOptions proves caller slice mutation cannot alter retained configuration. +func TestBuilderCopiesStaticFilteringOptions(t *testing.T) { + disabled := []codemode.CapabilityID{"cap.disabled"} + builder := codemode.New(codemode.Options{ + Authorizer: authz.AllowAll(), + DisabledCapabilities: disabled, + Limits: codemode.DefaultLimits(), + }) + disabled[0] = "cap.unknown" + require.NoError(t, codemode.Register(builder, validBuilderCapability("cap.disabled", "records.disabled"))) + + server, err := builder.Build() + + require.NoError(t, err) + results, err := server.Search("disabled") + require.NoError(t, err) + assert.Empty(t, results) +} + +// TestBuilderRejectsNilAuthorizersAndInvalidLimits proves required construction policy fails closed. +func TestBuilderRejectsNilAuthorizersAndInvalidLimits(t *testing.T) { + var typedNil *nilPolicy + tests := []struct { + // name identifies the invalid server options. + name string + + // options contains the construction input. + options codemode.Options + }{ + {name: "nil authorizer", options: codemode.Options{Limits: codemode.DefaultLimits()}}, + { + name: "typed nil authorizer", + options: codemode.Options{Authorizer: typedNil, Limits: codemode.DefaultLimits()}, + }, + {name: "zero limits", options: codemode.Options{Authorizer: authz.AllowAll()}}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + builder := codemode.New(tt.options) + + server, err := builder.Build() + + require.ErrorIs(t, err, codemode.ErrInvalidRegistration) + assert.Nil(t, server) + _, secondErr := builder.Build() + require.ErrorIs(t, secondErr, codemode.ErrInvalidRegistration) + }) + } +} + +// TestRegisterRejectsInvalidContractsBeforeRetention proves malformed capabilities never enter the builder. +func TestRegisterRejectsInvalidContractsBeforeRetention(t *testing.T) { + builder := codemode.New(codemode.Options{Authorizer: authz.AllowAll(), Limits: codemode.DefaultLimits()}) + invalidMetadata := validBuilderCapability("", "records.valid") + nilHandler := validBuilderCapability("cap.nil", "records.nil") + nilHandler.Handler = nil + + require.ErrorIs(t, codemode.Register(builder, invalidMetadata), codemode.ErrInvalidRegistration) + require.ErrorIs(t, codemode.Register(builder, nilHandler), codemode.ErrInvalidRegistration) + require.ErrorIs(t, codemode.Register(builder, codemode.Capability[invalidBuilderInput, builderOutput]{ + ID: "cap.unsupported", + Name: "records.unsupported", + Summary: "Unsupported input.", + Description: "Exercises compile-before-erasure validation.", + Handler: func(context.Context, authz.Subject, invalidBuilderInput) (builderOutput, error) { + return builderOutput{}, nil + }, + }), codemode.ErrInvalidRegistration) + require.ErrorIs(t, codemode.Register[builderInput, builderOutput](nil, validBuilderCapability( + "cap.nil_builder", + "records.nil_builder", + )), codemode.ErrInvalidRegistration) +} + +// TestRegisterRejectsObviousDuplicates proves duplicate stable and model-facing identities fail immediately. +func TestRegisterRejectsObviousDuplicates(t *testing.T) { + builder := codemode.New(codemode.Options{Authorizer: authz.AllowAll(), Limits: codemode.DefaultLimits()}) + require.NoError(t, codemode.Register(builder, validBuilderCapability("cap.one", "records.one"))) + + require.ErrorIs( + t, + codemode.Register(builder, validBuilderCapability("cap.one", "records.two")), + codemode.ErrInvalidRegistration, + ) + require.ErrorIs( + t, + codemode.Register(builder, validBuilderCapability("cap.two", "records.one")), + codemode.ErrInvalidRegistration, + ) +} + +// TestBuildRejectsWholeCatalogFailures proves filtering and namespace checks run before a server escapes. +func TestBuildRejectsWholeCatalogFailures(t *testing.T) { + t.Run("unknown disabled capability", func(t *testing.T) { + builder := codemode.New(codemode.Options{ + Authorizer: authz.AllowAll(), + DisabledCapabilities: []codemode.CapabilityID{"cap.unknown"}, + Limits: codemode.DefaultLimits(), + }) + + server, err := builder.Build() + + require.ErrorIs(t, err, codemode.ErrInvalidRegistration) + assert.Nil(t, server) + }) + + t.Run("function namespace collision", func(t *testing.T) { + builder := codemode.New(codemode.Options{Authorizer: authz.AllowAll(), Limits: codemode.DefaultLimits()}) + require.NoError(t, codemode.Register(builder, validBuilderCapability("cap.lookup", "records.lookup"))) + require.NoError(t, codemode.Register(builder, validBuilderCapability("cap.detail", "records.lookup.detail"))) + + server, err := builder.Build() + + require.ErrorIs(t, err, codemode.ErrInvalidRegistration) + assert.Nil(t, server) + }) +} + +// validBuilderCapability constructs one valid representative registration. +func validBuilderCapability( + id codemode.CapabilityID, + name codemode.CapabilityName, +) codemode.Capability[builderInput, builderOutput] { + return codemode.Capability[builderInput, builderOutput]{ + ID: id, + Name: name, + Summary: "Return one record.", + Description: "Returns the supplied record value.", + Handler: func(_ context.Context, _ authz.Subject, input builderInput) (builderOutput, error) { + return builderOutput(input), nil + }, + } +} diff --git a/doc.go b/doc.go index 2336313..ee95ce1 100644 --- a/doc.go +++ b/doc.go @@ -1,2 +1,3 @@ -// Package codemode defines capability and runtime configuration contracts for code-native MCP servers. +// Package codemode builds immutable catalogs of typed Go capabilities and executes +// bounded, authorized Starlark programs against them. package codemode diff --git a/go.mod b/go.mod index cf94f1b..1a35ef5 100644 --- a/go.mod +++ b/go.mod @@ -8,6 +8,7 @@ require ( ) require ( + github.com/stretchr/objx v0.5.3 // indirect go.yaml.in/yaml/v3 v3.0.5 // indirect golang.org/x/sys v0.42.0 // indirect ) diff --git a/go.sum b/go.sum index 7259203..64966d2 100644 --- a/go.sum +++ b/go.sum @@ -1,5 +1,7 @@ github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= +github.com/stretchr/objx v0.5.3 h1:jmXUvGomnU1o3W/V5h2VEradbpJDwGrzugQQvL0POH4= +github.com/stretchr/objx v0.5.3/go.mod h1:rDQraq+vQZU7Fde9LOZLr8Tax6zZvy4kuNKF+QYS+U0= github.com/stretchr/testify v1.12.1 h1:EuwCh5fleGS7H32xRwO3wRGT7DxrDhLAT6FF8MpWDWE= github.com/stretchr/testify v1.12.1/go.mod h1:MDEgiDPPsNp5cuIrHPPCyornHKgEVbtFUmoNlxoYthg= go.starlark.net v0.0.0-20260708150628-5395d018f003 h1:cAxcqHgW8fnmT0cEBU3TzvVYHIFt8IIGDMWUF6rImk4= diff --git a/internal/catalog/build.go b/internal/catalog/build.go index d06c5d6..c7f0cd4 100644 --- a/internal/catalog/build.go +++ b/internal/catalog/build.go @@ -17,11 +17,14 @@ var ( // ErrUnknownDisabledCapability classifies a disabled ID absent from the validated registration set. ErrUnknownDisabledCapability = errors.New("unknown disabled capability") + + // ErrInputTypeMismatch classifies an erased handler argument that cannot be the compiled input type. + ErrInputTypeMismatch = errors.New("capability input type mismatch") ) -// candidate is one fully validated and compiled registration before static filtering. +// candidate is one fully validated registration before static filtering. type candidate struct { - // entry contains copied metadata, the compiled plan, and the handler adapter. + // entry contains copied metadata, the supplied plan, and the handler adapter. entry Entry } @@ -38,7 +41,7 @@ func Build(registrations []Registration, options Options) (*Catalog, error) { ids := make(map[string]struct{}, len(registrations)) names := make(map[string]struct{}, len(registrations)) for index, registration := range registrations { - if err := validateRegistration(registration); err != nil { + if err := ValidateRegistration(registration); err != nil { return nil, fmt.Errorf("%w: registration %d: %w", ErrInvalidRegistration, index, err) } if _, duplicate := ids[registration.ID]; duplicate { @@ -47,10 +50,6 @@ func Build(registrations []Registration, options Options) (*Catalog, error) { if _, duplicate := names[registration.Name]; duplicate { return nil, fmt.Errorf("%w: duplicate name %q", ErrInvalidRegistration, registration.Name) } - plan, err := binding.Compile(registration.InputType, registration.OutputType) - if err != nil { - return nil, fmt.Errorf("%w: capability %q: %w", ErrInvalidRegistration, registration.Name, err) - } ids[registration.ID] = struct{}{} names[registration.Name] = struct{}{} candidates = append(candidates, candidate{entry: Entry{ @@ -58,13 +57,13 @@ func Build(registrations []Registration, options Options) (*Catalog, error) { Name: registration.Name, Summary: registration.Summary, Description: registration.Description, - Plan: plan, + Plan: registration.Plan, Invoke: registration.Invoke, - signature: plan.Signature(registration.Name), + signature: registration.Plan.Signature(registration.Name), searchName: strings.ToLower(registration.Name), searchSummary: strings.ToLower(registration.Summary), - inputShape: plan.InputShape(), - outputShape: plan.OutputShape(), + inputShape: registration.Plan.InputShape(), + outputShape: registration.Plan.OutputShape(), }}) } if err := validateNamespaceCollisions(names); err != nil { @@ -107,8 +106,8 @@ func Build(registrations []Registration, options Options) (*Catalog, error) { return catalog, nil } -// validateRegistration checks metadata and handler presence before any filtering occurs. -func validateRegistration(registration Registration) error { +// ValidateRegistration reports whether one copied registration has valid metadata, a compiled plan, and a handler. +func ValidateRegistration(registration Registration) error { if registration.ID == "" || registration.ID != strings.TrimSpace(registration.ID) { return errors.New("ID must be non-empty without surrounding whitespace") } @@ -121,6 +120,9 @@ func validateRegistration(registration Registration) error { if registration.Description == "" || registration.Description != strings.TrimSpace(registration.Description) { return errors.New("description must be non-empty without surrounding whitespace") } + if registration.Plan == nil { + return errors.New("plan must not be nil") + } if registration.Invoke == nil { return errors.New("handler must not be nil") } diff --git a/internal/catalog/catalog.go b/internal/catalog/catalog.go index bfc5ee6..973125a 100644 --- a/internal/catalog/catalog.go +++ b/internal/catalog/catalog.go @@ -2,7 +2,6 @@ package catalog import ( "context" - "reflect" "slices" "github.com/meigma/codemode/authz" @@ -12,7 +11,7 @@ import ( // Invoker is a type-erased adapter around one registered typed capability handler. type Invoker func(context.Context, authz.Subject, any) (any, error) -// Registration contains one capability before validation, compilation, and static filtering. +// Registration contains one capability before validation and static filtering. type Registration struct { // ID is the stable deployment and policy identity. ID string @@ -26,11 +25,8 @@ type Registration struct { // Description is the full exact-description text. Description string - // InputType is the exact typed handler input. - InputType reflect.Type - - // OutputType is the exact typed handler output. - OutputType reflect.Type + // Plan is the caller-compiled input, output, canonical-argument, and signature plan. + Plan *binding.Plan // Invoke calls the registered typed handler through its registration-time adapter. Invoke Invoker diff --git a/internal/catalog/catalog_test.go b/internal/catalog/catalog_test.go index 7b78afb..615e388 100644 --- a/internal/catalog/catalog_test.go +++ b/internal/catalog/catalog_test.go @@ -2,7 +2,6 @@ package catalog import ( "context" - "reflect" "sync" "testing" @@ -10,6 +9,7 @@ import ( "github.com/stretchr/testify/require" "github.com/meigma/codemode/authz" + "github.com/meigma/codemode/internal/binding" ) // testInput is the catalog test capability input. @@ -30,13 +30,13 @@ type testOutput struct { // TestBuildValidatesEveryRegistrationBeforeFiltering proves disabled entries cannot hide invalid contracts. func TestBuildValidatesEveryRegistrationBeforeFiltering(t *testing.T) { registration := validRegistration("cap.bad", "records.bad", "Bad record") - registration.InputType = reflect.TypeFor[int]() + registration.Plan = nil _, err := Build([]Registration{registration}, testOptions("cap.bad")) require.Error(t, err) require.ErrorIs(t, err, ErrInvalidRegistration) - assert.Contains(t, err.Error(), "binding plan") + assert.Contains(t, err.Error(), "plan") } // TestBuildRejectsInvalidMetadataDuplicatesAndNamespaceCollisions proves initialization fails closed. @@ -268,15 +268,14 @@ func TestCatalogSupportsConcurrentReadOnlyUse(t *testing.T) { assert.Equal(t, "records.alpha", catalog.Entries()[0].Name) } -// validRegistration constructs one valid test registration. +// validRegistration constructs one valid test registration from the once-compiled plan. func validRegistration(id string, name string, summary string) Registration { return Registration{ ID: id, Name: name, Summary: summary, Description: summary + " full description.", - InputType: reflect.TypeFor[testInput](), - OutputType: reflect.TypeFor[testOutput](), + Plan: mustCompileTestPlan(), Invoke: func(context.Context, authz.Subject, any) (any, error) { return testOutput{Name: name}, nil }, @@ -321,3 +320,12 @@ func withDescription(registration Registration, description string) Registration registration.Description = description return registration } + +// mustCompileTestPlan compiles the catalog test input and output types. +func mustCompileTestPlan() *binding.Plan { + plan, err := binding.CompileFor[testInput, testOutput]() + if err != nil { + panic(err) + } + return plan +} diff --git a/internal/execution/doc.go b/internal/execution/doc.go new file mode 100644 index 0000000..9fb2379 --- /dev/null +++ b/internal/execution/doc.go @@ -0,0 +1,2 @@ +// Package execution runs one restricted, authorized, bounded Starlark program at a time. +package execution diff --git a/internal/execution/engine.go b/internal/execution/engine.go new file mode 100644 index 0000000..5e308df --- /dev/null +++ b/internal/execution/engine.go @@ -0,0 +1,27 @@ +package execution + +import ( + "fmt" + + "go.starlark.net/starlark" + + "github.com/meigma/codemode/internal/catalog" +) + +// Engine owns the frozen capability namespace shared by fresh execution threads. +type Engine struct { + // predeclared contains immutable namespaced capability builtins. + predeclared starlark.StringDict +} + +// New compiles one immutable execution engine from a validated filtered catalog. +func New(capabilityCatalog *catalog.Catalog) (*Engine, error) { + if capabilityCatalog == nil { + return nil, fmt.Errorf("%w: nil catalog", ErrInternal) + } + predeclared, err := buildPredeclared(capabilityCatalog) + if err != nil { + return nil, err + } + return &Engine{predeclared: predeclared}, nil +} diff --git a/internal/execution/errors.go b/internal/execution/errors.go new file mode 100644 index 0000000..4de1ea5 --- /dev/null +++ b/internal/execution/errors.go @@ -0,0 +1,29 @@ +package execution + +import "errors" + +var ( + // ErrInvalidProgram classifies invalid source, entrypoint behavior, or loading-phase native calls. + ErrInvalidProgram = errors.New("invalid program") + + // ErrInvalidArguments classifies native arguments rejected before authorization. + ErrInvalidArguments = errors.New("invalid capability arguments") + + // ErrPermissionDenied classifies a recognized authorization denial. + ErrPermissionDenied = errors.New("permission denied") + + // ErrPolicyFailure classifies an authorization evaluation error or panic. + ErrPolicyFailure = errors.New("authorization policy failure") + + // ErrResourceLimit classifies source, step, time, call, depth, or result-size exhaustion. + ErrResourceLimit = errors.New("resource limit exceeded") + + // ErrCapabilityFailure classifies a native handler or handler-result failure. + ErrCapabilityFailure = errors.New("capability failed") + + // ErrInternal classifies an unexpected runtime or contract failure. + ErrInternal = errors.New("internal failure") + + // errRecoveredHandlerPanic distinguishes a recovered panic from an ordinary handler error. + errRecoveredHandlerPanic = errors.New("recovered handler panic") +) diff --git a/internal/execution/execute.go b/internal/execution/execute.go new file mode 100644 index 0000000..9499c6f --- /dev/null +++ b/internal/execution/execute.go @@ -0,0 +1,325 @@ +package execution + +import ( + "context" + "errors" + "fmt" + + "go.starlark.net/starlark" + "go.starlark.net/syntax" + + "github.com/meigma/codemode/authz" + "github.com/meigma/codemode/internal/binding" + "github.com/meigma/codemode/internal/catalog" +) + +const executionStateLocal = "codemode.execution.state" + +// executionOutcome carries one recovered top-level result without named returns. +type executionOutcome struct { + // value is main's converted final value. + value any + + // err is the classified execution failure. + err error +} + +// Execute runs source in a fresh restricted interpreter and converts only main's final value. +func (engine *Engine) Execute( + ctx context.Context, + subject authz.Subject, + source string, + authorizer authz.Authorizer, + limits Limits, +) (any, error) { + outcome := executionOutcome{} + func() { + defer func() { + if recover() != nil { + outcome.value = nil + outcome.err = ErrInternal + } + }() + outcome.value, outcome.err = execute( + ctx, + subject, + source, + engine.predeclared, + authorizer, + limits, + ) + }() + return outcome.value, outcome.err +} + +// execute performs one execution after the panic-recovery boundary is installed. +func execute( + ctx context.Context, + subject authz.Subject, + source string, + predeclared starlark.StringDict, + authorizer authz.Authorizer, + limits Limits, +) (any, error) { + if ctx == nil || predeclared == nil || authorizer == nil { + return nil, ErrInternal + } + if len(source) > limits.MaxSourceBytes { + return nil, ErrResourceLimit + } + if contextErr := contextFailure(ctx); contextErr != nil { + return nil, contextErr + } + + runCtx, cancel := context.WithTimeout(ctx, limits.MaxExecutionTime) + defer cancel() + if contextErr := contextFailure(runCtx); contextErr != nil { + return nil, contextErr + } + + state := &executionState{ + ctx: runCtx, + subject: subject, + authorizer: authorizer, + phase: phaseLoading, + nativeCalls: checkedCounter{ + maximum: limits.MaxNativeCalls, + }, + } + thread := newThread(state, limits.MaxExecutionSteps) + stopCancellation := watchCancellation(runCtx, thread) + defer stopCancellation() + + globals, executionErr := starlark.ExecFileOptions( + &syntax.FileOptions{}, + thread, + "", + source, + predeclared, + ) + if executionErr != nil { + return nil, classifyRuntimeError(runCtx, state, executionErr) + } + mainFunction, mainErr := requireMain(globals) + if mainErr != nil { + return nil, mainErr + } + if phaseErr := state.beginMain(); phaseErr != nil { + return nil, phaseErr + } + finalValue, callErr := starlark.Call(thread, mainFunction, nil, nil) + state.finishMain() + if callErr != nil { + return nil, classifyRuntimeError(runCtx, state, callErr) + } + if contextErr := contextFailure(runCtx); contextErr != nil { + return nil, contextErr + } + + converted, conversionErr := binding.ConvertFinal( + finalValue, + limits.MaxValueDepth, + limits.MaxResultBytes, + ) + if conversionErr != nil { + if errors.Is(conversionErr, binding.ErrValueLimit) { + return nil, fmt.Errorf("%w: %w", ErrResourceLimit, conversionErr) + } + return nil, fmt.Errorf("%w: %w", ErrInvalidProgram, conversionErr) + } + if contextErr := contextFailure(runCtx); contextErr != nil { + return nil, contextErr + } + return converted, nil +} + +// newThread constructs one restricted Starlark thread and stores its trusted execution state. +func newThread(state *executionState, maxSteps uint64) *starlark.Thread { + thread := &starlark.Thread{ + Name: "codemode", + Print: func(*starlark.Thread, string) { + // Program output is intentionally discarded; only main's final value may escape. + }, + Load: func(*starlark.Thread, string) (starlark.StringDict, error) { + return nil, fmt.Errorf("%w: module loading is disabled", ErrInvalidProgram) + }, + } + thread.SetLocal(executionStateLocal, state) + thread.OnMaxSteps = func(thread *starlark.Thread) { + state.stepLimited = true + thread.Cancel("execution step limit exceeded") + } + thread.SetMaxExecutionSteps(maxSteps) + return thread +} + +// watchCancellation interrupts Starlark evaluation when the request or elapsed deadline ends. +func watchCancellation(ctx context.Context, thread *starlark.Thread) func() { + stopped := make(chan struct{}) + go func() { + select { + case <-ctx.Done(): + thread.Cancel("execution canceled") + case <-stopped: + } + }() + return func() { + close(stopped) + } +} + +// requireMain returns the exact zero-argument main function required by the execution contract. +func requireMain(globals starlark.StringDict) (*starlark.Function, error) { + value, ok := globals["main"] + if !ok { + return nil, fmt.Errorf("%w: main is missing", ErrInvalidProgram) + } + mainFunction, ok := value.(*starlark.Function) + if !ok { + return nil, fmt.Errorf("%w: main must be a function", ErrInvalidProgram) + } + if mainFunction.NumParams() != 0 || mainFunction.NumKwonlyParams() != 0 || + mainFunction.HasVarargs() || mainFunction.HasKwargs() { + return nil, fmt.Errorf("%w: main must accept no arguments", ErrInvalidProgram) + } + return mainFunction, nil +} + +// callCapability validates, authorizes, dispatches, and converts one native invocation in order. +func callCapability( + thread *starlark.Thread, + entry catalog.Entry, + args starlark.Tuple, + kwargs []starlark.Tuple, +) (starlark.Value, error) { + state, ok := thread.Local(executionStateLocal).(*executionState) + if !ok || state == nil { + return nil, ErrInternal + } + if state.phase != phaseRunning { + return nil, ErrInvalidProgram + } + if counterErr := state.nativeCalls.increment(); counterErr != nil { + return nil, counterErr + } + input, canonical, bindingErr := entry.Plan.Bind(args, kwargs) + if bindingErr != nil { + if errors.Is(bindingErr, binding.ErrInvalidArguments) { + return nil, fmt.Errorf("%w: %w", ErrInvalidArguments, bindingErr) + } + return nil, fmt.Errorf("%w: %w", ErrInternal, bindingErr) + } + if contextErr := contextFailure(state.ctx); contextErr != nil { + return nil, contextErr + } + if policyErr := authorize(state, entry, canonical); policyErr != nil { + return nil, policyErr + } + if contextErr := contextFailure(state.ctx); contextErr != nil { + return nil, contextErr + } + output, invocationErr := invoke(state, entry, input) + if invocationErr != nil { + return nil, invocationErr + } + converted, conversionErr := entry.Plan.ConvertOutput(output) + if conversionErr != nil { + return nil, fmt.Errorf("%w: %w", ErrCapabilityFailure, conversionErr) + } + return converted, nil +} + +// authorize calls policy once and converts denial, error, and panic to safe internal classes. +func authorize(state *executionState, entry catalog.Entry, arguments map[string]any) (err error) { + defer func() { + if recover() != nil { + err = ErrPolicyFailure + } + }() + err = state.authorizer.Authorize(state.ctx, authz.AuthorizationInput{ + Subject: state.subject, + CapabilityID: entry.ID, + CapabilityName: entry.Name, + Arguments: arguments, + }) + if err == nil { + return nil + } + if errors.Is(err, authz.ErrDenied) { + return fmt.Errorf("%w: %w", ErrPermissionDenied, err) + } + return fmt.Errorf("%w: %w", ErrPolicyFailure, err) +} + +// invocationOutcome carries one recovered native handler result without named returns. +type invocationOutcome struct { + // output is the handler's typed result. + output any + + // err is the handler or recovery failure. + err error +} + +// invoke calls one typed handler and recovers handler panics at the native boundary. +func invoke(state *executionState, entry catalog.Entry, input any) (any, error) { + outcome := invocationOutcome{} + func() { + defer func() { + if recover() != nil { + outcome.output = nil + outcome.err = errRecoveredHandlerPanic + } + }() + outcome.output, outcome.err = entry.Invoke(state.ctx, state.subject, input) + }() + if outcome.err == nil { + return outcome.output, nil + } + if errors.Is(outcome.err, errRecoveredHandlerPanic) { + return nil, ErrInternal + } + if errors.Is(outcome.err, catalog.ErrInputTypeMismatch) { + return nil, fmt.Errorf("%w: %w", ErrInternal, outcome.err) + } + return nil, fmt.Errorf("%w: %w", ErrCapabilityFailure, outcome.err) +} + +// classifyRuntimeError prefers request and budget state before classifying evaluator causes. +func classifyRuntimeError(ctx context.Context, state *executionState, err error) error { + if state.stepLimited { + return ErrResourceLimit + } + if contextErr := contextFailure(ctx); contextErr != nil { + return contextErr + } + switch { + case errors.Is(err, ErrInvalidArguments): + return ErrInvalidArguments + case errors.Is(err, ErrPermissionDenied): + return ErrPermissionDenied + case errors.Is(err, ErrPolicyFailure): + return ErrPolicyFailure + case errors.Is(err, ErrResourceLimit): + return ErrResourceLimit + case errors.Is(err, ErrCapabilityFailure): + return ErrCapabilityFailure + case errors.Is(err, ErrInternal): + return ErrInternal + case errors.Is(err, ErrInvalidProgram): + return ErrInvalidProgram + default: + return fmt.Errorf("%w: %w", ErrInvalidProgram, err) + } +} + +// contextFailure projects request cancellation directly and deadlines as resource exhaustion. +func contextFailure(ctx context.Context) error { + switch err := ctx.Err(); { + case errors.Is(err, context.Canceled): + return context.Canceled + case errors.Is(err, context.DeadlineExceeded): + return fmt.Errorf("%w: %w", ErrResourceLimit, context.DeadlineExceeded) + default: + return nil + } +} diff --git a/internal/execution/execute_test.go b/internal/execution/execute_test.go new file mode 100644 index 0000000..7981cc4 --- /dev/null +++ b/internal/execution/execute_test.go @@ -0,0 +1,450 @@ +package execution_test + +import ( + "context" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/mock" + "github.com/stretchr/testify/require" + + "github.com/meigma/codemode/authz" + authzmocks "github.com/meigma/codemode/authz/mocks" + "github.com/meigma/codemode/internal/binding" + "github.com/meigma/codemode/internal/catalog" + "github.com/meigma/codemode/internal/execution" +) + +// testInput is the representative native capability input. +type testInput struct { + // Value is one required capability argument. + Value string `json:"value"` +} + +// testOutput is the representative native capability output. +type testOutput struct { + // Value is one capability result field. + Value string `json:"value"` +} + +// testExecutionResult carries one asynchronous execution outcome. +type testExecutionResult struct { + // value is main's converted final value. + value any + + // err is the classified execution failure. + err error +} + +// TestExecuteRequiresExactMainAndRejectsLoadingCalls proves entrypoint and phase validation precede side effects. +func TestExecuteRequiresExactMainAndRejectsLoadingCalls(t *testing.T) { + var handlerCalls atomic.Int64 + capabilityCatalog := buildEngine(t, func(context.Context, authz.Subject, any) (any, error) { + handlerCalls.Add(1) + return testOutput{}, nil + }) + tests := []struct { + // name identifies the invalid program shape. + name string + + // source is the submitted Starlark program. + source string + }{ + {name: "missing main", source: `value = 1`}, + {name: "main is not function", source: `main = 1`}, + {name: "main has argument", source: `def main(value): return value`}, + {name: "main has variadic arguments", source: `def main(*args): return args`}, + {name: "top-level native call", source: ` +value = records.lookup(value="alpha") +def main(): + return value +`}, + {name: "module loading", source: ` +load("module.star", "value") +def main(): + return value +`}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, err := capabilityCatalog.Execute( + t.Context(), + authz.Subject{ID: "subject-1"}, + tt.source, + authz.AllowAll(), + defaultExecutionLimits(), + ) + + require.ErrorIs(t, err, execution.ErrInvalidProgram) + }) + } + assert.Zero(t, handlerCalls.Load()) +} + +// TestExecuteBindsAuthorizesThenDispatches proves canonical validation and policy ordering. +func TestExecuteBindsAuthorizesThenDispatches(t *testing.T) { + authorizer := authzmocks.NewMockAuthorizer(t) + events := make([]string, 0, 2) + authorizer.EXPECT().Authorize(mock.Anything, mock.MatchedBy(func(input authz.AuthorizationInput) bool { + return input.Subject.ID == "subject-1" && + input.CapabilityID == "cap.lookup" && + input.CapabilityName == "records.lookup" && + assert.ObjectsAreEqual(map[string]any{"value": "alpha"}, input.Arguments) + })).Run(func(context.Context, authz.AuthorizationInput) { + events = append(events, "authorize") + }).Return(nil).Once() + capabilityCatalog := buildEngine(t, func(_ context.Context, _ authz.Subject, input any) (any, error) { + events = append(events, "handler") + typed, ok := input.(testInput) + if !ok { + return nil, catalog.ErrInputTypeMismatch + } + return testOutput(typed), nil + }) + + result, err := capabilityCatalog.Execute( + t.Context(), + authz.Subject{ID: "subject-1"}, + `def main(): return records.lookup(value="alpha")`, + authorizer, + defaultExecutionLimits(), + ) + + require.NoError(t, err) + assert.Equal(t, []string{"authorize", "handler"}, events) + assert.Equal(t, map[string]any{"value": "alpha"}, result) +} + +// TestExecuteRejectsMalformedArgumentsBeforePolicy proves invalid calls reach neither policy nor handler. +func TestExecuteRejectsMalformedArgumentsBeforePolicy(t *testing.T) { + authorizer := authzmocks.NewMockAuthorizer(t) + var handlerCalls atomic.Int64 + capabilityCatalog := buildEngine(t, func(context.Context, authz.Subject, any) (any, error) { + handlerCalls.Add(1) + return testOutput{}, nil + }) + + _, err := capabilityCatalog.Execute( + t.Context(), + authz.Subject{ID: "subject-1"}, + `def main(): return records.lookup()`, + authorizer, + defaultExecutionLimits(), + ) + + require.ErrorIs(t, err, execution.ErrInvalidArguments) + assert.Zero(t, handlerCalls.Load()) +} + +// TestExecuteCancellationAfterAuthorizationPreventsDispatch proves a stale allow cannot cross the handler boundary. +func TestExecuteCancellationAfterAuthorizationPreventsDispatch(t *testing.T) { + tests := []struct { + // name identifies the cancellation source. + name string + + // newContext creates the execution context, cleanup, and cancellation trigger. + newContext func(*testing.T) (context.Context, context.CancelFunc, func()) + + // targets are the required execution error classifications. + targets []error + }{ + { + name: "request cancellation", + newContext: func(t *testing.T) (context.Context, context.CancelFunc, func()) { + ctx, cancel := context.WithCancel(t.Context()) + return ctx, cancel, cancel + }, + targets: []error{context.Canceled}, + }, + { + name: "elapsed deadline", + newContext: func(t *testing.T) (context.Context, context.CancelFunc, func()) { + ctx, cancel := context.WithTimeout(t.Context(), 100*time.Millisecond) + return ctx, cancel, func() { + <-ctx.Done() + } + }, + targets: []error{execution.ErrResourceLimit, context.DeadlineExceeded}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctx, cancel, trigger := tt.newContext(t) + defer cancel() + authorizationStarted := make(chan struct{}) + releaseAuthorization := make(chan struct{}) + authorizer := authzmocks.NewMockAuthorizer(t) + authorizer.EXPECT().Authorize(mock.Anything, mock.Anything).Run( + func(context.Context, authz.AuthorizationInput) { + close(authorizationStarted) + <-releaseAuthorization + }, + ).Return(nil).Once() + var handlerCalls atomic.Int64 + engine := buildEngine(t, func(context.Context, authz.Subject, any) (any, error) { + handlerCalls.Add(1) + return testOutput{}, nil + }) + result := make(chan testExecutionResult, 1) + go func() { + value, err := engine.Execute( + ctx, + authz.Subject{ID: "subject-1"}, + `def main(): return records.lookup(value="alpha")`, + authorizer, + defaultExecutionLimits(), + ) + result <- testExecutionResult{value: value, err: err} + }() + + <-authorizationStarted + trigger() + close(releaseAuthorization) + outcome := <-result + + assert.Nil(t, outcome.value) + for _, target := range tt.targets { + require.ErrorIs(t, outcome.err, target) + } + assert.Zero(t, handlerCalls.Load()) + }) + } +} + +// TestExecuteCancelsInFlightStarlark proves the watcher interrupts evaluation after execution begins. +func TestExecuteCancelsInFlightStarlark(t *testing.T) { + handlerStarted := make(chan struct{}) + releaseHandler := make(chan struct{}) + handlerReturned := make(chan struct{}) + engine := buildEngine(t, func(context.Context, authz.Subject, any) (any, error) { + close(handlerStarted) + <-releaseHandler + close(handlerReturned) + return testOutput{}, nil + }) + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + limits := defaultExecutionLimits() + limits.MaxExecutionSteps = ^uint64(0) + result := make(chan testExecutionResult, 1) + go func() { + value, err := engine.Execute( + ctx, + authz.Subject{ID: "subject-1"}, + ` +def main(): + records.lookup(value="alpha") + total = 0 + for item in range(1000000000): + total += item + return total +`, + authz.AllowAll(), + limits, + ) + result <- testExecutionResult{value: value, err: err} + }() + + <-handlerStarted + close(releaseHandler) + <-handlerReturned + cancel() + outcome := <-result + + assert.Nil(t, outcome.value) + require.ErrorIs(t, outcome.err, context.Canceled) +} + +// TestExecuteCreatesFreshStateAndReturnsOnlyMain proves executions share no globals or counters. +func TestExecuteCreatesFreshStateAndReturnsOnlyMain(t *testing.T) { + capabilityCatalog := buildEngine(t, func(_ context.Context, _ authz.Subject, input any) (any, error) { + typed, ok := input.(testInput) + if !ok { + return nil, catalog.ErrInputTypeMismatch + } + return testOutput(typed), nil + }) + const source = ` +hidden = "discarded" +print(hidden) +def main(): + return records.lookup(value="visible") +` + limits := defaultExecutionLimits() + limits.MaxNativeCalls = 1 + + first, err := capabilityCatalog.Execute( + t.Context(), + authz.Subject{ID: "subject-1"}, + source, + authz.AllowAll(), + limits, + ) + require.NoError(t, err) + second, err := capabilityCatalog.Execute( + t.Context(), + authz.Subject{ID: "subject-1"}, + source, + authz.AllowAll(), + limits, + ) + + require.NoError(t, err) + assert.Equal(t, map[string]any{"value": "visible"}, first) + assert.Equal(t, map[string]any{"value": "visible"}, second) +} + +// TestExecuteEnforcesEveryRuntimeBudget proves source, steps, calls, time, depth, and bytes fail safely. +func TestExecuteEnforcesEveryRuntimeBudget(t *testing.T) { + var handlerCalls atomic.Int64 + capabilityCatalog := buildEngine(t, func(_ context.Context, _ authz.Subject, input any) (any, error) { + handlerCalls.Add(1) + typed, ok := input.(testInput) + if !ok { + return nil, catalog.ErrInputTypeMismatch + } + return testOutput(typed), nil + }) + + t.Run("source bytes", func(t *testing.T) { + limits := defaultExecutionLimits() + limits.MaxSourceBytes = 8 + _, err := capabilityCatalog.Execute( + t.Context(), + authz.Subject{ID: "subject-1"}, + `def main(): return None`, + authz.AllowAll(), + limits, + ) + require.ErrorIs(t, err, execution.ErrResourceLimit) + }) + + t.Run("execution steps", func(t *testing.T) { + limits := defaultExecutionLimits() + limits.MaxExecutionSteps = 20 + _, err := capabilityCatalog.Execute(t.Context(), authz.Subject{ID: "subject-1"}, ` +def main(): + total = 0 + for item in range(1000000): + total += item + return total +`, authz.AllowAll(), limits) + require.ErrorIs(t, err, execution.ErrResourceLimit) + }) + + t.Run("native calls", func(t *testing.T) { + handlerCalls.Store(0) + limits := defaultExecutionLimits() + limits.MaxNativeCalls = 1 + _, err := capabilityCatalog.Execute(t.Context(), authz.Subject{ID: "subject-1"}, ` +def main(): + records.lookup(value="first") + return records.lookup(value="second") +`, authz.AllowAll(), limits) + require.ErrorIs(t, err, execution.ErrResourceLimit) + assert.Equal(t, int64(1), handlerCalls.Load()) + }) + + t.Run("elapsed time", func(t *testing.T) { + limits := defaultExecutionLimits() + limits.MaxExecutionTime = time.Nanosecond + _, err := capabilityCatalog.Execute( + t.Context(), + authz.Subject{ID: "subject-1"}, + `def main(): return None`, + authz.AllowAll(), + limits, + ) + require.ErrorIs(t, err, execution.ErrResourceLimit) + require.ErrorIs(t, err, context.DeadlineExceeded) + }) + + t.Run("value depth", func(t *testing.T) { + limits := defaultExecutionLimits() + limits.MaxValueDepth = 2 + _, err := capabilityCatalog.Execute( + t.Context(), + authz.Subject{ID: "subject-1"}, + `def main(): return [[["deep"]]]`, + authz.AllowAll(), + limits, + ) + require.ErrorIs(t, err, execution.ErrResourceLimit) + }) + + t.Run("result bytes", func(t *testing.T) { + limits := defaultExecutionLimits() + limits.MaxResultBytes = 4 + _, err := capabilityCatalog.Execute( + t.Context(), + authz.Subject{ID: "subject-1"}, + `def main(): return "oversized"`, + authz.AllowAll(), + limits, + ) + require.ErrorIs(t, err, execution.ErrResourceLimit) + }) +} + +// TestExecutePreservesCancellationAndRejectsUnsupportedFinalValues proves terminal failures remain classified. +func TestExecutePreservesCancellationAndRejectsUnsupportedFinalValues(t *testing.T) { + capabilityCatalog := buildEngine(t, func(context.Context, authz.Subject, any) (any, error) { + return testOutput{}, nil + }) + ctx, cancel := context.WithCancel(t.Context()) + cancel() + + _, err := capabilityCatalog.Execute( + ctx, + authz.Subject{ID: "subject-1"}, + `def main(): return None`, + authz.AllowAll(), + defaultExecutionLimits(), + ) + require.ErrorIs(t, err, context.Canceled) + _, err = capabilityCatalog.Execute( + t.Context(), + authz.Subject{ID: "subject-1"}, + `def main(): return main`, + authz.AllowAll(), + defaultExecutionLimits(), + ) + require.ErrorIs(t, err, execution.ErrInvalidProgram) +} + +// buildEngine creates one immutable representative execution engine. +func buildEngine(t *testing.T, invoke catalog.Invoker) *execution.Engine { + t.Helper() + plan, err := binding.CompileFor[testInput, testOutput]() + require.NoError(t, err) + capabilityCatalog, err := catalog.Build([]catalog.Registration{ + { + ID: "cap.lookup", + Name: "records.lookup", + Summary: "Look up one record.", + Description: "Returns one record by value.", + Plan: plan, + Invoke: invoke, + }, + }, catalog.Options{MaxSearchQueryBytes: 1024, MaxSearchResults: 10}) + require.NoError(t, err) + engine, err := execution.New(capabilityCatalog) + require.NoError(t, err) + return engine +} + +// defaultExecutionLimits returns positive representative execution budgets. +func defaultExecutionLimits() execution.Limits { + return execution.Limits{ + MaxSourceBytes: 64 * 1024, + MaxExecutionSteps: 1_000_000, + MaxExecutionTime: 5 * time.Second, + MaxNativeCalls: 10, + MaxValueDepth: 16, + MaxResultBytes: 64 * 1024, + } +} diff --git a/internal/execution/limits.go b/internal/execution/limits.go new file mode 100644 index 0000000..45e248f --- /dev/null +++ b/internal/execution/limits.go @@ -0,0 +1,24 @@ +package execution + +import "time" + +// Limits contains the prevalidated positive budgets for one execution. +type Limits struct { + // MaxSourceBytes bounds the submitted Starlark source length. + MaxSourceBytes int + + // MaxExecutionSteps bounds abstract Starlark interpreter work. + MaxExecutionSteps uint64 + + // MaxExecutionTime bounds elapsed execution time. + MaxExecutionTime time.Duration + + // MaxNativeCalls bounds attempted capability invocations. + MaxNativeCalls uint64 + + // MaxValueDepth bounds recursive final-result conversion. + MaxValueDepth int + + // MaxResultBytes bounds encoded final-result size. + MaxResultBytes int +} diff --git a/internal/execution/namespace.go b/internal/execution/namespace.go new file mode 100644 index 0000000..6a3b0c1 --- /dev/null +++ b/internal/execution/namespace.go @@ -0,0 +1,75 @@ +package execution + +import ( + "fmt" + "maps" + + "go.starlark.net/starlark" + "go.starlark.net/starlarkstruct" + + "github.com/meigma/codemode/internal/catalog" +) + +// namespaceNode is one temporary registration-time path used to assemble frozen modules. +type namespaceNode struct { + // children contains nested namespace segments. + children map[string]*namespaceNode + + // functions contains native capability leaves. + functions starlark.StringDict +} + +// buildPredeclared creates the frozen capability namespace for one immutable Engine. +func buildPredeclared(capabilityCatalog *catalog.Catalog) (starlark.StringDict, error) { + root := newNamespaceNode() + for _, namespaceBinding := range capabilityCatalog.NamespaceBindings() { + node := root + for _, segment := range namespaceBinding.Segments { + child, ok := node.children[segment] + if !ok { + child = newNamespaceNode() + node.children[segment] = child + } + node = child + } + if _, duplicate := node.functions[namespaceBinding.Function]; duplicate { + return nil, fmt.Errorf("%w: duplicate namespace function", ErrInternal) + } + entry := namespaceBinding.Capability + node.functions[namespaceBinding.Function] = starlark.NewBuiltin(entry.Name, func( + thread *starlark.Thread, + _ *starlark.Builtin, + args starlark.Tuple, + kwargs []starlark.Tuple, + ) (starlark.Value, error) { + return callCapability(thread, entry, args, kwargs) + }) + } + + predeclared := make(starlark.StringDict, len(root.children)) + for name, node := range root.children { + predeclared[name] = freezeModule(name, node) + } + predeclared.Freeze() + return predeclared, nil +} + +// newNamespaceNode creates an empty mutable namespace assembly node. +func newNamespaceNode() *namespaceNode { + return &namespaceNode{ + children: make(map[string]*namespaceNode), + functions: make(starlark.StringDict), + } +} + +// freezeModule recursively converts one assembly node to an immutable Starlark module. +func freezeModule(name string, node *namespaceNode) *starlarkstruct.Module { + members := make(starlark.StringDict, len(node.children)+len(node.functions)) + maps.Copy(members, node.functions) + for segment, child := range node.children { + members[segment] = freezeModule(name+"."+segment, child) + } + module := &starlarkstruct.Module{Name: name, Members: members} + module.Freeze() + return module +} diff --git a/internal/execution/phase.go b/internal/execution/phase.go new file mode 100644 index 0000000..f2ecda7 --- /dev/null +++ b/internal/execution/phase.go @@ -0,0 +1,75 @@ +package execution + +import ( + "context" + "fmt" + + "github.com/meigma/codemode/authz" +) + +// executionPhase identifies whether native capability calls are currently permitted. +type executionPhase uint8 + +const ( + // phaseLoading executes top-level source while rejecting every native capability call. + phaseLoading executionPhase = iota + + // phaseRunning permits native capability calls from the validated main function. + phaseRunning + + // phaseDone rejects calls after the one main invocation completes. + phaseDone +) + +// checkedCounter bounds one monotonically increasing per-execution quantity. +type checkedCounter struct { + // current is the number of accepted increments. + current uint64 + + // maximum is the largest accepted count. + maximum uint64 +} + +// increment advances the counter or returns ErrResourceLimit before exceeding its maximum. +func (counter *checkedCounter) increment() error { + if counter.current >= counter.maximum { + return ErrResourceLimit + } + counter.current++ + return nil +} + +// executionState contains request-local trusted state unavailable to Starlark programs. +type executionState struct { + // ctx carries request cancellation and the derived elapsed deadline. + ctx context.Context + + // subject is the trusted authenticated identity for every native call. + subject authz.Subject + + // authorizer decides each validated invocation immediately before dispatch. + authorizer authz.Authorizer + + // phase guards the side-effectful native-call boundary. + phase executionPhase + + // nativeCalls bounds attempted capability invocations. + nativeCalls checkedCounter + + // stepLimited records interpreter cancellation caused by the configured step budget. + stepLimited bool +} + +// beginMain transitions exactly once from loading to running. +func (state *executionState) beginMain() error { + if state.phase != phaseLoading { + return fmt.Errorf("%w: invalid execution phase", ErrInternal) + } + state.phase = phaseRunning + return nil +} + +// finishMain closes native dispatch for this execution. +func (state *executionState) finishMain() { + state.phase = phaseDone +} diff --git a/server.go b/server.go new file mode 100644 index 0000000..a586227 --- /dev/null +++ b/server.go @@ -0,0 +1,141 @@ +package codemode + +import ( + "context" + "errors" + "fmt" + + "github.com/meigma/codemode/authz" + "github.com/meigma/codemode/internal/catalog" + "github.com/meigma/codemode/internal/execution" +) + +// Program is one bounded Starlark source program executed by a Server. +type Program string + +// SearchResult is one compact enabled-capability discovery record. +type SearchResult = catalog.SearchResult + +// Description is one exact enabled-capability description and supported binding shape. +type Description = catalog.Description + +// Server is an immutable, concurrency-safe capability catalog and Starlark execution service. +// +// Every Execute call owns a fresh interpreter and budgets. The elapsed budget cancels Starlark +// evaluation; registered Authorizer and Handler implementations must honor their context, return +// promptly, and be safe for the caller's concurrency. +type Server struct { + // catalog is the immutable statically filtered capability set. + catalog *catalog.Catalog + + // engine owns the immutable precompiled native capability namespace. + engine *execution.Engine + + // authorizer is called once before each valid native handler dispatch. + authorizer authz.Authorizer + + // limits is the immutable execution budget configuration. + limits execution.Limits +} + +// newServer compiles a Server from fully validated retained state. +func newServer( + capabilityCatalog *catalog.Catalog, + authorizer authz.Authorizer, + limits Limits, +) (*Server, error) { + engine, err := execution.New(capabilityCatalog) + if err != nil { + return nil, err + } + return &Server{ + catalog: capabilityCatalog, + engine: engine, + authorizer: authorizer, + limits: execution.Limits{ + MaxSourceBytes: limits.MaxSourceBytes, + MaxExecutionSteps: limits.MaxExecutionSteps, + MaxExecutionTime: limits.MaxExecutionTime, + MaxNativeCalls: limits.MaxNativeCalls, + MaxValueDepth: limits.MaxValueDepth, + MaxResultBytes: limits.MaxResultBytes, + }, + }, nil +} + +// Search returns a bounded name-sorted scan of enabled capability names and summaries. +func (server *Server) Search(query string) ([]SearchResult, error) { + if server == nil || server.catalog == nil { + return nil, ErrInternal + } + results, err := server.catalog.Search(query) + if err != nil { + if errors.Is(err, catalog.ErrSearchQueryLimit) { + return nil, ErrResourceLimit + } + return nil, ErrInternal + } + return results, nil +} + +// Describe returns one exact enabled capability description or ErrNotFound. +func (server *Server) Describe(name CapabilityName) (Description, error) { + if server == nil || server.catalog == nil { + return Description{}, ErrInternal + } + description, ok := server.catalog.Describe(string(name)) + if !ok { + return Description{}, ErrNotFound + } + return description, nil +} + +// Execute runs one bounded program for a trusted authenticated subject and returns only main's final value. +// +// The elapsed budget cancels Starlark evaluation. Authorizer and Handler implementations must +// honor ctx and return promptly; CodeMode does not detach or forcibly interrupt blocking Go code. +func (server *Server) Execute(ctx context.Context, subject authz.Subject, program Program) (any, error) { + if server == nil || server.engine == nil || server.authorizer == nil { + return nil, ErrInternal + } + if ctx == nil { + return nil, ErrInternal + } + if subject.ID == "" { + return nil, ErrUnauthenticated + } + result, err := server.engine.Execute(ctx, subject, string(program), server.authorizer, server.limits) + if err != nil { + return nil, projectExecutionError(err) + } + return result, nil +} + +// projectExecutionError converts trusted internal execution failures to the public safe taxonomy. +func projectExecutionError(err error) error { + switch { + case errors.Is(err, execution.ErrInvalidProgram): + return ErrInvalidProgram + case errors.Is(err, execution.ErrInvalidArguments): + return ErrInvalidArguments + case errors.Is(err, execution.ErrPermissionDenied): + return ErrPermissionDenied + case errors.Is(err, execution.ErrPolicyFailure): + return ErrPolicyFailure + case errors.Is(err, execution.ErrResourceLimit): + if errors.Is(err, context.DeadlineExceeded) { + return fmt.Errorf("%w: %w", ErrResourceLimit, context.DeadlineExceeded) + } + return ErrResourceLimit + case errors.Is(err, execution.ErrCapabilityFailure): + return ErrCapabilityFailure + case errors.Is(err, execution.ErrInternal): + return ErrInternal + case errors.Is(err, context.Canceled): + return context.Canceled + case errors.Is(err, context.DeadlineExceeded): + return fmt.Errorf("%w: %w", ErrResourceLimit, context.DeadlineExceeded) + default: + return ErrInternal + } +} diff --git a/server_test.go b/server_test.go new file mode 100644 index 0000000..f13bbb2 --- /dev/null +++ b/server_test.go @@ -0,0 +1,423 @@ +package codemode_test + +import ( + "context" + "errors" + "sync" + "sync/atomic" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/mock" + "github.com/stretchr/testify/require" + + "github.com/meigma/codemode" + "github.com/meigma/codemode/authz" + authzmocks "github.com/meigma/codemode/authz/mocks" +) + +// TestServerSearchAndDescribeExposeOnlyEnabledCapabilities proves public discovery uses the filtered catalog. +func TestServerSearchAndDescribeExposeOnlyEnabledCapabilities(t *testing.T) { + builder := codemode.New(codemode.Options{ + Authorizer: authz.AllowAll(), + DisabledCapabilities: []codemode.CapabilityID{"cap.disabled"}, + Limits: codemode.DefaultLimits(), + }) + require.NoError(t, codemode.Register(builder, validBuilderCapability("cap.alpha", "records.alpha"))) + require.NoError(t, codemode.Register(builder, validBuilderCapability("cap.disabled", "records.disabled"))) + server, err := builder.Build() + require.NoError(t, err) + + results, err := server.Search("record") + require.NoError(t, err) + require.Len(t, results, 1) + assert.Equal(t, "records.alpha", results[0].Name) + description, err := server.Describe("records.alpha") + require.NoError(t, err) + assert.Equal(t, "records.alpha", description.Name) + _, err = server.Describe("records.disabled") + require.ErrorIs(t, err, codemode.ErrNotFound) +} + +// TestServerSearchProjectsQueryLimits proves internal discovery details do not escape the public taxonomy. +func TestServerSearchProjectsQueryLimits(t *testing.T) { + limits := codemode.DefaultLimits() + limits.MaxSearchQueryBytes = 4 + builder := codemode.New(codemode.Options{Authorizer: authz.AllowAll(), Limits: limits}) + server, err := builder.Build() + require.NoError(t, err) + + _, err = server.Search("oversized") + + require.ErrorIs(t, err, codemode.ErrResourceLimit) + assert.Equal(t, codemode.ErrResourceLimit.Error(), err.Error()) +} + +// TestServerExecuteAuthorizesCanonicalArgumentsBeforeDispatch proves the complete native-call ordering. +func TestServerExecuteAuthorizesCanonicalArgumentsBeforeDispatch(t *testing.T) { + authorizer := authzmocks.NewMockAuthorizer(t) + events := make([]string, 0, 2) + authorizer.EXPECT().Authorize(mock.Anything, mock.MatchedBy(func(input authz.AuthorizationInput) bool { + return input.Subject.ID == "subject-1" && + input.CapabilityID == "cap.lookup" && + input.CapabilityName == "records.lookup" && + assert.ObjectsAreEqual(map[string]any{"value": "alpha"}, input.Arguments) + })).Run(func(context.Context, authz.AuthorizationInput) { + events = append(events, "authorize") + }).Return(nil).Once() + capability := validBuilderCapability("cap.lookup", "records.lookup") + capability.Handler = func(_ context.Context, _ authz.Subject, input builderInput) (builderOutput, error) { + events = append(events, "handler") + return builderOutput(input), nil + } + server := buildTestServer(t, authorizer, codemode.DefaultLimits(), capability) + + result, err := server.Execute(t.Context(), authz.Subject{ID: "subject-1"}, ` +def main(): + return records.lookup(value="alpha") +`) + + require.NoError(t, err) + assert.Equal(t, []string{"authorize", "handler"}, events) + assert.Equal(t, map[string]any{"value": "alpha"}, result) +} + +// TestServerExecuteRejectsArgumentsBeforeAuthorization proves malformed calls cannot reach policy or handlers. +func TestServerExecuteRejectsArgumentsBeforeAuthorization(t *testing.T) { + authorizer := authzmocks.NewMockAuthorizer(t) + 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{}, nil + } + server := buildTestServer(t, authorizer, codemode.DefaultLimits(), capability) + + _, err := server.Execute(t.Context(), authz.Subject{ID: "subject-1"}, ` +def main(): + return records.lookup() +`) + + require.ErrorIs(t, err, codemode.ErrInvalidArguments) + assert.Zero(t, handlerCalls.Load()) +} + +// TestServerExecuteRejectsTopLevelNativeCalls proves loading code cannot bypass the running-phase guard. +func TestServerExecuteRejectsTopLevelNativeCalls(t *testing.T) { + authorizer := authzmocks.NewMockAuthorizer(t) + 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{}, nil + } + server := buildTestServer(t, authorizer, codemode.DefaultLimits(), capability) + + _, err := server.Execute(t.Context(), authz.Subject{ID: "subject-1"}, ` +value = records.lookup(value="alpha") +def main(): + return value +`) + + require.ErrorIs(t, err, codemode.ErrInvalidProgram) + assert.Zero(t, handlerCalls.Load()) +} + +// TestServerExecuteFailsClosedBeforeHandlerDispatch proves denial and policy failures have no handler side effects. +func TestServerExecuteFailsClosedBeforeHandlerDispatch(t *testing.T) { + tests := []struct { + // name identifies the policy failure. + name string + + // configure installs one generated mock expectation. + configure func(*authzmocks.MockAuthorizer) + + // target is the expected public classification. + target error + }{ + { + name: "recognized denial", + configure: func(authorizer *authzmocks.MockAuthorizer) { + authorizer.EXPECT().Authorize(mock.Anything, mock.Anything).Return(authz.ErrDenied).Once() + }, + target: codemode.ErrPermissionDenied, + }, + { + name: "policy error", + configure: func(authorizer *authzmocks.MockAuthorizer) { + authorizer.EXPECT(). + Authorize(mock.Anything, mock.Anything). + Return(errors.New("trusted policy detail")). + Once() + }, + target: codemode.ErrPolicyFailure, + }, + { + name: "policy cancellation", + configure: func(authorizer *authzmocks.MockAuthorizer) { + authorizer.EXPECT().Authorize(mock.Anything, mock.Anything).Return(context.Canceled).Once() + }, + target: codemode.ErrPolicyFailure, + }, + { + name: "policy panic", + configure: func(authorizer *authzmocks.MockAuthorizer) { + authorizer.EXPECT(). + Authorize(mock.Anything, mock.Anything). + Run(func(context.Context, authz.AuthorizationInput) { + panic("trusted panic detail") + }). + Return(nil). + Once() + }, + target: codemode.ErrPolicyFailure, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + authorizer := authzmocks.NewMockAuthorizer(t) + tt.configure(authorizer) + 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{}, nil + } + server := buildTestServer(t, authorizer, codemode.DefaultLimits(), capability) + + _, err := server.Execute(t.Context(), authz.Subject{ID: "subject-1"}, ` +def main(): + return records.lookup(value="alpha") +`) + + require.ErrorIs(t, err, tt.target) + assert.Equal(t, tt.target.Error(), err.Error()) + assert.Zero(t, handlerCalls.Load()) + }) + } +} + +// TestServerExecuteProjectsHandlerFailuresWithoutTrustedDetail proves native failures remain opaque. +func TestServerExecuteProjectsHandlerFailuresWithoutTrustedDetail(t *testing.T) { + tests := []struct { + // name identifies the handler failure. + name string + + // handler is the failing native implementation. + handler codemode.Handler[builderInput, builderOutput] + + // target is the expected public classification. + target error + }{ + { + name: "handler error", + handler: func(context.Context, authz.Subject, builderInput) (builderOutput, error) { + return builderOutput{}, errors.New("trusted handler detail") + }, + target: codemode.ErrCapabilityFailure, + }, + { + name: "handler panic", + handler: func(context.Context, authz.Subject, builderInput) (builderOutput, error) { + panic("trusted handler panic") + }, + target: codemode.ErrInternal, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + capability := validBuilderCapability("cap.lookup", "records.lookup") + capability.Handler = tt.handler + server := buildTestServer(t, authz.AllowAll(), codemode.DefaultLimits(), capability) + + _, err := server.Execute(t.Context(), authz.Subject{ID: "subject-1"}, ` +def main(): + return records.lookup(value="alpha") +`) + + require.ErrorIs(t, err, tt.target) + assert.Equal(t, tt.target.Error(), err.Error()) + }) + } +} + +// TestServerExecuteEnforcesSourceCallStepAndResultLimits proves each public budget fails safely. +func TestServerExecuteEnforcesSourceCallStepAndResultLimits(t *testing.T) { + t.Run("source bytes", func(t *testing.T) { + limits := codemode.DefaultLimits() + limits.MaxSourceBytes = 8 + server := buildTestServer( + t, + authz.AllowAll(), + limits, + validBuilderCapability("cap.lookup", "records.lookup"), + ) + + _, err := server.Execute(t.Context(), authz.Subject{ID: "subject-1"}, `def main(): return None`) + + require.ErrorIs(t, err, codemode.ErrResourceLimit) + }) + + t.Run("native calls", func(t *testing.T) { + limits := codemode.DefaultLimits() + limits.MaxNativeCalls = 1 + var handlerCalls atomic.Int64 + capability := validBuilderCapability("cap.lookup", "records.lookup") + capability.Handler = func(_ context.Context, _ authz.Subject, input builderInput) (builderOutput, error) { + handlerCalls.Add(1) + return builderOutput(input), 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, int64(1), handlerCalls.Load()) + }) + + t.Run("execution steps", func(t *testing.T) { + limits := codemode.DefaultLimits() + limits.MaxExecutionSteps = 20 + server := buildTestServer( + t, + authz.AllowAll(), + limits, + validBuilderCapability("cap.lookup", "records.lookup"), + ) + + _, err := server.Execute(t.Context(), authz.Subject{ID: "subject-1"}, ` +def main(): + total = 0 + for item in range(1000000): + total += item + return total +`) + + require.ErrorIs(t, err, codemode.ErrResourceLimit) + }) + + t.Run("result bytes", func(t *testing.T) { + limits := codemode.DefaultLimits() + limits.MaxResultBytes = 4 + server := buildTestServer( + t, + authz.AllowAll(), + limits, + validBuilderCapability("cap.lookup", "records.lookup"), + ) + + _, err := server.Execute(t.Context(), authz.Subject{ID: "subject-1"}, ` +def main(): + return records.lookup(value="oversized") +`) + + require.ErrorIs(t, err, codemode.ErrResourceLimit) + }) +} + +// TestServerExecuteReturnsOnlyMainResult proves top-level values and printed text do not escape. +func TestServerExecuteReturnsOnlyMainResult(t *testing.T) { + server := buildTestServer( + t, + authz.AllowAll(), + codemode.DefaultLimits(), + validBuilderCapability("cap.lookup", "records.lookup"), + ) + + result, err := server.Execute(t.Context(), authz.Subject{ID: "subject-1"}, ` +print("discarded") +intermediate = "hidden" +def main(): + return records.lookup(value="visible") +`) + + require.NoError(t, err) + assert.Equal(t, map[string]any{"value": "visible"}, result) +} + +// TestServerExecuteRequiresTrustedInvocationInputs proves nil contexts and empty subjects fail before execution. +func TestServerExecuteRequiresTrustedInvocationInputs(t *testing.T) { + server := buildTestServer( + t, + authz.AllowAll(), + codemode.DefaultLimits(), + validBuilderCapability("cap.lookup", "records.lookup"), + ) + + //nolint:staticcheck // A nil context is the behavior under test. + _, err := server.Execute(nil, authz.Subject{ID: "subject-1"}, `def main(): return None`) + require.ErrorIs(t, err, codemode.ErrInternal) + _, err = server.Execute(t.Context(), authz.Subject{}, `def main(): return None`) + require.ErrorIs(t, err, codemode.ErrUnauthenticated) + ctx, cancel := context.WithCancel(t.Context()) + cancel() + _, err = server.Execute(ctx, authz.Subject{ID: "subject-1"}, `def main(): return None`) + require.ErrorIs(t, err, context.Canceled) + deadlineCtx, deadlineCancel := context.WithTimeout(t.Context(), 0) + defer deadlineCancel() + _, err = server.Execute(deadlineCtx, authz.Subject{ID: "subject-1"}, `def main(): return None`) + require.ErrorIs(t, err, codemode.ErrResourceLimit) + require.ErrorIs(t, err, context.DeadlineExceeded) +} + +// TestServerSupportsConcurrentDiscoveryAndExecution proves immutable server state is race-safe for parallel reads. +func TestServerSupportsConcurrentDiscoveryAndExecution(t *testing.T) { + var handlerCalls atomic.Int64 + capability := validBuilderCapability("cap.lookup", "records.lookup") + capability.Handler = func(_ context.Context, _ authz.Subject, input builderInput) (builderOutput, error) { + handlerCalls.Add(1) + return builderOutput(input), nil + } + server := buildTestServer(t, authz.AllowAll(), codemode.DefaultLimits(), capability) + + const workers = 32 + var workersDone sync.WaitGroup + failures := make(chan error, workers) + for range workers { + workersDone.Go(func() { + if _, err := server.Search("lookup"); err != nil { + failures <- err + return + } + if _, err := server.Describe("records.lookup"); err != nil { + failures <- err + return + } + _, err := server.Execute(t.Context(), authz.Subject{ID: "subject-1"}, ` +def main(): + return records.lookup(value="alpha") +`) + if err != nil { + failures <- err + } + }) + } + workersDone.Wait() + close(failures) + + for err := range failures { + require.NoError(t, err) + } + assert.Equal(t, int64(workers), handlerCalls.Load()) +} + +// buildTestServer registers one capability and returns a successfully built server. +func buildTestServer( + t *testing.T, + authorizer authz.Authorizer, + limits codemode.Limits, + capability codemode.Capability[builderInput, builderOutput], +) *codemode.Server { + t.Helper() + builder := codemode.New(codemode.Options{Authorizer: authorizer, Limits: limits}) + require.NoError(t, codemode.Register(builder, capability)) + server, err := builder.Build() + require.NoError(t, err) + return server +}