Exclude pre-call roots from post-call contract rebinding

Author: mgomes

vibes/capability_contracts_test.go

@@ -430,6 +430,28 @@ func (argPassThroughContractCapability) CapabilityContracts() map[string]Capabil
 	}
 }
 
+type hashMergeContractLeakProbeCapability struct{}
+
+func (hashMergeContractLeakProbeCapability) Bind(binding CapabilityBinding) (map[string]Value, error) {
+	return map[string]Value{
+		"cap": NewObject(map[string]Value{
+			"touch": NewBuiltin("cap.touch", func(exec *Execution, receiver Value, args []Value, kwargs map[string]Value, block Value) (Value, error) {
+				return NewString("ok"), nil
+			}),
+		}),
+	}, nil
+}
+
+func (hashMergeContractLeakProbeCapability) CapabilityContracts() map[string]CapabilityMethodContract {
+	return map[string]CapabilityMethodContract{
+		"hash.merge": {
+			ValidateArgs: func(args []Value, kwargs map[string]Value, block Value) error {
+				return fmt.Errorf("hash.merge contract should not bind to foreign builtin")
+			},
+		},
+	}
+}
+
 func TestCapabilityContractRejectsInvalidArguments(t *testing.T) {
 	engine := MustNewEngine(Config{})
 	script, err := engine.Compile(`def run()
@@ -778,3 +800,30 @@ end`)
 		t.Fatalf("unexpected result: %#v", result)
 	}
 }
+
+func TestCapabilityContractsDoNotHijackReceiverStoredForeignBuiltins(t *testing.T) {
+	engine := MustNewEngine(Config{})
+	script, err := engine.Compile(`def run()
+  cap.foreign = { a: 1 }.merge
+  cap.touch()
+  cap.foreign({ b: 2 })
+end`)
+	if err != nil {
+		t.Fatalf("compile failed: %v", err)
+	}
+
+	result, err := script.Call(context.Background(), "run", nil, CallOptions{
+		Capabilities: []CapabilityAdapter{
+			hashMergeContractLeakProbeCapability{},
+		},
+	})
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if result.Kind() != KindHash {
+		t.Fatalf("expected hash result, got %v", result.Kind())
+	}
+	if got, ok := result.Hash()["b"]; !ok || got.Kind() != KindInt || got.Int() != 2 {
+		t.Fatalf("unexpected merge result: %#v", result.Hash())
+	}
+}

vibes/execution.go

@@ -615,14 +615,20 @@ func (exec *Execution) invokeCallable(callee Value, receiver Value, args []Value
 	case KindBuiltin:
 		builtin := callee.Builtin()
 		scope := exec.capabilityContractScopes[builtin]
-		var preCallArgBuiltins map[*Builtin]struct{}
+		var preCallKnownBuiltins map[*Builtin]struct{}
 		if scope != nil && len(scope.contracts) > 0 {
-			preCallArgBuiltins = make(map[*Builtin]struct{})
+			preCallKnownBuiltins = make(map[*Builtin]struct{})
+			if receiver.Kind() != KindNil {
+				collectCapabilityBuiltins(receiver, preCallKnownBuiltins)
+			}
+			for _, root := range scope.roots {
+				collectCapabilityBuiltins(root, preCallKnownBuiltins)
+			}
 			for _, arg := range args {
-				collectCapabilityBuiltins(arg, preCallArgBuiltins)
+				collectCapabilityBuiltins(arg, preCallKnownBuiltins)
 			}
 			for _, kwarg := range kwargs {
-				collectCapabilityBuiltins(kwarg, preCallArgBuiltins)
+				collectCapabilityBuiltins(kwarg, preCallKnownBuiltins)
 			}
 		}
 		contract, hasContract := exec.capabilityContracts[builtin]
@@ -645,22 +651,22 @@ func (exec *Execution) invokeCallable(callee Value, receiver Value, args []Value
 			// Capability methods can lazily publish additional builtins at runtime
 			// (e.g. through factory return values or receiver mutation). Re-scan
 			// these values so future calls still enforce declared contracts.
-			bindCapabilityContracts(result, scope, exec.capabilityContracts, exec.capabilityContractScopes)
+			bindCapabilityContractsExcluding(result, scope, exec.capabilityContracts, exec.capabilityContractScopes, preCallKnownBuiltins)
 			if receiver.Kind() != KindNil {
-				bindCapabilityContracts(receiver, scope, exec.capabilityContracts, exec.capabilityContractScopes)
+				bindCapabilityContractsExcluding(receiver, scope, exec.capabilityContracts, exec.capabilityContractScopes, preCallKnownBuiltins)
 			}
 			// Methods can mutate sibling scope roots via captured references; refresh
 			// all adapter roots so newly exposed builtins also get bound.
 			for _, root := range scope.roots {
-				bindCapabilityContracts(root, scope, exec.capabilityContracts, exec.capabilityContractScopes)
+				bindCapabilityContractsExcluding(root, scope, exec.capabilityContracts, exec.capabilityContractScopes, preCallKnownBuiltins)
 			}
 			// Methods can also publish builtins by mutating positional or keyword
 			// argument objects supplied by script code.
 			for _, arg := range args {
-				bindCapabilityContractsExcluding(arg, scope, exec.capabilityContracts, exec.capabilityContractScopes, preCallArgBuiltins)
+				bindCapabilityContractsExcluding(arg, scope, exec.capabilityContracts, exec.capabilityContractScopes, preCallKnownBuiltins)
 			}
 			for _, kwarg := range kwargs {
-				bindCapabilityContractsExcluding(kwarg, scope, exec.capabilityContracts, exec.capabilityContractScopes, preCallArgBuiltins)
+				bindCapabilityContractsExcluding(kwarg, scope, exec.capabilityContracts, exec.capabilityContractScopes, preCallKnownBuiltins)
 			}
 		}
 		return result, nil