feat(server/v2/cometbft): Implement necessary handlers and helpers fo…

…r vote extensions (#22830)

runtime/v2/builder.go

@@ -30,6 +30,7 @@ type AppBuilder[T transaction.Tx] struct {
 	branch      func(state store.ReaderMap) store.WriterMap
 	txValidator func(ctx context.Context, tx T) error
 	postTxExec  func(ctx context.Context, tx T, success bool) error
+	preblocker  func(ctx context.Context, txs []T, mmPreblocker func() error) error
 }
 
 // RegisterModules registers the provided modules with the module manager.
@@ -95,11 +96,22 @@ func (a *AppBuilder[T]) Build(opts ...AppBuilderOption[T]) (*App[T], error) {
 
 	endBlocker, valUpdate := a.app.moduleManager.EndBlock()
 
+	preblockerFn := func(ctx context.Context, txs []T) error {
+		if a.preblocker != nil {
+			return a.preblocker(ctx, txs, func() error {
+				return a.app.moduleManager.PreBlocker()(ctx, txs)
+			})
+		}
+
+		// if there is no preblocker set, call the module manager's preblocker directly
+		return a.app.moduleManager.PreBlocker()(ctx, txs)
+	}
+
 	stf, err := stf.New[T](
 		a.app.logger.With("module", "stf"),
 		a.app.msgRouterBuilder,
 		a.app.queryRouterBuilder,
-		a.app.moduleManager.PreBlocker(),
+		preblockerFn,
 		a.app.moduleManager.BeginBlock(),
 		endBlocker,
 		a.txValidator,
@@ -219,3 +231,17 @@ func AppBuilderWithPostTxExec[T transaction.Tx](
 		a.postTxExec = postTxExec
 	}
 }
+
+// AppBuilderWithPreblocker sets logic that will be executed before each block.
+// mmPreblocker can be used to call module manager's preblocker, so that it can be
+// called before or after depending on the app's logic.
+// This is especially useful when implementing vote extensions.
+func AppBuilderWithPreblocker[T transaction.Tx](
+	preblocker func(
+		ctx context.Context, txs []T, mmPreblocker func() error,
+	) error,
+) AppBuilderOption[T] {
+	return func(a *AppBuilder[T]) {
+		a.preblocker = preblocker
+	}
+}

server/v2/cometbft/abci.go

@@ -68,7 +68,7 @@ type consensus[T transaction.Tx] struct {
 
 	prepareProposalHandler handlers.PrepareHandler[T]
 	processProposalHandler handlers.ProcessHandler[T]
-	verifyVoteExt          handlers.VerifyVoteExtensionhandler
+	verifyVoteExt          handlers.VerifyVoteExtensionHandler
 	extendVote             handlers.ExtendVoteHandler
 	checkTxHandler         handlers.CheckTxHandler[T]
 
@@ -142,7 +142,7 @@ func (c *consensus[T]) Info(ctx context.Context, _ *abciproto.InfoRequest) (*abc
 	// if height is 0, we dont know the consensus params
 	var appVersion uint64 = 0
 	if version > 0 {
-		cp, err := c.GetConsensusParams(ctx)
+		cp, err := GetConsensusParams(ctx, c.app)
 		// if the consensus params are not found, we set the app version to 0
 		// in the case that the start version is > 0
 		if cp == nil || errors.Is(err, collections.ErrNotFound) {
@@ -387,7 +387,7 @@ func (c *consensus[T]) PrepareProposal(
 		LastCommit:      toCoreExtendedCommitInfo(req.LocalLastCommit),
 	})
 
-	txs, err := c.prepareProposalHandler(ciCtx, c.app, c.appCodecs.TxCodec, req)
+	txs, err := c.prepareProposalHandler(ciCtx, c.app, c.appCodecs.TxCodec, req, c.chainID)
 	if err != nil {
 		return nil, err
 	}
@@ -433,7 +433,7 @@ func (c *consensus[T]) ProcessProposal(
 		LastCommit:      toCoreCommitInfo(req.ProposedLastCommit),
 	})
 
-	err := c.processProposalHandler(ciCtx, c.app, c.appCodecs.TxCodec, req)
+	err := c.processProposalHandler(ciCtx, c.app, c.appCodecs.TxCodec, req, c.chainID)
 	if err != nil {
 		c.logger.Error("failed to process proposal", "height", req.Height, "time", req.Time, "hash", fmt.Sprintf("%X", req.Hash), "err", err)
 		return &abciproto.ProcessProposalResponse{
@@ -533,7 +533,7 @@ func (c *consensus[T]) FinalizeBlock(
 
 	c.lastCommittedHeight.Store(req.Height)
 
-	cp, err := c.GetConsensusParams(ctx) // we get the consensus params from the latest state because we committed state above
+	cp, err := GetConsensusParams(ctx, c.app) // we get the consensus params from the latest state because we committed state above
 	if err != nil {
 		return nil, err
 	}
@@ -600,7 +600,7 @@ func (c *consensus[T]) Commit(ctx context.Context, _ *abciproto.CommitRequest) (
 
 	c.snapshotManager.SnapshotIfApplicable(lastCommittedHeight)
 
-	cp, err := c.GetConsensusParams(ctx)
+	cp, err := GetConsensusParams(ctx, c.app)
 	if err != nil {
 		return nil, err
 	}
@@ -619,7 +619,7 @@ func (c *consensus[T]) VerifyVoteExtension(
 ) (*abciproto.VerifyVoteExtensionResponse, error) {
 	// If vote extensions are not enabled, as a safety precaution, we return an
 	// error.
-	cp, err := c.GetConsensusParams(ctx)
+	cp, err := GetConsensusParams(ctx, c.app)
 	if err != nil {
 		return nil, err
 	}
@@ -658,7 +658,7 @@ func (c *consensus[T]) VerifyVoteExtension(
 func (c *consensus[T]) ExtendVote(ctx context.Context, req *abciproto.ExtendVoteRequest) (*abciproto.ExtendVoteResponse, error) {
 	// If vote extensions are not enabled, as a safety precaution, we return an
 	// error.
-	cp, err := c.GetConsensusParams(ctx)
+	cp, err := GetConsensusParams(ctx, c.app)
 	if err != nil {
 		return nil, err
 	}
@@ -702,6 +702,8 @@ func decodeTxs[T transaction.Tx](logger log.Logger, rawTxs [][]byte, codec trans
 		if err != nil {
 			// do not return an error here, as we want to deliver the block even if some txs are invalid
 			logger.Debug("failed to decode tx", "err", err)
+			txs[i] = RawTx(rawTx).(T) // allows getting the raw bytes down the line
+			continue
 		}
 		txs[i] = tx
 	}

server/v2/cometbft/handlers/defaults.go

@@ -32,7 +32,7 @@ func NewDefaultProposalHandler[T transaction.Tx](mp mempool.Mempool[T]) *Default
 }
 
 func (h *DefaultProposalHandler[T]) PrepareHandler() PrepareHandler[T] {
-	return func(ctx context.Context, app AppManager[T], codec transaction.Codec[T], req *abci.PrepareProposalRequest) ([]T, error) {
+	return func(ctx context.Context, app AppManager[T], codec transaction.Codec[T], req *abci.PrepareProposalRequest, chainID string) ([]T, error) {
 		var maxBlockGas uint64
 
 		res, err := app.Query(ctx, 0, &consensustypes.QueryParamsRequest{})
@@ -98,7 +98,7 @@ func (h *DefaultProposalHandler[T]) PrepareHandler() PrepareHandler[T] {
 }
 
 func (h *DefaultProposalHandler[T]) ProcessHandler() ProcessHandler[T] {
-	return func(ctx context.Context, app AppManager[T], codec transaction.Codec[T], req *abci.ProcessProposalRequest) error {
+	return func(ctx context.Context, app AppManager[T], codec transaction.Codec[T], req *abci.ProcessProposalRequest, chainID string) error {
 		// If the mempool is nil we simply return ACCEPT,
 		// because PrepareProposal may have included txs that could fail verification.
 		_, isNoOp := h.mempool.(mempool.NoOpMempool[T])
@@ -174,15 +174,15 @@ func decodeTxs[T transaction.Tx](codec transaction.Codec[T], txsBz [][]byte) []T
 // NoOpPrepareProposal defines a no-op PrepareProposal handler. It will always
 // return the transactions sent by the client's request.
 func NoOpPrepareProposal[T transaction.Tx]() PrepareHandler[T] {
-	return func(ctx context.Context, app AppManager[T], codec transaction.Codec[T], req *abci.PrepareProposalRequest) ([]T, error) {
+	return func(ctx context.Context, app AppManager[T], codec transaction.Codec[T], req *abci.PrepareProposalRequest, chainID string) ([]T, error) {
 		return decodeTxs(codec, req.Txs), nil
 	}
 }
 
 // NoOpProcessProposal defines a no-op ProcessProposal Handler. It will always
 // return ACCEPT.
 func NoOpProcessProposal[T transaction.Tx]() ProcessHandler[T] {
-	return func(context.Context, AppManager[T], transaction.Codec[T], *abci.ProcessProposalRequest) error {
+	return func(context.Context, AppManager[T], transaction.Codec[T], *abci.ProcessProposalRequest, string) error {
 		return nil
 	}
 }
@@ -197,7 +197,7 @@ func NoOpExtendVote() ExtendVoteHandler {
 
 // NoOpVerifyVoteExtensionHandler defines a no-op VerifyVoteExtension handler. It
 // will always return an ACCEPT status with no error.
-func NoOpVerifyVoteExtensionHandler() VerifyVoteExtensionhandler {
+func NoOpVerifyVoteExtensionHandler() VerifyVoteExtensionHandler {
 	return func(context.Context, store.ReaderMap, *abci.VerifyVoteExtensionRequest) (*abci.VerifyVoteExtensionResponse, error) {
 		return &abci.VerifyVoteExtensionResponse{Status: abci.VERIFY_VOTE_EXTENSION_STATUS_ACCEPT}, nil
 	}

server/v2/cometbft/handlers/handlers.go

@@ -13,16 +13,16 @@ import (
 type (
 	// PrepareHandler passes in the list of Txs that are being proposed. The app can then do stateful operations
 	// over the list of proposed transactions. It can return a modified list of txs to include in the proposal.
-	PrepareHandler[T transaction.Tx] func(context.Context, AppManager[T], transaction.Codec[T], *abci.PrepareProposalRequest) ([]T, error)
+	PrepareHandler[T transaction.Tx] func(ctx context.Context, app AppManager[T], cdc transaction.Codec[T], req *abci.PrepareProposalRequest, chainID string) ([]T, error)
 
 	// ProcessHandler is a function that takes a list of transactions and returns a boolean and an error.
 	// If the verification of a transaction fails, the boolean is false and the error is non-nil.
-	ProcessHandler[T transaction.Tx] func(context.Context, AppManager[T], transaction.Codec[T], *abci.ProcessProposalRequest) error
+	ProcessHandler[T transaction.Tx] func(ctx context.Context, app AppManager[T], cdc transaction.Codec[T], req *abci.ProcessProposalRequest, chainID string) error
 
-	// VerifyVoteExtensionhandler is a function type that handles the verification of a vote extension request.
+	// VerifyVoteExtensionHandler is a function type that handles the verification of a vote extension request.
 	// It takes a context, a store reader map, and a request to verify a vote extension.
 	// It returns a response to verify the vote extension and an error if any.
-	VerifyVoteExtensionhandler func(context.Context, store.ReaderMap, *abci.VerifyVoteExtensionRequest) (*abci.VerifyVoteExtensionResponse, error)
+	VerifyVoteExtensionHandler func(context.Context, store.ReaderMap, *abci.VerifyVoteExtensionRequest) (*abci.VerifyVoteExtensionResponse, error)
 
 	// ExtendVoteHandler is a function type that handles the extension of a vote.
 	// It takes a context, a store reader map, and a request to extend a vote.

server/v2/cometbft/options.go

@@ -20,7 +20,7 @@ type ServerOptions[T transaction.Tx] struct {
 	PrepareProposalHandler     handlers.PrepareHandler[T]
 	ProcessProposalHandler     handlers.ProcessHandler[T]
 	CheckTxHandler             handlers.CheckTxHandler[T]
-	VerifyVoteExtensionHandler handlers.VerifyVoteExtensionhandler
+	VerifyVoteExtensionHandler handlers.VerifyVoteExtensionHandler
 	ExtendVoteHandler          handlers.ExtendVoteHandler
 	KeygenF                    keyGenF