Traverse class and instance values for capability contracts

Author: Mauricio Gomes

vibes/capability_contracts.go

@@ -6,14 +6,18 @@ import (
 )
 
 type capabilityContractScanner struct {
-	seenArrays map[sliceIdentity]struct{}
-	seenMaps   map[uintptr]struct{}
+	seenArrays    map[sliceIdentity]struct{}
+	seenMaps      map[uintptr]struct{}
+	seenClasses   map[*ClassDef]struct{}
+	seenInstances map[*Instance]struct{}
 }
 
 func newCapabilityContractScanner() *capabilityContractScanner {
 	return &capabilityContractScanner{
-		seenArrays: make(map[sliceIdentity]struct{}),
-		seenMaps:   make(map[uintptr]struct{}),
+		seenArrays:    make(map[sliceIdentity]struct{}),
+		seenMaps:      make(map[uintptr]struct{}),
+		seenClasses:   make(map[*ClassDef]struct{}),
+		seenInstances: make(map[*Instance]struct{}),
 	}
 }
 
@@ -103,5 +107,32 @@ func (s *capabilityContractScanner) bindContracts(val Value, contractsByName map
 		for _, item := range entries {
 			s.bindContracts(item, contractsByName, target)
 		}
+	case KindClass:
+		classDef := val.Class()
+		if classDef == nil {
+			return
+		}
+		if _, seen := s.seenClasses[classDef]; seen {
+			return
+		}
+		s.seenClasses[classDef] = struct{}{}
+		for _, item := range classDef.ClassVars {
+			s.bindContracts(item, contractsByName, target)
+		}
+	case KindInstance:
+		instance := val.Instance()
+		if instance == nil {
+			return
+		}
+		if _, seen := s.seenInstances[instance]; seen {
+			return
+		}
+		s.seenInstances[instance] = struct{}{}
+		for _, item := range instance.Ivars {
+			s.bindContracts(item, contractsByName, target)
+		}
+		if instance.Class != nil {
+			s.bindContracts(NewClass(instance.Class), contractsByName, target)
+		}
 	}
 }

vibes/capability_contracts_test.go

@@ -82,6 +82,74 @@ func (unrelatedNamedContractCapability) CapabilityContracts() map[string]Capabil
 	}
 }
 
+type instanceIvarContractCapability struct {
+	invokeCount *int
+}
+
+func (c instanceIvarContractCapability) Bind(binding CapabilityBinding) (map[string]Value, error) {
+	classDef := &ClassDef{
+		Name:         "CapabilityBox",
+		Methods:      map[string]*ScriptFunction{},
+		ClassMethods: map[string]*ScriptFunction{},
+		ClassVars:    map[string]Value{},
+	}
+	instance := &Instance{
+		Class: classDef,
+		Ivars: map[string]Value{
+			"call": NewBuiltin("probe.call", func(exec *Execution, receiver Value, args []Value, kwargs map[string]Value, block Value) (Value, error) {
+				*c.invokeCount = *c.invokeCount + 1
+				return NewString("ok"), nil
+			}),
+		},
+	}
+	return map[string]Value{"box": NewInstance(instance)}, nil
+}
+
+func (c instanceIvarContractCapability) CapabilityContracts() map[string]CapabilityMethodContract {
+	return map[string]CapabilityMethodContract{
+		"probe.call": {
+			ValidateArgs: func(args []Value, kwargs map[string]Value, block Value) error {
+				if len(args) != 1 || args[0].Kind() != KindInt {
+					return fmt.Errorf("probe.call expects int")
+				}
+				return nil
+			},
+		},
+	}
+}
+
+type classVarContractCapability struct {
+	invokeCount *int
+}
+
+func (c classVarContractCapability) Bind(binding CapabilityBinding) (map[string]Value, error) {
+	classDef := &ClassDef{
+		Name:         "CapabilityHolder",
+		Methods:      map[string]*ScriptFunction{},
+		ClassMethods: map[string]*ScriptFunction{},
+		ClassVars: map[string]Value{
+			"call": NewBuiltin("probe.class_call", func(exec *Execution, receiver Value, args []Value, kwargs map[string]Value, block Value) (Value, error) {
+				*c.invokeCount = *c.invokeCount + 1
+				return NewString("ok"), nil
+			}),
+		},
+	}
+	return map[string]Value{"holder": NewClass(classDef)}, nil
+}
+
+func (c classVarContractCapability) CapabilityContracts() map[string]CapabilityMethodContract {
+	return map[string]CapabilityMethodContract{
+		"probe.class_call": {
+			ValidateArgs: func(args []Value, kwargs map[string]Value, block Value) error {
+				if len(args) != 1 || args[0].Kind() != KindInt {
+					return fmt.Errorf("probe.class_call expects int")
+				}
+				return nil
+			},
+		},
+	}
+}
+
 func TestCapabilityContractRejectsInvalidArguments(t *testing.T) {
 	engine := MustNewEngine(Config{})
 	script, err := engine.Compile(`def run()
@@ -184,3 +252,55 @@ end`)
 		t.Fatalf("unexpected merge result: %#v", result.Hash())
 	}
 }
+
+func TestCapabilityContractsTraverseInstanceValues(t *testing.T) {
+	engine := MustNewEngine(Config{})
+	script, err := engine.Compile(`def run()
+  box.call("bad")
+end`)
+	if err != nil {
+		t.Fatalf("compile failed: %v", err)
+	}
+
+	invocations := 0
+	_, err = script.Call(context.Background(), "run", nil, CallOptions{
+		Capabilities: []CapabilityAdapter{
+			instanceIvarContractCapability{invokeCount: &invocations},
+		},
+	})
+	if err == nil {
+		t.Fatalf("expected instance contract validation error")
+	}
+	if got := err.Error(); !strings.Contains(got, "probe.call expects int") {
+		t.Fatalf("unexpected error: %s", got)
+	}
+	if invocations != 0 {
+		t.Fatalf("instance capability should not execute when contract fails")
+	}
+}
+
+func TestCapabilityContractsTraverseClassValues(t *testing.T) {
+	engine := MustNewEngine(Config{})
+	script, err := engine.Compile(`def run()
+  holder.call("bad")
+end`)
+	if err != nil {
+		t.Fatalf("compile failed: %v", err)
+	}
+
+	invocations := 0
+	_, err = script.Call(context.Background(), "run", nil, CallOptions{
+		Capabilities: []CapabilityAdapter{
+			classVarContractCapability{invokeCount: &invocations},
+		},
+	})
+	if err == nil {
+		t.Fatalf("expected class contract validation error")
+	}
+	if got := err.Error(); !strings.Contains(got, "probe.class_call expects int") {
+		t.Fatalf("unexpected error: %s", got)
+	}
+	if invocations != 0 {
+		t.Fatalf("class capability should not execute when contract fails")
+	}
+}