fix!: Branch antehandlers and commit to state before runMsgs

x/auth/middleware/branch_store.go

@@ -8,47 +8,72 @@ import (
 	tmtypes "github.com/tendermint/tendermint/types"
 )
 
-type branchStoreHandler struct {
-	next tx.Handler
+// WithAnteBranch creates a store branch (cache-wrapped multistore) for Antehandlers.
+//
+// Usage:
+// beginBranch, endBranch := WithAnteBranch()
+//
+// ComposeMiddlewares(
+//   beginBranch,
+//   // some middlewares representing antehandlers
+//   endBranch, // will write to state before all middlewares below
+//   // some other middlewares
+// )
+func WithAnteBranch() (tx.Middleware, tx.Middleware) {
+	// We'll store the original sdk.Context multistore, as well as a
+	// cache-wrapped multistore, inside the sdk.Context's Values (it's just
+	// a simple key-value map). We define those keys here, and pass them to the
+	// middlewares.
+	originalMSKey := sdk.ContextKey("ante-original-ms")
+	cacheMSKey := sdk.ContextKey("ante-cache-ms")
+
+	return func(h tx.Handler) tx.Handler {
+			return anteBranchBegin{next: h, originalMSKey: originalMSKey, cacheMSKey: cacheMSKey}
+		}, func(h tx.Handler) tx.Handler {
+			return anteBranchWrite{next: h, originalMSKey: originalMSKey, cacheMSKey: cacheMSKey}
+		}
 }
 
-// WithBranchedStore creates a new MultiStore branch and commits the store if the downstream
-// returned no error. It cancels writes from the failed transactions.
-func WithBranchedStore(txh tx.Handler) tx.Handler {
-	return branchStoreHandler{next: txh}
+// anteBranchBegin is the tx.Handler that creates a new branched store.
+type anteBranchBegin struct {
+	next          tx.Handler
+	originalMSKey sdk.ContextKey
+	cacheMSKey    sdk.ContextKey
 }
 
 // CheckTx implements tx.Handler.CheckTx method.
 // Do nothing during CheckTx.
-func (sh branchStoreHandler) CheckTx(ctx context.Context, req tx.Request, checkReq tx.RequestCheckTx) (tx.Response, tx.ResponseCheckTx, error) {
+func (sh anteBranchBegin) CheckTx(ctx context.Context, req tx.Request, checkReq tx.RequestCheckTx) (tx.Response, tx.ResponseCheckTx, error) {
 	return sh.next.CheckTx(ctx, req, checkReq)
 }
 
 // DeliverTx implements tx.Handler.DeliverTx method.
-func (sh branchStoreHandler) DeliverTx(ctx context.Context, req tx.Request) (tx.Response, error) {
-	return branchAndRun(ctx, req, sh.next.DeliverTx)
+func (sh anteBranchBegin) DeliverTx(ctx context.Context, req tx.Request) (tx.Response, error) {
+	return anteCreateBranch(ctx, sh.originalMSKey, sh.cacheMSKey, req, sh.next.DeliverTx)
 }
 
 // SimulateTx implements tx.Handler.SimulateTx method.
-func (sh branchStoreHandler) SimulateTx(ctx context.Context, req tx.Request) (tx.Response, error) {
-	return branchAndRun(ctx, req, sh.next.SimulateTx)
+func (sh anteBranchBegin) SimulateTx(ctx context.Context, req tx.Request) (tx.Response, error) {
+	return anteCreateBranch(ctx, sh.originalMSKey, sh.cacheMSKey, req, sh.next.SimulateTx)
 }
 
 type nextFn func(ctx context.Context, req tx.Request) (tx.Response, error)
 
-// branchAndRun creates a new Context based on the existing Context with a MultiStore branch
+// anteCreateBranch creates a new Context based on the existing Context with a MultiStore branch
 // in case message processing fails.
-func branchAndRun(ctx context.Context, req tx.Request, fn nextFn) (tx.Response, error) {
+func anteCreateBranch(ctx context.Context, originalMSKey, cacheMSKey sdk.ContextKey, req tx.Request, fn nextFn) (tx.Response, error) {
 	sdkCtx := sdk.UnwrapSDKContext(ctx)
-	runMsgCtx, branchedStore := branchStore(sdkCtx, tmtypes.Tx(req.TxBytes))
 
-	rsp, err := fn(sdk.WrapSDKContext(runMsgCtx), req)
-	if err == nil {
-		// commit storage iff no error
-		branchedStore.Write()
-	}
+	newSdkCtx, branchedStore := branchStore(sdkCtx, tmtypes.Tx(req.TxBytes))
+	// Put the stores inside the new sdk.Context, as we will need them when
+	// writing.
+	newSdkCtx = newSdkCtx.
+		WithValue(cacheMSKey, branchedStore).
+		WithValue(originalMSKey, sdkCtx.MultiStore())
 
-	return rsp, err
+	newCtx := sdk.WrapSDKContext(newSdkCtx)
+
+	return fn(newCtx, req)
 }
 
 // branchStore returns a new context based off of the provided context with
@@ -68,3 +93,105 @@ func branchStore(sdkCtx sdk.Context, tx tmtypes.Tx) (sdk.Context, sdk.CacheMulti
 
 	return sdkCtx.WithMultiStore(msCache), msCache
 }
+
+// anteBranchWrite is the tx.Handler that commits the state writes of a previously
+// created anteBranchBegin tx.Handler.
+type anteBranchWrite struct {
+	next          tx.Handler
+	originalMSKey sdk.ContextKey
+	cacheMSKey    sdk.ContextKey
+}
+
+// CheckTx implements tx.Handler.CheckTx method.
+// Do nothing during CheckTx.
+func (sh anteBranchWrite) CheckTx(ctx context.Context, req tx.Request, checkReq tx.RequestCheckTx) (tx.Response, tx.ResponseCheckTx, error) {
+	return sh.next.CheckTx(ctx, req, checkReq)
+}
+
+// DeliverTx implements tx.Handler.DeliverTx method.
+func (sh anteBranchWrite) DeliverTx(ctx context.Context, req tx.Request) (tx.Response, error) {
+	return anteWrite(ctx, sh.originalMSKey, sh.cacheMSKey, req, sh.next.DeliverTx)
+}
+
+// SimulateTx implements tx.Handler.SimulateTx method.
+func (sh anteBranchWrite) SimulateTx(ctx context.Context, req tx.Request) (tx.Response, error) {
+	return anteWrite(ctx, sh.originalMSKey, sh.cacheMSKey, req, sh.next.SimulateTx)
+}
+
+// branchAndRun creates a new Context based on the existing Context with a MultiStore branch
+// in case message processing fails.
+func anteWrite(ctx context.Context, originalMSKey, cacheMSKey sdk.ContextKey, req tx.Request, fn nextFn) (tx.Response, error) {
+	sdkCtx := sdk.UnwrapSDKContext(ctx)
+	originalStore := sdkCtx.Value(originalMSKey).(sdk.MultiStore)
+	branchedStore := sdkCtx.Value(cacheMSKey).(sdk.CacheMultiStore)
+
+	if !sdkCtx.IsZero() {
+		// At this point, newCtx.MultiStore() is a store branch, or something else
+		// replaced by the AnteHandler. We want the original multistore.
+		//
+		// Also, in the case of the tx aborting, we need to track gas consumed via
+		// the instantiated gas meter in the AnteHandler, so we update the context
+		// prior to returning.
+		sdkCtx = sdkCtx.WithMultiStore(originalStore)
+	}
+
+	branchedStore.Write()
+
+	// We don't need references to the 2 multistores anymore.
+	sdkCtx = sdkCtx.WithValue(originalMSKey, nil).WithValue(cacheMSKey, nil)
+
+	return fn(sdk.WrapSDKContext(sdkCtx), req)
+}
+
+// WithRunMsgsBranch creates a store branch (cache store) for running Msgs.
+//
+// Usage:
+//
+// ComposeMiddlewares(
+//   // some middlewares
+//
+//   // Creates a new MultiStore branch, discards downstream writes if the downstream returns error.
+//   WithRunMsgsBranch,
+//   // optionally, some other middlewares who should also discard writes when
+//   // middleware fails.
+// )
+func WithRunMsgsBranch(h tx.Handler) tx.Handler {
+	return runMsgsBranch{next: h}
+}
+
+// runMsgsBranch is the tx.Handler that commits the state writes of a previously
+// created anteBranchBegin tx.Handler.
+type runMsgsBranch struct {
+	next tx.Handler
+}
+
+// CheckTx implements tx.Handler.CheckTx method.
+// Do nothing during CheckTx.
+func (sh runMsgsBranch) CheckTx(ctx context.Context, req tx.Request, checkReq tx.RequestCheckTx) (tx.Response, tx.ResponseCheckTx, error) {
+	return sh.next.CheckTx(ctx, req, checkReq)
+}
+
+// DeliverTx implements tx.Handler.DeliverTx method.
+func (sh runMsgsBranch) DeliverTx(ctx context.Context, req tx.Request) (tx.Response, error) {
+	return branchAndRun(ctx, req, sh.next.DeliverTx)
+}
+
+// SimulateTx implements tx.Handler.SimulateTx method.
+func (sh runMsgsBranch) SimulateTx(ctx context.Context, req tx.Request) (tx.Response, error) {
+	return branchAndRun(ctx, req, sh.next.SimulateTx)
+}
+
+// branchAndRun creates a new Context based on the existing Context with a MultiStore branch
+// in case message processing fails.
+func branchAndRun(ctx context.Context, req tx.Request, fn nextFn) (tx.Response, error) {
+	sdkCtx := sdk.UnwrapSDKContext(ctx)
+	runMsgCtx, branchedStore := branchStore(sdkCtx, tmtypes.Tx(req.TxBytes))
+
+	rsp, err := fn(sdk.WrapSDKContext(runMsgCtx), req)
+	if err == nil {
+		// commit storage iff no error
+		branchedStore.Write()
+	}
+
+	return rsp, err
+}

x/auth/middleware/branch_store_test.go

@@ -67,14 +67,19 @@ func (s *MWTestSuite) TestBranchStore() {
 				return tx.Response{}, nil
 			}}
 
+			beginAnteBranch, endAnteBranch := middleware.WithAnteBranch()
+
 			txHandler := middleware.ComposeMiddlewares(
 				testMsgTxHandler,
+				beginAnteBranch,
 				middleware.NewTxDecoderMiddleware(s.clientCtx.TxConfig.TxDecoder()),
 				middleware.GasTxMiddleware,
 				middleware.RecoveryTxMiddleware,
 				middleware.DeductFeeMiddleware(s.app.AccountKeeper, s.app.BankKeeper, s.app.FeeGrantKeeper, nil),
 				middleware.IncrementSequenceMiddleware(s.app.AccountKeeper),
-				middleware.WithBranchedStore,
+				endAnteBranch,
+
+				middleware.WithRunMsgsBranch,
 				middleware.ConsumeBlockGasMiddleware,
 			)
 

x/auth/middleware/middleware.go

@@ -71,8 +71,22 @@ func NewDefaultTxHandler(options TxHandlerOptions) (tx.Handler, error) {
 		return nil, sdkerrors.Wrap(sdkerrors.ErrLogic, "sign mode handler is required for middlewares")
 	}
 
+	// Antehandlers have been migrated to middlewares as part of ADR-045.
+	// However, we still call "antehandlers" the middlewares that are run before
+	// the runMsgs tx.Handler.
+	//
+	// After successful antehandlers, we actually do a state write. This
+	// `WithAnteBranch` middleware denotes where to start the store branching
+	// (i.e. before the antehanndler block) and where to write back the branch
+	// to the main store (i.e. right after the antehandler block, before runMsgs).
+	beginAnteBranch, endAnteBranch := WithAnteBranch()
+
 	return ComposeMiddlewares(
 		NewRunMsgsTxHandler(options.MsgServiceRouter, options.LegacyRouter),
+
+		// Creates a new MultiStore branch, discards downstream writes if the downstream returns error.
+		// This block of middlewares correspond to what was called "antehandlers" before.
+		beginAnteBranch,
 		NewTxDecoderMiddleware(options.TxDecoder),
 		// Set a new GasMeter on sdk.Context.
 		//
@@ -102,11 +116,14 @@ func NewDefaultTxHandler(options TxHandlerOptions) (tx.Handler, error) {
 		SigGasConsumeMiddleware(options.AccountKeeper, options.SigGasConsumer),
 		SigVerificationMiddleware(options.AccountKeeper, options.SignModeHandler),
 		IncrementSequenceMiddleware(options.AccountKeeper),
+		endAnteBranch,
+
 		// Creates a new MultiStore branch, discards downstream writes if the downstream returns error.
 		// These kinds of middlewares should be put under this:
 		// - Could return error after messages executed succesfully.
-		// - Storage writes should be discarded together when tx failed.
-		WithBranchedStore,
+		// - Storage writes below this middleware should ALWAYS be discarded
+		//   together when middleware fails, but those above can have state writes.
+		WithRunMsgsBranch,
 		// Consume block gas. All middlewares whose gas consumption after their `next` handler
 		// should be accounted for, should go below this middleware.
 		ConsumeBlockGasMiddleware,

x/auth/middleware/middleware_test.go

@@ -7,6 +7,9 @@ import (
 	"strings"
 	"testing"
 
+	"github.com/stretchr/testify/require"
+	abci "github.com/tendermint/tendermint/abci/types"
+
 	"github.com/cosmos/cosmos-sdk/crypto/keys/ed25519"
 	kmultisig "github.com/cosmos/cosmos-sdk/crypto/keys/multisig"
 	"github.com/cosmos/cosmos-sdk/crypto/keys/secp256k1"
@@ -20,8 +23,6 @@ import (
 	"github.com/cosmos/cosmos-sdk/x/auth/types"
 	"github.com/cosmos/cosmos-sdk/x/bank/testutil"
 	minttypes "github.com/cosmos/cosmos-sdk/x/mint/types"
-	"github.com/stretchr/testify/require"
-	abci "github.com/tendermint/tendermint/abci/types"
 )
 
 var testCoins = sdk.Coins{sdk.NewInt64Coin("atom", 10000000)}
@@ -877,12 +878,12 @@ func (s *MWTestSuite) TestTxHandlerSetPubKey() {
 			sdkerrors.ErrWrongSequence,
 		},
 		{
-			"make sure previous public key has been set after wrong signature",
+			"make sure previous public key has NOT been set after wrong signature",
 			func() {
 				// Make sure public key has been set, as SetPubKeyMiddleware
 				// is called before all signature verification middlewares.
 				acc1 := s.app.AccountKeeper.GetAccount(ctx, accounts[1].acc.GetAddress())
-				s.Require().Equal(acc1.GetPubKey(), accounts[1].priv.PubKey())
+				s.Require().Equal(nil, acc1.GetPubKey())
 			},
 			false,
 			false,
@@ -1153,3 +1154,41 @@ func (s *MWTestSuite) TestTxHandlerReCheck() {
 	_, _, err = s.txHandler.CheckTx(sdk.WrapSDKContext(ctx), tx.Request{Tx: testTx}, tx.RequestCheckTx{})
 	s.Require().NotNil(err, "txhandler on recheck did not fail once feePayer no longer has sufficient funds")
 }
+
+func (s *MWTestSuite) TestTxReplayFeeDeduction() {
+	ctx := s.SetupTest(false) // reset
+	txBuilder := s.clientCtx.TxConfig.NewTxBuilder()
+	accounts := s.createTestAccounts(ctx, 1, testCoins)
+
+	privs, accNums, accSeqs := []cryptotypes.PrivKey{accounts[0].priv}, []uint64{0}, []uint64{0}
+	msgs := []sdk.Msg{testdata.NewTestMsg(accounts[0].acc.GetAddress())}
+	s.Require().NoError(txBuilder.SetMsgs(msgs...))
+
+	txBuilder.SetFeeAmount(testdata.NewTestFeeAmount())
+	txBuilder.SetGasLimit(testdata.NewTestGasLimit())
+
+	testTx, txBytes, err := s.createTestTx(txBuilder, privs, accNums, accSeqs, ctx.ChainID())
+	s.Require().NoError(err)
+
+	// execute a valid tx and require that it succeeds
+	_, err = s.txHandler.DeliverTx(sdk.WrapSDKContext(ctx), tx.Request{Tx: testTx, TxBytes: txBytes})
+	s.Require().NoError(err)
+
+	// require the balance be adjusted (10_000_000 - 150 (fee) = 9_999_850)
+	balance := s.app.BankKeeper.GetBalance(ctx, accounts[0].acc.GetAddress(), testCoins[0].Denom)
+	s.Require().Equal(sdk.NewInt(9_999_850), balance.Amount)
+
+	// Execute the exact same tx, which should fail as expected. Imagine a faulty
+	// modified mempool or bad actors proposing blocks with already committed
+	// transactions.
+	_, err = s.txHandler.DeliverTx(sdk.WrapSDKContext(ctx), tx.Request{Tx: testTx, TxBytes: txBytes})
+	s.Require().Error(err)
+
+	// Require the balance DOES NOT change, even though the network already accepted
+	// the exact same message from the same sender.
+	//
+	// Note, this would be considered faulty behavior and can result in accounts
+	// being drained of their funds.
+	balance = s.app.BankKeeper.GetBalance(ctx, accounts[0].acc.GetAddress(), testCoins[0].Denom)
+	s.Require().Equal(sdk.NewInt(9_999_850), balance.Amount)
+}

x/auth/middleware/testutil_test.go

@@ -38,6 +38,7 @@ type MWTestSuite struct {
 	app       *simapp.SimApp
 	clientCtx client.Context
 	txHandler txtypes.Handler
+	mOpts     middleware.TxHandlerOptions
 }
 
 // returns context and app with params set on account keeper
@@ -80,7 +81,7 @@ func (s *MWTestSuite) SetupTest(isCheckTx bool) sdk.Context {
 			MsgResponses: []*codectypes.Any{any},
 		}, nil
 	}))
-	txHandler, err := middleware.NewDefaultTxHandler(middleware.TxHandlerOptions{
+	s.mOpts = middleware.TxHandlerOptions{
 		Debug:            s.app.Trace(),
 		MsgServiceRouter: msr,
 		LegacyRouter:     legacyRouter,
@@ -90,8 +91,11 @@ func (s *MWTestSuite) SetupTest(isCheckTx bool) sdk.Context {
 		SignModeHandler:  encodingConfig.TxConfig.SignModeHandler(),
 		SigGasConsumer:   middleware.DefaultSigVerificationGasConsumer,
 		TxDecoder:        s.clientCtx.TxConfig.TxDecoder(),
-	})
+	}
+
+	txHandler, err := middleware.NewDefaultTxHandler(s.mOpts)
 	s.Require().NoError(err)
+
 	s.txHandler = txHandler
 
 	return ctx