Fix KZG Precompile Oracle (#9577)

* op-e2e: Fix flaky TestOutputCannonStepWithKZGPointEvaluation test

* fault-proofs: Fix on-chain KZG precompile oracle

* fix for flakes in CI

and add nil check when logging oracleKey

cannon/mipsevm/evm_test.go

@@ -9,7 +9,6 @@ import (
 	"math/big"
 	"os"
 	"path"
-	"strings"
 	"testing"
 	"time"
 
@@ -42,20 +41,25 @@ func MarkdownTracer() vm.EVMLogger {
 }
 
 type MIPSEVM struct {
-	env      *vm.EVM
-	evmState *state.StateDB
-	addrs    *Addresses
+	env         *vm.EVM
+	evmState    *state.StateDB
+	addrs       *Addresses
+	localOracle PreimageOracle
 }
 
 func NewMIPSEVM(contracts *Contracts, addrs *Addresses) *MIPSEVM {
 	env, evmState := NewEVMEnv(contracts, addrs)
-	return &MIPSEVM{env, evmState, addrs}
+	return &MIPSEVM{env, evmState, addrs, nil}
 }
 
 func (m *MIPSEVM) SetTracer(tracer vm.EVMLogger) {
 	m.env.Config.Tracer = tracer
 }
 
+func (m *MIPSEVM) SetLocalOracle(oracle PreimageOracle) {
+	m.localOracle = oracle
+}
+
 // Step is a pure function that computes the poststate from the VM state encoded in the StepWitness.
 func (m *MIPSEVM) Step(t *testing.T, stepWitness *StepWitness) []byte {
 	sender := common.Address{0x13, 0x37}
@@ -66,7 +70,7 @@ func (m *MIPSEVM) Step(t *testing.T, stepWitness *StepWitness) []byte {
 
 	if stepWitness.HasPreimage() {
 		t.Logf("reading preimage key %x at offset %d", stepWitness.PreimageKey, stepWitness.PreimageOffset)
-		poInput, err := encodePreimageOracleInput(t, stepWitness, LocalContext{})
+		poInput, err := encodePreimageOracleInput(t, stepWitness, LocalContext{}, m.localOracle)
 		require.NoError(t, err, "encode preimage oracle input")
 		_, leftOverGas, err := m.env.Call(vm.AccountRef(sender), m.addrs.Oracle, poInput, startingGas, big.NewInt(0))
 		require.NoErrorf(t, err, "evm should not fail, took %d gas", startingGas-leftOverGas)
@@ -100,7 +104,7 @@ func encodeStepInput(t *testing.T, wit *StepWitness, localContext LocalContext)
 	return input
 }
 
-func encodePreimageOracleInput(t *testing.T, wit *StepWitness, localContext LocalContext) ([]byte, error) {
+func encodePreimageOracleInput(t *testing.T, wit *StepWitness, localContext LocalContext, localOracle PreimageOracle) ([]byte, error) {
 	if wit.PreimageKey == ([32]byte{}) {
 		return nil, errors.New("cannot encode pre-image oracle input, witness has no pre-image to proof")
 	}
@@ -132,6 +136,18 @@ func encodePreimageOracleInput(t *testing.T, wit *StepWitness, localContext Loca
 			wit.PreimageValue[8:])
 		require.NoError(t, err)
 		return input, nil
+	case preimage.KZGPointEvaluationKeyType:
+		if localOracle == nil {
+			return nil, fmt.Errorf("local oracle is required for point evaluation preimages")
+		}
+		preimage := localOracle.GetPreimage(preimage.Keccak256Key(wit.PreimageKey).PreimageKey())
+		input, err := preimageAbi.Pack(
+			"loadKZGPointEvaluationPreimagePart",
+			new(big.Int).SetUint64(uint64(wit.PreimageOffset)),
+			preimage,
+		)
+		require.NoError(t, err)
+		return input, nil
 	default:
 		return nil, fmt.Errorf("unsupported pre-image type %d, cannot prepare preimage with key %x offset %d for oracle",
 			wit.PreimageKey[0], wit.PreimageKey, wit.PreimageOffset)
@@ -147,15 +163,13 @@ func TestEVM(t *testing.T) {
 
 	for _, f := range testFiles {
 		t.Run(f.Name(), func(t *testing.T) {
-			var oracle PreimageOracle
-			if strings.HasPrefix(f.Name(), "oracle") {
-				oracle = staticOracle(t, []byte("hello world"))
-			}
+			oracle := selectOracleFixture(t, f.Name())
 			// Short-circuit early for exit_group.bin
 			exitGroup := f.Name() == "exit_group.bin"
 
 			evm := NewMIPSEVM(contracts, addrs)
 			evm.SetTracer(tracer)
+			evm.SetLocalOracle(oracle)
 
 			fn := path.Join("open_mips_tests/test/bin", f.Name())
 			programMem, err := os.ReadFile(fn)

cannon/mipsevm/open_mips_tests/test/oracle_kzg.asm

+.section .test, "x"
+    .balign 4
+    .set    noreorder
+    .global test
+    .ent    test
+
+# load hash at 0x30001000
+# 0x47173285 a8d7341e 5e972fc6 77286384 f802f8ef 42a5ec5f 03bbfa25 4cb01fad = keccak("hello world")
+# 0x06173285 a8d7341e 5e972fc6 77286384 f802f8ef 42a5ec5f 03bbfa25 4cb01fad = keccak("hello world").key (kzg)
+test:
+  lui $s0, 0x3000
+  ori $s0, 0x1000
+
+  lui $t0, 0x0617
+  ori $t0, 0x3285
+  sw $t0, 0($s0)
+  lui $t0, 0xa8d7
+  ori $t0, 0x341e
+  sw $t0, 4($s0)
+  lui $t0, 0x5e97
+  ori $t0, 0x2fc6
+  sw $t0, 8($s0)
+  lui $t0, 0x7728
+  ori $t0, 0x6384
+  sw $t0, 0xc($s0)
+  lui $t0, 0xf802
+  ori $t0, 0xf8ef
+  sw $t0, 0x10($s0)
+  lui $t0, 0x42a5
+  ori $t0, 0xec5f
+  sw $t0, 0x14($s0)
+  lui $t0, 0x03bb
+  ori $t0, 0xfa25
+  sw $t0, 0x18($s0)
+  lui $t0, 0x4cb0
+  ori $t0, 0x1fad
+  sw $t0, 0x1c($s0)
+
+# preimage request - write(fdPreimageWrite, preimageData, 32)
+  li $a0, 6
+  li $a1, 0x30001000
+  li $t0, 8
+  li $a2, 4
+$writeloop:
+  li $v0, 4004
+  syscall
+  addiu $a1, $a1, 4
+  addiu $t0, $t0, -1
+  bnez $t0, $writeloop
+  nop
+
+# preimage response to 0x30002000 - read(fdPreimageRead, addr, count)
+# read preimage length
+  li $a0, 5
+  li $a1, 0x31000000
+  li $a2, 4
+  li $v0, 4003
+  syscall
+  li $a1, 0x31000004
+  li $v0, 4003
+  syscall
+# read the 1 byte preimage data
+  li $a1, 0x31000008
+  li $v0, 4003
+  syscall
+  nop
+
+# length at 0x31000000. We also check that the lower 32 bits are zero
+  lui $s1, 0x3100
+  lw $t0, 0($s1)
+  sltiu $t6, $t0, 1
+  li $s1, 0x31000004
+  lw $t0, 0($s1)
+# len should be 1
+  li $t4, 1
+  subu $t5, $t0, $t4
+  sltiu $v0, $t5, 1
+  and $v0, $v0, $t6
+
+# most likely broken. need to check pc for exact case
+# save results
+  lui     $s0, 0xbfff         # Load the base address 0xbffffff0
+  ori     $s0, 0xfff0
+  ori     $s1, $0, 1          # Prepare the 'done' status
+
+  sw      $v0, 8($s0)         # Set the test result
+  sw      $s1, 4($s0)         # Set 'done'
+
+$done:
+  jr      $ra
+  nop
+
+  .end test

cannon/mipsevm/state_test.go

@@ -31,10 +31,7 @@ func TestState(t *testing.T) {
 
 	for _, f := range testFiles {
 		t.Run(f.Name(), func(t *testing.T) {
-			var oracle PreimageOracle
-			if strings.HasPrefix(f.Name(), "oracle") {
-				oracle = staticOracle(t, []byte("hello world"))
-			}
+			oracle := selectOracleFixture(t, f.Name())
 			// Short-circuit early for exit_group.bin
 			exitGroup := f.Name() == "exit_group.bin"
 
@@ -269,3 +266,34 @@ func staticOracle(t *testing.T, preimageData []byte) *testOracle {
 		},
 	}
 }
+
+func staticPrecompileOracle(t *testing.T, preimageData []byte, result []byte) *testOracle {
+	return &testOracle{
+		hint: func(v []byte) {},
+		getPreimage: func(k [32]byte) []byte {
+			switch k[0] {
+			case byte(preimage.Keccak256KeyType):
+				if k != preimage.Keccak256Key(crypto.Keccak256Hash(preimageData)).PreimageKey() {
+					t.Fatalf("invalid preimage request for %x", k)
+				}
+				return preimageData
+			case byte(preimage.KZGPointEvaluationKeyType):
+				if k != preimage.KZGPointEvaluationKey(crypto.Keccak256Hash(preimageData)).PreimageKey() {
+					t.Fatalf("invalid preimage request for %x", k)
+				}
+				return result
+			}
+			panic("unreachable")
+		},
+	}
+}
+
+func selectOracleFixture(t *testing.T, programName string) PreimageOracle {
+	if strings.HasPrefix(programName, "oracle_kzg") {
+		return staticPrecompileOracle(t, []byte("hello world"), []byte{0x1})
+	} else if strings.HasPrefix(programName, "oracle") {
+		return staticOracle(t, []byte("hello world"))
+	} else {
+		return nil
+	}
+}

op-challenger/game/fault/agent.go

@@ -70,7 +70,15 @@ func (a *Agent) Act(ctx context.Context) error {
 		if action.Type == types.ActionTypeStep {
 			containsOracleData := action.OracleData != nil
 			isLocal := containsOracleData && action.OracleData.IsLocal
-			actionLog = actionLog.New("prestate", common.Bytes2Hex(action.PreState), "proof", common.Bytes2Hex(action.ProofData), "containsOracleData", containsOracleData, "isLocalPreimage", isLocal)
+			actionLog = actionLog.New(
+				"prestate", common.Bytes2Hex(action.PreState),
+				"proof", common.Bytes2Hex(action.ProofData),
+				"containsOracleData", containsOracleData,
+				"isLocalPreimage", isLocal,
+			)
+			if action.OracleData != nil {
+				actionLog = actionLog.New("oracleKey", common.Bytes2Hex(action.OracleData.OracleKey))
+			}
 		} else {
 			actionLog = actionLog.New("value", action.Value)
 		}

op-challenger/game/fault/contracts/oracle.go

@@ -20,25 +20,25 @@ import (
 )
 
 const (
-	methodInitLPP                        = "initLPP"
-	methodAddLeavesLPP                   = "addLeavesLPP"
-	methodSqueezeLPP                     = "squeezeLPP"
-	methodLoadKeccak256PreimagePart      = "loadKeccak256PreimagePart"
-	methodLoadSha256PreimagePart         = "loadSha256PreimagePart"
-	methodLoadBlobPreimagePart           = "loadBlobPreimagePart"
-	methodLoadKZGPointEvaluationPreimage = "loadKZGPointEvaluationPreimage"
-	methodProposalCount                  = "proposalCount"
-	methodProposals                      = "proposals"
-	methodProposalMetadata               = "proposalMetadata"
-	methodProposalBlocksLen              = "proposalBlocksLen"
-	methodProposalBlocks                 = "proposalBlocks"
-	methodPreimagePartOk                 = "preimagePartOk"
-	methodMinProposalSize                = "minProposalSize"
-	methodChallengeFirstLPP              = "challengeFirstLPP"
-	methodChallengeLPP                   = "challengeLPP"
-	methodChallengePeriod                = "challengePeriod"
-	methodGetTreeRootLPP                 = "getTreeRootLPP"
-	methodMinBondSizeLPP                 = "MIN_BOND_SIZE"
+	methodInitLPP                            = "initLPP"
+	methodAddLeavesLPP                       = "addLeavesLPP"
+	methodSqueezeLPP                         = "squeezeLPP"
+	methodLoadKeccak256PreimagePart          = "loadKeccak256PreimagePart"
+	methodLoadSha256PreimagePart             = "loadSha256PreimagePart"
+	methodLoadBlobPreimagePart               = "loadBlobPreimagePart"
+	methodLoadKZGPointEvaluationPreimagePart = "loadKZGPointEvaluationPreimagePart"
+	methodProposalCount                      = "proposalCount"
+	methodProposals                          = "proposals"
+	methodProposalMetadata                   = "proposalMetadata"
+	methodProposalBlocksLen                  = "proposalBlocksLen"
+	methodProposalBlocks                     = "proposalBlocks"
+	methodPreimagePartOk                     = "preimagePartOk"
+	methodMinProposalSize                    = "minProposalSize"
+	methodChallengeFirstLPP                  = "challengeFirstLPP"
+	methodChallengeLPP                       = "challengeLPP"
+	methodChallengePeriod                    = "challengePeriod"
+	methodGetTreeRootLPP                     = "getTreeRootLPP"
+	methodMinBondSizeLPP                     = "MIN_BOND_SIZE"
 )
 
 var (
@@ -108,7 +108,7 @@ func (c *PreimageOracleContract) AddGlobalDataTx(data *types.PreimageOracleData)
 			new(big.Int).SetUint64(uint64(data.OracleOffset)))
 		return call.ToTxCandidate()
 	case preimage.KZGPointEvaluationKeyType:
-		call := c.contract.Call(methodLoadKZGPointEvaluationPreimage, data.GetPreimageWithoutSize())
+		call := c.contract.Call(methodLoadKZGPointEvaluationPreimagePart, new(big.Int).SetUint64(uint64(data.OracleOffset)), data.GetPreimageWithoutSize())
 		return call.ToTxCandidate()
 	default:
 		return txmgr.TxCandidate{}, fmt.Errorf("%w: %v", ErrUnsupportedKeyType, keyType)

op-challenger/game/fault/contracts/oracle_test.go

@@ -74,8 +74,9 @@ func TestPreimageOracleContract_AddGlobalDataTx(t *testing.T) {
 	t.Run("KZGPointEvaluation", func(t *testing.T) {
 		stubRpc, oracle := setupPreimageOracleTest(t)
 		input := testutils.RandomData(rand.New(rand.NewSource(23)), 200)
-		data := types.NewPreimageOracleKZGPointEvaluationData(common.Hash{byte(preimage.KZGPointEvaluationKeyType), 0xcc}.Bytes(), input)
-		stubRpc.SetResponse(oracleAddr, methodLoadKZGPointEvaluationPreimage, batching.BlockLatest, []interface{}{
+		data := types.NewPreimageOracleData(common.Hash{byte(preimage.KZGPointEvaluationKeyType), 0xcc}.Bytes(), input, uint32(545))
+		stubRpc.SetResponse(oracleAddr, methodLoadKZGPointEvaluationPreimagePart, batching.BlockLatest, []interface{}{
+			new(big.Int).SetUint64(uint64(data.OracleOffset)),
 			data.GetPreimageWithoutSize(),
 		}, nil)
 

op-challenger/game/fault/trace/cannon/preimage.go

@@ -109,7 +109,7 @@ func (l *preimageLoader) loadKZGPointEvaluationPreimage(proof *proofData) (*type
 		return nil, fmt.Errorf("failed to get key preimage: %w", err)
 	}
 	inputWithLength := lengthPrefixed(input)
-	return types.NewPreimageOracleKZGPointEvaluationData(proof.OracleKey, inputWithLength), nil
+	return types.NewPreimageOracleData(proof.OracleKey, inputWithLength, proof.OracleOffset), nil
 }
 
 func lengthPrefixed(data []byte) []byte {

op-challenger/game/fault/trace/cannon/preimage_test.go

@@ -172,7 +172,7 @@ func TestPreimageLoader_KZGPointEvaluationPreimage(t *testing.T) {
 		actual, err := loader.LoadPreimage(proof)
 		require.NoError(t, err)
 		inputWithLength := lengthPrefixed(input)
-		expected := types.NewPreimageOracleKZGPointEvaluationData(proof.OracleKey, inputWithLength)
+		expected := types.NewPreimageOracleData(proof.OracleKey, inputWithLength, proof.OracleOffset)
 		require.Equal(t, expected, actual)
 	})
 }

op-challenger/game/fault/types/types.go

@@ -79,15 +79,6 @@ func NewPreimageOracleBlobData(key []byte, data []byte, offset uint32, fieldInde
 	}
 }
 
-func NewPreimageOracleKZGPointEvaluationData(key []byte, input []byte) *PreimageOracleData {
-	return &PreimageOracleData{
-		IsLocal:      false,
-		OracleKey:    key,
-		oracleData:   input,
-		OracleOffset: 0,
-	}
-}
-
 // StepCallData encapsulates the data needed to perform a step.
 type StepCallData struct {
 	ClaimIndex uint64

op-e2e/e2eutils/disputegame/output_game_helper.go

@@ -582,7 +582,7 @@ func (g *OutputGameHelper) uploadPreimage(ctx context.Context, data *types.Preim
 	var tx *gethtypes.Transaction
 	switch data.OracleKey[0] {
 	case byte(preimage.KZGPointEvaluationKeyType):
-		tx, err = boundOracle.LoadKZGPointEvaluationPreimage(g.opts, data.GetPreimageWithoutSize())
+		tx, err = boundOracle.LoadKZGPointEvaluationPreimagePart(g.opts, new(big.Int).SetUint64(uint64(data.OracleOffset)), data.GetPreimageWithoutSize())
 	default:
 		tx, err = boundOracle.LoadKeccak256PreimagePart(g.opts, new(big.Int).SetUint64(uint64(data.OracleOffset)), data.GetPreimageWithoutSize())
 	}

op-e2e/faultproofs/output_cannon_test.go

@@ -3,6 +3,7 @@ package faultproofs
 import (
 	"context"
 	"fmt"
+	"math/big"
 	"testing"
 
 	"github.com/ethereum-optimism/optimism/op-challenger/game/fault/trace/cannon"
@@ -13,6 +14,7 @@ import (
 	"github.com/ethereum-optimism/optimism/op-e2e/e2eutils/disputegame/preimage"
 	"github.com/ethereum-optimism/optimism/op-e2e/e2eutils/wait"
 	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/rpc"
 	"github.com/stretchr/testify/require"
 )
 
@@ -298,7 +300,6 @@ func TestOutputCannonStepWithPreimage(t *testing.T) {
 }
 
 func TestOutputCannonStepWithKZGPointEvaluation(t *testing.T) {
-	t.Skip("TODO: Fix flaky test")
 	op_e2e.InitParallel(t, op_e2e.UsesCannon)
 
 	testPreimageStep := func(t *testing.T, preloadPreimage bool) {
@@ -308,6 +309,12 @@ func TestOutputCannonStepWithKZGPointEvaluation(t *testing.T) {
 		sys, _ := startFaultDisputeSystem(t, withEcotone())
 		t.Cleanup(sys.Close)
 
+		// NOTE: Flake prevention
+		// Ensure that the L1 origin including the point eval tx isn't on the genesis epoch.
+		safeBlock, err := sys.Clients["sequencer"].BlockByNumber(ctx, big.NewInt(int64(rpc.SafeBlockNumber)))
+		require.NoError(t, err)
+		require.NoError(t, wait.ForSafeBlock(ctx, sys.RollupClient("sequencer"), safeBlock.NumberU64()+3))
+
 		receipt := sendKZGPointEvaluationTx(t, sys, "sequencer", sys.Cfg.Secrets.Alice)
 		precompileBlock := receipt.BlockNumber
 		t.Logf("KZG Point Evaluation block number: %d", precompileBlock)

packages/contracts-bedrock/snapshots/abi/PreimageOracle.json

@@ -316,13 +316,18 @@
   },
   {
     "inputs": [
+      {
+        "internalType": "uint256",
+        "name": "_partOffset",
+        "type": "uint256"
+      },
       {
         "internalType": "bytes",
         "name": "_input",
         "type": "bytes"
       }
     ],
-    "name": "loadKZGPointEvaluationPreimage",
+    "name": "loadKZGPointEvaluationPreimagePart",
     "outputs": [],
     "stateMutability": "nonpayable",
     "type": "function"

packages/contracts-bedrock/src/cannon/PreimageOracle.sol

@@ -334,10 +334,11 @@ contract PreimageOracle is IPreimageOracle {
     }
 
     /// @inheritdoc IPreimageOracle
-    function loadKZGPointEvaluationPreimage(bytes calldata _input) external {
+    function loadKZGPointEvaluationPreimagePart(uint256 _partOffset, bytes calldata _input) external {
         // Prior to Cancun activation, the blob preimage precompile is not available.
         if (block.timestamp < CANCUN_ACTIVATION) revert CancunNotActive();
 
+        bytes32 res;
         bytes32 key;
         bytes32 part;
         assembly {
@@ -353,7 +354,7 @@ contract PreimageOracle is IPreimageOracle {
 
             // Verify the KZG proof by calling the point evaluation precompile.
             // Capture the verification result
-            part :=
+            res :=
                 staticcall(
                     gas(), // forward all gas
                     0x0A, // point evaluation precompile address
@@ -362,13 +363,21 @@ contract PreimageOracle is IPreimageOracle {
                     0x00, // output ptr
                     0x00 // output size
                 )
-            // "part" will be 0 on error, and 1 on success, of the KZG Point-evaluation precompile call
+            // "res" will be 0 on error, and 1 on success, of the KZG Point-evaluation precompile call
             // We do have to shift it to the left-most byte of the bytes32 however, since we only read that byte.
-            part := shl(248, part)
+            res := shl(248, res)
+
+            // Reuse the `ptr` to store the preimage part including size prefix.
+            // put size as big-endian uint64 at the start of pre-image
+            mstore(ptr, shl(192, 1))
+            ptr := add(ptr, 0x08)
+            mstore(ptr, res)
+            // compute part given ofset
+            part := mload(add(sub(ptr, 0x08), _partOffset))
         }
-        // the part offset is always 0
-        preimagePartOk[key][0] = true;
-        preimageParts[key][0] = part;
+        preimagePartOk[key][_partOffset] = true;
+        preimageParts[key][_partOffset] = part;
+        // size is always 1
         preimageLengths[key] = 1;
     }
 

packages/contracts-bedrock/src/cannon/interfaces/IPreimageOracle.sol

@@ -72,5 +72,5 @@ interface IPreimageOracle {
 
     /// @notice Prepares a point evaluation precompile result to be read by the keccak256 of its input.
     /// @param _input The point evaluation precompile input.
-    function loadKZGPointEvaluationPreimage(bytes calldata _input) external;
+    function loadKZGPointEvaluationPreimagePart(uint256 _partOffset, bytes calldata _input) external;
 }