From 22f0c71153ad025dd371fd923685a00e7c2acf46 Mon Sep 17 00:00:00 2001
From: Minhyuk Kim <kimminhyuk1004@gmail.com>
Date: Tue, 22 Oct 2024 16:31:55 -0600
Subject: [PATCH] Add offset validation and remove unnecessary code per code
 review

---
 rvgo/fast/instrumented.go |  3 +++
 rvgo/fast/memory.go       |  5 +----
 rvgo/fast/vm.go           |  6 +-----
 rvgo/slow/vm.go           | 12 +++++++-----
 rvsol/src/RISCV.sol       |  5 -----
 5 files changed, 12 insertions(+), 19 deletions(-)

diff --git a/rvgo/fast/instrumented.go b/rvgo/fast/instrumented.go
index 8898c4d0..80187cc2 100644
--- a/rvgo/fast/instrumented.go
+++ b/rvgo/fast/instrumented.go
@@ -83,6 +83,9 @@ func (m *InstrumentedState) readPreimage(key [32]byte, offset uint64) (dat [32]b
 		m.lastPreimage = preimage
 	}
 	m.lastPreimageOffset = offset
+	if offset >= uint64(len(preimage)) {
+		panic("Preimage offset out-of-bounds")
+	}
 	datLen = uint64(copy(dat[:], preimage[offset:]))
 	return
 }
diff --git a/rvgo/fast/memory.go b/rvgo/fast/memory.go
index 8b45ee73..7db1bf9a 100644
--- a/rvgo/fast/memory.go
+++ b/rvgo/fast/memory.go
@@ -195,10 +195,7 @@ func (m *Memory) SetAligned(addr uint64, dat []byte) {
 		m.Invalidate(addr) // invalidate this branch of memory, now that the value changed
 	}
 
-	d := copy(p.Data[pageAddr:], dat)
-	if d == len(dat) {
-		return // if all the data fitted in the page, we're done
-	}
+	copy(p.Data[pageAddr:], dat)
 }
 
 func (m *Memory) GetUnaligned(addr uint64, dest []byte) {
diff --git a/rvgo/fast/vm.go b/rvgo/fast/vm.go
index f6d4feb4..ba2f17d2 100644
--- a/rvgo/fast/vm.go
+++ b/rvgo/fast/vm.go
@@ -266,10 +266,6 @@ func (inst *InstrumentedState) riscvStep() (outErr error) {
 		s.Memory.SetAligned(rightAddr, bytez[leftSize:size])
 	}
 
-	//
-	// CSR (control and status registers) functions
-	//
-
 	//
 	// Preimage oracle interactions
 	//
@@ -378,7 +374,7 @@ func (inst *InstrumentedState) riscvStep() (outErr error) {
 			// ensure MAP_ANONYMOUS is set and fd == -1
 			if (flags&0x20) == 0 || fd != u64Mask() {
 				addr = u64Mask()
-				errCode = toU64(0x4d) // no error
+				errCode = toU64(0x4d) // EBADF
 			} else {
 				// ignore: prot, flags, fd, offset
 				switch addr {
diff --git a/rvgo/slow/vm.go b/rvgo/slow/vm.go
index 42103401..458f00c6 100644
--- a/rvgo/slow/vm.go
+++ b/rvgo/slow/vm.go
@@ -128,7 +128,13 @@ func Step(calldata []byte, po PreimageOracle) (stateHash common.Hash, outErr err
 	}
 
 	proofContentOffset := shortToU64(uint16(stateContentOffset) + paddedStateSize + 32)
-	// TODO: validate abi offset values?
+
+	if and(b32asBEWord(calldataload(shortToU64(uint16(stateContentOffset)+paddedStateSize))), shortToU256(60-1)) != (U256{}) {
+		// proof offset must be stateContentOffset+paddedStateSize+32
+		// proof size: 64-5+1=60 * 32 byte leaf,
+		// but multiple memProof can be used, so the proofSize must be a multiple of 60
+		panic("invalid proof offset input")
+	}
 
 	//
 	// State loading
@@ -450,10 +456,6 @@ func Step(calldata []byte, po PreimageOracle) (stateHash common.Hash, outErr err
 		storeMemUnaligned(addr, size, u64ToU256(value), proofIndexL, proofIndexR)
 	}
 
-	//
-	// CSR (control and status registers) functions
-	//
-
 	//
 	// Preimage oracle interactions
 	//
diff --git a/rvsol/src/RISCV.sol b/rvsol/src/RISCV.sol
index 80b04d46..c53d6c75 100644
--- a/rvsol/src/RISCV.sol
+++ b/rvsol/src/RISCV.sol
@@ -333,7 +333,6 @@ contract RISCV {
                 out := 548
             }
             if iszero(eq(proof.offset, proofContentOffset())) { revert(0, 0) }
-            // TODO: validate abi offset values?
 
             //
             // State loading
@@ -728,10 +727,6 @@ contract RISCV {
                 storeMemUnaligned(addr, size, u64ToU256(value), proofIndexL, proofIndexR)
             }
 
-            //
-            // CSR (control and status registers) functions
-            //
-
             //
             // Preimage oracle interactions
             //