From 44293e1864339686a87c76822fb114c7526400f6 Mon Sep 17 00:00:00 2001
From: codchen <codchen03@gmail.com>
Date: Tue, 30 Apr 2024 23:43:22 +0800
Subject: [PATCH] associate for traceBlock

---
 evmrpc/simulate.go       | 23 ++++++++++++++++++++++-
 x/evm/ante/preprocess.go |  6 +++---
 2 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/evmrpc/simulate.go b/evmrpc/simulate.go
index 60a65a18ef..971c215427 100644
--- a/evmrpc/simulate.go
+++ b/evmrpc/simulate.go
@@ -23,6 +23,7 @@ import (
 	"github.com/ethereum/go-ethereum/params"
 	"github.com/ethereum/go-ethereum/rpc"
 	"github.com/sei-protocol/sei-chain/utils"
+	"github.com/sei-protocol/sei-chain/x/evm/ante"
 	"github.com/sei-protocol/sei-chain/x/evm/keeper"
 	"github.com/sei-protocol/sei-chain/x/evm/state"
 	"github.com/sei-protocol/sei-chain/x/evm/types"
@@ -297,7 +298,9 @@ func (b *Backend) StateAtTransaction(ctx context.Context, block *ethtypes.Block,
 			if !associatedNow {
 				return nil, vm.BlockContext{}, nil, nil, fmt.Errorf("address %s is not associated in the latest height", msg.From.Hex())
 			}
-			b.keeper.SetAddressMapping(statedb.Ctx(), seiAddr, msg.From)
+			if err := ante.NewEVMPreprocessDecorator(b.keeper, b.keeper.AccountKeeper()).AssociateAddresses(statedb.Ctx(), seiAddr, msg.From, nil); err != nil {
+				return nil, vm.BlockContext{}, nil, nil, err
+			}
 		}
 
 		if idx == txIndex {
@@ -319,6 +322,24 @@ func (b *Backend) StateAtTransaction(ctx context.Context, block *ethtypes.Block,
 func (b *Backend) StateAtBlock(ctx context.Context, block *ethtypes.Block, reexec uint64, base vm.StateDB, readOnly bool, preferDisk bool) (vm.StateDB, tracers.StateReleaseFunc, error) {
 	emptyRelease := func() {}
 	statedb := state.NewDBImpl(b.ctxProvider(block.Number().Int64()-1), b.keeper, true)
+	signer := ethtypes.MakeSigner(b.ChainConfig(), block.Number(), block.Time())
+	for _, tx := range block.Transactions() {
+		msg, _ := core.TransactionToMessage(tx, signer, block.BaseFee())
+
+		// set address association for the sender if not present. Note that here we take the shortcut
+		// of querying from the latest height with the assumption that if this tx has been processed
+		// at all then its association must be present in the latest height
+		_, associated := b.keeper.GetSeiAddress(statedb.Ctx(), msg.From)
+		if !associated {
+			seiAddr, associatedNow := b.keeper.GetSeiAddress(b.ctxProvider(LatestCtxHeight), msg.From)
+			if !associatedNow {
+				return nil, emptyRelease, fmt.Errorf("address %s is not associated in the latest height", msg.From.Hex())
+			}
+			if err := ante.NewEVMPreprocessDecorator(b.keeper, b.keeper.AccountKeeper()).AssociateAddresses(statedb.Ctx(), seiAddr, msg.From, nil); err != nil {
+				return nil, emptyRelease, err
+			}
+		}
+	}
 	return statedb, emptyRelease, nil
 }
 
diff --git a/x/evm/ante/preprocess.go b/x/evm/ante/preprocess.go
index 15eb9aaa4a..f7b41220ee 100644
--- a/x/evm/ante/preprocess.go
+++ b/x/evm/ante/preprocess.go
@@ -88,7 +88,7 @@ func (p *EVMPreprocessDecorator) AnteHandle(ctx sdk.Context, tx sdk.Tx, simulate
 				return ctx, sdkerrors.Wrap(sdkerrors.ErrInsufficientFunds, "account needs to have at least 1Sei to force association")
 			}
 		}
-		if err := p.associateAddresses(ctx, seiAddr, evmAddr, pubkey); err != nil {
+		if err := p.AssociateAddresses(ctx, seiAddr, evmAddr, pubkey); err != nil {
 			return ctx, err
 		}
 		return ctx.WithPriority(antedecorators.EVMAssociatePriority), nil // short-circuit without calling next
@@ -96,7 +96,7 @@ func (p *EVMPreprocessDecorator) AnteHandle(ctx sdk.Context, tx sdk.Tx, simulate
 		// noop; for readability
 	} else {
 		// not associatedTx and not already associated
-		if err := p.associateAddresses(ctx, seiAddr, evmAddr, pubkey); err != nil {
+		if err := p.AssociateAddresses(ctx, seiAddr, evmAddr, pubkey); err != nil {
 			return ctx, err
 		}
 		if p.evmKeeper.EthReplayConfig.Enabled {
@@ -107,7 +107,7 @@ func (p *EVMPreprocessDecorator) AnteHandle(ctx sdk.Context, tx sdk.Tx, simulate
 	return next(ctx, tx, simulate)
 }
 
-func (p *EVMPreprocessDecorator) associateAddresses(ctx sdk.Context, seiAddr sdk.AccAddress, evmAddr common.Address, pubkey cryptotypes.PubKey) error {
+func (p *EVMPreprocessDecorator) AssociateAddresses(ctx sdk.Context, seiAddr sdk.AccAddress, evmAddr common.Address, pubkey cryptotypes.PubKey) error {
 	p.evmKeeper.SetAddressMapping(ctx, seiAddr, evmAddr)
 	if acc := p.accountKeeper.GetAccount(ctx, seiAddr); acc.GetPubKey() == nil {
 		if err := acc.SetPubKey(pubkey); err != nil {