Skip to content

Commit d8f2e2d

Browse files
committed
fix(vibes): enforce step checks during db.each iteration
1 parent 4a1df71 commit d8f2e2d

2 files changed

Lines changed: 32 additions & 0 deletions

File tree

vibes/capability_db.go

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -176,6 +176,9 @@ func (c *dbCapability) callEach(exec *Execution, receiver Value, args []Value, k
176176
return NewNil(), err
177177
}
178178
for idx, row := range rows {
179+
if err := exec.step(); err != nil {
180+
return NewNil(), err
181+
}
179182
if err := validateCapabilityDataOnlyValue(fmt.Sprintf("%s.each row %d", c.name, idx), row); err != nil {
180183
return NewNil(), err
181184
}

vibes/capability_foundations_test.go

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -99,6 +99,35 @@ end`)
9999
}
100100
}
101101

102+
func TestCapabilityFoundationsEachNoopBlockRespectsStepQuota(t *testing.T) {
103+
rows := make([]Value, 120)
104+
for i := range rows {
105+
rows[i] = NewHash(map[string]Value{"amount": NewInt(1)})
106+
}
107+
db := &dbCapabilityStub{eachRows: rows}
108+
109+
engine := MustNewEngine(Config{StepQuota: 20})
110+
script, err := engine.Compile(`def run()
111+
db.each("ScoreEntry") do |row|
112+
end
113+
end`)
114+
if err != nil {
115+
t.Fatalf("compile failed: %v", err)
116+
}
117+
118+
_, err = script.Call(context.Background(), "run", nil, CallOptions{
119+
Capabilities: []CapabilityAdapter{
120+
MustNewDBCapability("db", db),
121+
},
122+
})
123+
if err == nil {
124+
t.Fatalf("expected step quota error")
125+
}
126+
if got := err.Error(); !strings.Contains(got, "step quota exceeded") {
127+
t.Fatalf("unexpected error: %s", got)
128+
}
129+
}
130+
102131
func TestCapabilityFoundationsEachRespectsRecursionLimit(t *testing.T) {
103132
db := &dbCapabilityStub{
104133
eachRows: []Value{

0 commit comments

Comments
 (0)