feat(ctb): `PreimageOracle` LPP Diff Testing (#9147)

* feat(ctb): merkle tree diff testing

* Fix `go-ffi` merkle diff tester ++ add more tests

* fix(op-challenger): merkle tree signed integer conversion

---------
Co-authored-by: refcell <abigger87@gmail.com>

op-challenger/game/keccak/merkle/tree.go

@@ -56,7 +56,7 @@ func (m *merkleNode) IsRightChild(o *merkleNode) bool {
 // It is an append-only tree, where leaves are added from left to right.
 type BinaryMerkleTree struct {
 	Root      *merkleNode
-	LeafCount int
+	LeafCount uint64
 }
 
 func NewBinaryMerkleTree() *BinaryMerkleTree {
@@ -72,11 +72,11 @@ func (m *BinaryMerkleTree) RootHash() (rootHash common.Hash) {
 }
 
 // walkDownToMaxLeaf walks down the tree to the max leaf node.
-func (m *BinaryMerkleTree) walkDownToLeafCount(subtreeLeafCount int) *merkleNode {
-	maxSubtreeLeafCount := MaxLeafCount + 1
+func (m *BinaryMerkleTree) walkDownToLeafCount(subtreeLeafCount uint64) *merkleNode {
+	maxSubtreeLeafCount := uint64(MaxLeafCount) + 1
 	levelNode := m.Root
 	for height := 0; height < BinaryMerkleTreeDepth; height++ {
-		if subtreeLeafCount*2 <= maxSubtreeLeafCount {
+		if subtreeLeafCount*2 <= uint64(maxSubtreeLeafCount) {
 			if levelNode.Left == nil {
 				levelNode.Left = &merkleNode{
 					Label:  zeroHashes[height],
@@ -101,12 +101,17 @@ func (m *BinaryMerkleTree) walkDownToLeafCount(subtreeLeafCount int) *merkleNode
 
 // AddLeaf adds a leaf to the binary merkle tree.
 func (m *BinaryMerkleTree) AddLeaf(leaf types.Leaf) {
+	m.AddRawLeaf(leaf.Hash())
+}
+
+// AddRawLeaf adds a raw 32 byte leaf to the binary merkle tree.
+func (m *BinaryMerkleTree) AddRawLeaf(hash common.Hash) {
 	// Walk down to the new max leaf node.
 	m.LeafCount += 1
 	levelNode := m.walkDownToLeafCount(m.LeafCount)
 
 	// Set the leaf node data.
-	levelNode.Label = leaf.Hash()
+	levelNode.Label = hash
 
 	// Walk back up the tree, updating the hashes with its sibling hash.
 	for height := 0; height < BinaryMerkleTreeDepth; height++ {
@@ -137,7 +142,7 @@ func (m *BinaryMerkleTree) ProofAtIndex(index uint64) (proof Proof, err error) {
 		return proof, IndexOutOfBoundsError
 	}
 
-	levelNode := m.walkDownToLeafCount(int(index) + 1)
+	levelNode := m.walkDownToLeafCount(index + 1)
 	for height := 0; height < BinaryMerkleTreeDepth; height++ {
 		if levelNode.Parent.IsLeftChild(levelNode) {
 			if levelNode.Parent.Right == nil {

packages/contracts-bedrock/scripts/go-ffi/bin.go

@@ -11,6 +11,8 @@ func main() {
 		DiffTestUtils()
 	case "trie":
 		FuzzTrie()
+	case "merkle":
+		DiffMerkle()
 	default:
 		log.Fatal("Must pass a subcommand")
 	}

packages/contracts-bedrock/scripts/go-ffi/merkle.go

+package main
+
+import (
+	"fmt"
+	"log"
+	"os"
+	"strconv"
+
+	"github.com/ethereum-optimism/optimism/op-challenger/game/keccak/merkle"
+
+	"github.com/ethereum/go-ethereum/accounts/abi"
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/common/hexutil"
+	"github.com/ethereum/go-ethereum/crypto"
+)
+
+// VerifyMerkleProof verifies a merkle proof against the root hash and the leaf hash.
+// Reference: https://github.com/ethereum/consensus-specs/blob/dev/specs/phase0/beacon-chain.md#is_valid_merkle_branch
+func VerifyMerkleProof(root, leaf common.Hash, index uint64, proof [merkle.BinaryMerkleTreeDepth]common.Hash) bool {
+	value := leaf
+	for i := 0; i < merkle.BinaryMerkleTreeDepth; i++ {
+		if ((index >> i) & 1) == 1 {
+			value = crypto.Keccak256Hash(proof[i][:], value[:])
+		} else {
+			value = crypto.Keccak256Hash(value[:], proof[i][:])
+		}
+	}
+	return value == root
+}
+
+const (
+	// GenProof generates a merkle proof for a given leaf index by reconstructing the merkle tree from the passed
+	// leaves.
+	genProof = "gen_proof"
+)
+
+var (
+	rootAndProof, _ = abi.NewType("tuple", "", []abi.ArgumentMarshaling{
+		{Name: "root", Type: "bytes32"},
+		{Name: "proof", Type: "bytes32[]"},
+	})
+
+	merkleEncoder = abi.Arguments{
+		{Type: rootAndProof},
+	}
+)
+
+// DiffMerkle generates an abi-encoded `merkleTestCase` of a specified variant.
+func DiffMerkle() {
+	variant := os.Args[2]
+	if len(variant) == 0 {
+		log.Fatal("Must pass a variant to the merkle diff tester!")
+	}
+
+	switch variant {
+	case genProof:
+		if len(os.Args) < 5 {
+			log.Fatal("Invalid arguments to `gen_proof` variant.")
+		}
+
+		rawLeaves, err := hexutil.Decode(os.Args[3])
+		if err != nil {
+			log.Fatal("Failed to decode leaves: ", err)
+		}
+		index, err := strconv.ParseInt(os.Args[4], 10, 64)
+		if err != nil {
+			log.Fatal("Failed to parse leaf index: ", err)
+		}
+		merkleTree := merkle.NewBinaryMerkleTree()
+
+		// Append all leaves to the merkle tree.
+		for i := 0; i < len(rawLeaves)/32; i++ {
+			leaf := common.BytesToHash(rawLeaves[i<<5 : (i+1)<<5])
+			merkleTree.AddRawLeaf(leaf)
+		}
+
+		// Generate the proof for the given index.
+		proof, err := merkleTree.ProofAtIndex(uint64(index))
+		if err != nil {
+			log.Fatal("Failed to generate proof: ", err)
+		}
+
+		// Generate the merkle root.
+		root := merkleTree.RootHash()
+
+		// Return "abi.encode(root, proof)"
+		packed, err := merkleEncoder.Pack(struct {
+			Root  common.Hash
+			Proof [merkle.BinaryMerkleTreeDepth]common.Hash
+		}{
+			Root:  root,
+			Proof: proof,
+		})
+		if err != nil {
+			log.Fatal("Failed to ABI encode root and proof: ", err)
+		}
+		fmt.Print(hexutil.Encode(packed[32:]))
+	default:
+		log.Fatal("Invalid variant passed to merkle diff tester!")
+	}
+}

packages/contracts-bedrock/test/cannon/PreimageOracle.t.sol

@@ -671,6 +671,79 @@ contract PreimageOracle_LargePreimageProposals_Test is Test {
         });
     }
 
+    /// @notice Tests that squeezing a large preimage proposal after the challenge period has passed always succeeds and
+    ///         persists the correct data.
+    function testFuzz_squeeze_succeeds(uint256 _numBlocks, uint32 _partOffset) public {
+        _numBlocks = bound(_numBlocks, 1, 2 ** 8);
+        _partOffset = uint32(bound(_partOffset, 0, _numBlocks * LibKeccak.BLOCK_SIZE_BYTES + 8 - 1));
+
+        // Allocate the preimage data.
+        bytes memory data = new bytes(136 * _numBlocks);
+        for (uint256 i; i < data.length; i++) {
+            data[i] = bytes1(uint8(i % 256));
+        }
+
+        // Propose and squeeze a large preimage.
+        {
+            // Initialize the proposal.
+            oracle.initLPP(TEST_UUID, _partOffset, uint32(data.length));
+
+            // Add the leaves to the tree with correct state commitments.
+            LibKeccak.StateMatrix memory matrixA;
+            bytes32[] memory stateCommitments = _generateStateCommitments(matrixA, data);
+            oracle.addLeavesLPP(TEST_UUID, 0, data, stateCommitments, true);
+
+            // Construct the leaf preimage data for the blocks added.
+            LibKeccak.StateMatrix memory matrixB;
+            PreimageOracle.Leaf[] memory leaves = _generateLeaves(matrixB, data);
+
+            // Fetch the merkle proofs for the pre/post state leaves in the proposal tree.
+            bytes32 canonicalRoot = oracle.getTreeRootLPP(address(this), TEST_UUID);
+            (bytes32 rootA, bytes32[] memory preProof) = _generateProof(leaves.length - 2, leaves);
+            assertEq(rootA, canonicalRoot);
+            (bytes32 rootB, bytes32[] memory postProof) = _generateProof(leaves.length - 1, leaves);
+            assertEq(rootB, canonicalRoot);
+
+            // Warp past the challenge period.
+            vm.warp(block.timestamp + CHALLENGE_PERIOD + 1 seconds);
+
+            // Squeeze the LPP.
+            LibKeccak.StateMatrix memory preMatrix = _stateMatrixAtBlockIndex(data, leaves.length - 1);
+            oracle.squeezeLPP({
+                _claimant: address(this),
+                _uuid: TEST_UUID,
+                _stateMatrix: preMatrix,
+                _preState: leaves[leaves.length - 2],
+                _preStateProof: preProof,
+                _postState: leaves[leaves.length - 1],
+                _postStateProof: postProof
+            });
+        }
+
+        // Validate the preimage part
+        {
+            bytes32 finalDigest = keccak256(data);
+            bytes32 expectedPart;
+            assembly {
+                switch lt(_partOffset, 0x08)
+                case true {
+                    mstore(0x00, shl(192, mload(data)))
+                    mstore(0x08, mload(add(data, 0x20)))
+                    expectedPart := mload(_partOffset)
+                }
+                default {
+                    // Clean the word after `data` so we don't get any dirty bits.
+                    mstore(add(add(data, 0x20), mload(data)), 0x00)
+                    expectedPart := mload(add(add(data, 0x20), sub(_partOffset, 0x08)))
+                }
+            }
+
+            assertTrue(oracle.preimagePartOk(finalDigest, _partOffset));
+            assertEq(oracle.preimageLengths(finalDigest), data.length);
+            assertEq(oracle.preimageParts(finalDigest, _partOffset), expectedPart);
+        }
+    }
+
     /// @notice Tests that a valid leaf cannot be countered with the `challengeFirst` function.
     function test_challengeFirst_validCommitment_reverts() public {
         // Allocate the preimage data.
@@ -792,6 +865,99 @@ contract PreimageOracle_LargePreimageProposals_Test is Test {
         assertTrue(metaData.countered());
     }
 
+    /// @notice Tests that challenging the first divergence in a large preimage proposal at an arbitrary location
+    ///         in the leaf values always succeeds.
+    function testFuzz_challenge_arbitraryLocation_succeeds(uint256 _lastCorrectLeafIdx, uint256 _numBlocks) public {
+        _numBlocks = bound(_numBlocks, 1, 2 ** 8);
+        _lastCorrectLeafIdx = bound(_lastCorrectLeafIdx, 0, _numBlocks - 1);
+
+        // Allocate the preimage data.
+        bytes memory data = new bytes(136 * _numBlocks);
+
+        // Initialize the proposal.
+        oracle.initLPP(TEST_UUID, 0, uint32(data.length));
+
+        // Add the leaves to the tree with corrupted state commitments.
+        LibKeccak.StateMatrix memory matrixA;
+        bytes32[] memory stateCommitments = _generateStateCommitments(matrixA, data);
+        for (uint256 i = _lastCorrectLeafIdx + 1; i < stateCommitments.length; i++) {
+            stateCommitments[i] = 0;
+        }
+        oracle.addLeavesLPP(TEST_UUID, 0, data, stateCommitments, true);
+
+        // Construct the leaf preimage data for the blocks added and corrupt the state commitments.
+        LibKeccak.StateMatrix memory matrixB;
+        PreimageOracle.Leaf[] memory leaves = _generateLeaves(matrixB, data);
+        for (uint256 i = _lastCorrectLeafIdx + 1; i < leaves.length; i++) {
+            leaves[i].stateCommitment = 0;
+        }
+
+        // Avoid stack too deep
+        uint256 agreedLeafIdx = _lastCorrectLeafIdx;
+        uint256 disputedLeafIdx = agreedLeafIdx + 1;
+
+        // Fetch the merkle proofs for the pre/post state leaves in the proposal tree.
+        bytes32 canonicalRoot = oracle.getTreeRootLPP(address(this), TEST_UUID);
+        (bytes32 rootA, bytes32[] memory preProof) = _generateProof(agreedLeafIdx, leaves);
+        assertEq(rootA, canonicalRoot);
+        (bytes32 rootB, bytes32[] memory postProof) = _generateProof(disputedLeafIdx, leaves);
+        assertEq(rootB, canonicalRoot);
+
+        LibKeccak.StateMatrix memory preMatrix = _stateMatrixAtBlockIndex(data, disputedLeafIdx);
+        oracle.challengeLPP({
+            _claimant: address(this),
+            _uuid: TEST_UUID,
+            _stateMatrix: preMatrix,
+            _preState: leaves[agreedLeafIdx],
+            _preStateProof: preProof,
+            _postState: leaves[disputedLeafIdx],
+            _postStateProof: postProof
+        });
+
+        LPPMetaData metaData = oracle.proposalMetadata(address(this), TEST_UUID);
+        assertTrue(metaData.countered());
+    }
+
+    /// @notice Tests that challenging the a divergence in a large preimage proposal at the first leaf always succeeds.
+    function testFuzz_challengeFirst_succeeds(uint256 _numBlocks) public {
+        _numBlocks = bound(_numBlocks, 1, 2 ** 8);
+
+        // Allocate the preimage data.
+        bytes memory data = new bytes(136 * _numBlocks);
+
+        // Initialize the proposal.
+        oracle.initLPP(TEST_UUID, 0, uint32(data.length));
+
+        // Add the leaves to the tree with corrupted state commitments.
+        bytes32[] memory stateCommitments = new bytes32[](_numBlocks + 1);
+        for (uint256 i = 0; i < stateCommitments.length; i++) {
+            stateCommitments[i] = 0;
+        }
+        oracle.addLeavesLPP(TEST_UUID, 0, data, stateCommitments, true);
+
+        // Construct the leaf preimage data for the blocks added and corrupt the state commitments.
+        LibKeccak.StateMatrix memory matrixB;
+        PreimageOracle.Leaf[] memory leaves = _generateLeaves(matrixB, data);
+        for (uint256 i = 0; i < leaves.length; i++) {
+            leaves[i].stateCommitment = 0;
+        }
+
+        // Fetch the merkle proofs for the pre/post state leaves in the proposal tree.
+        bytes32 canonicalRoot = oracle.getTreeRootLPP(address(this), TEST_UUID);
+        (bytes32 rootA, bytes32[] memory postProof) = _generateProof(0, leaves);
+        assertEq(rootA, canonicalRoot);
+
+        oracle.challengeFirstLPP({
+            _claimant: address(this),
+            _uuid: TEST_UUID,
+            _postState: leaves[0],
+            _postStateProof: postProof
+        });
+
+        LPPMetaData metaData = oracle.proposalMetadata(address(this), TEST_UUID);
+        assertTrue(metaData.countered());
+    }
+
     /// @notice Tests that a valid leaf cannot be countered with the `challenge` function in the middle of the tree.
     function test_challenge_validCommitment_reverts() public {
         // Allocate the preimage data.
@@ -1019,16 +1185,19 @@ contract PreimageOracle_LargePreimageProposals_Test is Test {
         returns (PreimageOracle.Leaf[] memory leaves_)
     {
         bytes memory data = LibKeccak.padMemory(_data);
-        uint256 numLeaves = data.length / LibKeccak.BLOCK_SIZE_BYTES;
+        uint256 numCommitments = data.length / LibKeccak.BLOCK_SIZE_BYTES;
 
-        leaves_ = new PreimageOracle.Leaf[](numLeaves);
-        for (uint256 i = 0; i < numLeaves; i++) {
+        leaves_ = new PreimageOracle.Leaf[](numCommitments);
+        for (uint256 i = 0; i < numCommitments; i++) {
             bytes memory blockSlice = Bytes.slice(data, i * LibKeccak.BLOCK_SIZE_BYTES, LibKeccak.BLOCK_SIZE_BYTES);
             LibKeccak.absorb(_stateMatrix, blockSlice);
             LibKeccak.permutation(_stateMatrix);
-            bytes32 stateCommitment = keccak256(abi.encode(_stateMatrix));
 
-            leaves_[i] = PreimageOracle.Leaf({ input: blockSlice, index: i, stateCommitment: stateCommitment });
+            leaves_[i] = PreimageOracle.Leaf({
+                input: blockSlice,
+                index: uint32(i),
+                stateCommitment: keccak256(abi.encode(_stateMatrix))
+            });
         }
     }
 
@@ -1071,4 +1240,27 @@ contract PreimageOracle_LargePreimageProposals_Test is Test {
             stateCommitments_[i] = keccak256(abi.encode(_stateMatrix));
         }
     }
+
+    /// @notice Calls out to the `go-ffi` tool to generate a merkle proof for the leaf at `_leafIdx` in a merkle tree
+    ///         constructed with `_leaves`.
+    function _generateProof(
+        uint256 _leafIdx,
+        PreimageOracle.Leaf[] memory _leaves
+    )
+        internal
+        returns (bytes32 root_, bytes32[] memory proof_)
+    {
+        bytes32[] memory leaves = new bytes32[](_leaves.length);
+        for (uint256 i = 0; i < _leaves.length; i++) {
+            leaves[i] = _hashLeaf(_leaves[i]);
+        }
+
+        string[] memory commands = new string[](5);
+        commands[0] = "scripts/go-ffi/go-ffi";
+        commands[1] = "merkle";
+        commands[2] = "gen_proof";
+        commands[3] = vm.toString(abi.encodePacked(leaves));
+        commands[4] = vm.toString(_leafIdx);
+        (root_, proof_) = abi.decode(vm.ffi(commands), (bytes32, bytes32[]));
+    }
 }