Fix incorrect depth set bug (#170)

* fix(vm): fix incorrect depth set bug

* fix(tests): fix failing test

core/vm/instructions.go

@@ -1082,13 +1082,6 @@ func opReturn(pc *uint64, interpreter *EVMInterpreter, scope *ScopeContext) ([]b
 			// If a handle is returned, automatically make it available to the caller.
 			verifiedCiphertext.verifiedDepths.add(interpreter.evm.depth - 1)
 		}
-		if interpreter.evm.Commit {
-			interpreter.evm.Logger.Info("opReturn removing ciphertext from depth",
-				"handle", verifiedCiphertext.ciphertext.getHash().Hex(),
-				"depth", interpreter.evm.depth)
-		}
-		// Delete the current EVM depth. Do that after the `isVerifiedAtCurrentDepth()` check.
-		verifiedCiphertext.verifiedDepths.del(interpreter.evm.depth)
 	}
 
 	return ret, errStopToken

core/vm/instructions_test.go

@@ -978,9 +978,6 @@ func TestOpReturnDelegation(t *testing.T) {
 	if !bytes.Equal(ct.serialize(), ctAfterOp.ciphertext.serialize()) {
 		t.Fatalf("expected ciphertext after the return op is the same as original")
 	}
-	if ctAfterOp.verifiedDepths.count() != 1 || !ctAfterOp.verifiedDepths.has(interpreter.evm.depth) {
-		t.Fatalf("expected ciphertext to be verified only at depth - 1 after the return op")
-	}
 }
 
 func TestOpReturnUnverifyIfNotReturned(t *testing.T) {

core/vm/interpreter.go

@@ -165,7 +165,18 @@ func (in *EVMInterpreter) Run(contract *Contract, input []byte, readOnly bool) (
 
 	// Increment the call depth which is restricted to 1024
 	in.evm.depth++
-	defer func() { in.evm.depth-- }()
+	defer func() {
+		for _, verifiedCiphertext := range in.verifiedCiphertexts {
+			if in.evm.Commit {
+				in.evm.Logger.Info("Run removing ciphertext from depth",
+					"handle", verifiedCiphertext.ciphertext.getHash().Hex(),
+					"depth", in.evm.depth)
+			}
+			// Delete the current EVM depth from the set of verified depths.
+			verifiedCiphertext.verifiedDepths.del(in.evm.depth)
+		}
+		in.evm.depth--
+	}()
 
 	// Make sure the readOnly is only set if we aren't in readOnly yet.
 	// This also makes sure that the readOnly flag isn't removed for child calls.

core/vm/interpreter_test.go

@@ -75,3 +75,46 @@ func TestLoopInterrupt(t *testing.T) {
 	}
 
 }
+
+// Generates a contract that reverts immediately.
+func newRevertingContract() *Contract {
+	addr := AccountRef{}
+	c := NewContract(addr, addr, big.NewInt(0), 100000)
+	c.Code = make([]byte, 5)
+	c.Code[0] = byte(PUSH1)
+	c.Code[1] = byte(0)
+	c.Code[2] = byte(PUSH1)
+	c.Code[3] = byte(0)
+	c.Code[4] = byte(REVERT)
+	return c
+}
+
+func TestDepthRemovalOnRevert(t *testing.T) {
+	statedb, _ := state.New(common.Hash{}, state.NewDatabase(rawdb.NewMemoryDatabase()), nil)
+	evm := NewEVM(BlockContext{}, TxContext{}, statedb, params.AllEthashProtocolChanges, Config{})
+	evm.depth = 1 // simulate first "call"
+
+	contract := newRevertingContract()
+
+	h := common.BytesToHash([]byte("1337"))
+
+	verifiedCt := verifiedCiphertext{
+		newDepthSet(),
+		&tfheCiphertext{
+			make([]byte, 0),
+			&h,
+			FheUint8,
+		},
+	}
+	verifiedCt.verifiedDepths.add(1)
+	verifiedCt.verifiedDepths.add(2) // simulate passing `h` as a call argument to `contract`
+
+	evm.interpreter.verifiedCiphertexts[h] = &verifiedCt
+
+	evm.interpreter.Run(contract, make([]byte, 0), false)
+
+	resultingVerifiedDepths := evm.interpreter.verifiedCiphertexts[h].verifiedDepths
+	if resultingVerifiedDepths.has(2) {
+		t.Fatalf("verified depth should have been removed after revert")
+	}
+}