feat: add vote extension support (#1641)

<!--
Please read and fill out this form before submitting your PR.

Please make sure you have reviewed our contributors guide before
submitting your
first PR.

NOTE: PR titles should follow semantic commits:
https://www.conventionalcommits.org/en/v1.0.0/
-->

## Overview

Resolves #1593

<!-- 
Please provide an explanation of the PR, including the appropriate
context,
background, goal, and rationale. If there is an issue with this
information,
please provide a tl;dr and link the issue. 

Ex: Closes #<issue number>
-->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

- **New Features**
- Enhanced voting extension capabilities in the blockchain management
system, including saving and retrieving extended commit information.
- Added a new test function to verify the functionality of extended
commits in the blockchain store.
- Introduced a new field in the state initialization to enable vote
extensions at a specified block height.

- **Enhancements**
- Updated the block creation process to include extended commit
information, improving the integration with voting extensions.
- Streamlined the `SetHeight` method in the store for better performance
and simplicity.

- **Bug Fixes**
- Adjusted test cases to accommodate new parameters related to extended
commit information, ensuring consistency and reliability in block
creation tests.

- **Documentation**
- Updated code comments and versioning information to reflect new
functionalities and changes in mock functions.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Author: tzdybal

block/manager.go

@@ -10,16 +10,18 @@ import (
 	"sync/atomic"
 	"time"
 
-	goheaderstore "github.com/celestiaorg/go-header/store"
 	abci "github.com/cometbft/cometbft/abci/types"
 	cmcrypto "github.com/cometbft/cometbft/crypto"
 	"github.com/cometbft/cometbft/crypto/merkle"
+	cmproto "github.com/cometbft/cometbft/proto/tendermint/types"
 	"github.com/cometbft/cometbft/proxy"
 	cmtypes "github.com/cometbft/cometbft/types"
 	ds "github.com/ipfs/go-datastore"
 	"github.com/libp2p/go-libp2p/core/crypto"
 	pkgErrors "github.com/pkg/errors"
 
+	goheaderstore "github.com/celestiaorg/go-header/store"
+
 	"github.com/rollkit/rollkit/config"
 	"github.com/rollkit/rollkit/da"
 	"github.com/rollkit/rollkit/mempool"
@@ -811,7 +813,11 @@ func (m *Manager) publishBlock(ctx context.Context) error {
 		block = pendingBlock
 	} else {
 		m.logger.Info("Creating and publishing block", "height", newHeight)
-		block, err = m.createBlock(newHeight, lastCommit, lastHeaderHash)
+		extendedCommit, err := m.getExtendedCommit(ctx, height)
+		if err != nil {
+			return fmt.Errorf("failed to load extended commit for height %d: %w", height, err)
+		}
+		block, err = m.createBlock(newHeight, lastCommit, lastHeaderHash, extendedCommit)
 		if err != nil {
 			return err
 		}
@@ -863,6 +869,10 @@ func (m *Manager) publishBlock(ctx context.Context) error {
 		return err
 	}
 
+	if err := m.processVoteExtension(ctx, block, newHeight); err != nil {
+		return err
+	}
+
 	// set the commit to current block's signed header
 	block.SignedHeader.Commit = *commit
 	// Validate the created block before storing
@@ -926,6 +936,60 @@ func (m *Manager) publishBlock(ctx context.Context) error {
 	return nil
 }
 
+func (m *Manager) processVoteExtension(ctx context.Context, block *types.Block, newHeight uint64) error {
+	if !m.voteExtensionEnabled(newHeight) {
+		return nil
+	}
+
+	extension, err := m.executor.ExtendVote(ctx, block)
+	if err != nil {
+		return fmt.Errorf("error returned by ExtendVote: %w", err)
+	}
+	sign, err := m.proposerKey.Sign(extension)
+	if err != nil {
+		return fmt.Errorf("error signing vote extension: %w", err)
+	}
+	extendedCommit := buildExtendedCommit(block, extension, sign)
+	err = m.store.SaveExtendedCommit(ctx, newHeight, extendedCommit)
+	if err != nil {
+		return fmt.Errorf("failed to save extended commit: %w", err)
+	}
+	return nil
+}
+
+func (m *Manager) voteExtensionEnabled(newHeight uint64) bool {
+	enableHeight := m.lastState.ConsensusParams.Abci.VoteExtensionsEnableHeight
+	return m.lastState.ConsensusParams.Abci != nil && enableHeight != 0 && uint64(enableHeight) <= newHeight
+}
+
+func (m *Manager) getExtendedCommit(ctx context.Context, height uint64) (abci.ExtendedCommitInfo, error) {
+	emptyExtendedCommit := abci.ExtendedCommitInfo{}
+	if !m.voteExtensionEnabled(height) || height <= uint64(m.genesis.InitialHeight) {
+		return emptyExtendedCommit, nil
+	}
+	extendedCommit, err := m.store.GetExtendedCommit(ctx, height)
+	if err != nil {
+		return emptyExtendedCommit, err
+	}
+	return *extendedCommit, nil
+}
+
+func buildExtendedCommit(block *types.Block, extension []byte, sign []byte) *abci.ExtendedCommitInfo {
+	extendedCommit := &abci.ExtendedCommitInfo{
+		Round: 0,
+		Votes: []abci.ExtendedVoteInfo{{
+			Validator: abci.Validator{
+				Address: block.SignedHeader.Validators.GetProposer().Address,
+				Power:   block.SignedHeader.Validators.GetProposer().VotingPower,
+			},
+			VoteExtension:      extension,
+			ExtensionSignature: sign,
+			BlockIdFlag:        cmproto.BlockIDFlagCommit,
+		}},
+	}
+	return extendedCommit
+}
+
 func (m *Manager) recordMetrics(block *types.Block) {
 	m.metrics.NumTxs.Set(float64(len(block.Data.Txs)))
 	m.metrics.TotalTxs.Add(float64(len(block.Data.Txs)))
@@ -1057,10 +1121,10 @@ func (m *Manager) getLastBlockTime() time.Time {
 	return m.lastState.LastBlockTime
 }
 
-func (m *Manager) createBlock(height uint64, lastCommit *types.Commit, lastHeaderHash types.Hash) (*types.Block, error) {
+func (m *Manager) createBlock(height uint64, lastCommit *types.Commit, lastHeaderHash types.Hash, extendedCommit abci.ExtendedCommitInfo) (*types.Block, error) {
 	m.lastStateMtx.RLock()
 	defer m.lastStateMtx.RUnlock()
-	return m.executor.CreateBlock(height, lastCommit, lastHeaderHash, m.lastState)
+	return m.executor.CreateBlock(height, lastCommit, extendedCommit, lastHeaderHash, m.lastState)
 }
 
 func (m *Manager) applyBlock(ctx context.Context, block *types.Block) (types.State, *abci.ResponseFinalizeBlock, error) {

node/full_node_test.go

@@ -11,26 +11,27 @@ import (
 	"testing"
 	"time"
 
+	abci "github.com/cometbft/cometbft/abci/types"
 	cmconfig "github.com/cometbft/cometbft/config"
+	cmproto "github.com/cometbft/cometbft/proto/tendermint/types"
 	"github.com/cometbft/cometbft/proxy"
-
-	goDA "github.com/rollkit/go-da"
-	"github.com/rollkit/rollkit/config"
-	test "github.com/rollkit/rollkit/test/log"
+	cmtypes "github.com/cometbft/cometbft/types"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
 
-	abci "github.com/cometbft/cometbft/abci/types"
 	"github.com/libp2p/go-libp2p/core/crypto"
 	"github.com/libp2p/go-libp2p/core/peer"
 
 	testutils "github.com/celestiaorg/utils/test"
 
+	goDA "github.com/rollkit/go-da"
 	"github.com/rollkit/rollkit/block"
+	"github.com/rollkit/rollkit/config"
 	"github.com/rollkit/rollkit/da"
 	"github.com/rollkit/rollkit/mempool"
+	test "github.com/rollkit/rollkit/test/log"
 	"github.com/rollkit/rollkit/test/mocks"
 	"github.com/rollkit/rollkit/types"
 )
@@ -262,6 +263,63 @@ func TestPendingBlocks(t *testing.T) {
 
 }
 
+func TestVoteExtension(t *testing.T) {
+	require := require.New(t)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	const voteExtensionEnableHeight = 5
+	const expectedExtension = "vote extension from height %d"
+
+	// Create & configure node with app. Get signing key for mock functions.
+	app := &mocks.Application{}
+	app.On("InitChain", mock.Anything, mock.Anything).Return(&abci.ResponseInitChain{}, nil)
+	node, signingKey := createAggregatorWithApp(ctx, app, voteExtensionEnableHeight, t)
+	require.NotNil(node)
+	require.NotNil(signingKey)
+
+	prepareProposalVoteExtChecker := func(_ context.Context, req *abci.RequestPrepareProposal) (*abci.ResponsePrepareProposal, error) {
+		if req.Height <= voteExtensionEnableHeight {
+			require.Empty(req.LocalLastCommit.Votes)
+		} else {
+			require.Len(req.LocalLastCommit.Votes, 1)
+			extendedCommit := req.LocalLastCommit.Votes[0]
+			require.NotNil(extendedCommit)
+			require.Equal(extendedCommit.BlockIdFlag, cmproto.BlockIDFlagCommit)
+			// during PrepareProposal at height h, vote extensions from previous block (h-1) is available
+			require.Equal([]byte(fmt.Sprintf(expectedExtension, req.Height-1)), extendedCommit.VoteExtension)
+			require.NotNil(extendedCommit.Validator)
+			require.NotNil(extendedCommit.Validator.Address)
+			require.NotEmpty(extendedCommit.ExtensionSignature)
+			ok, err := signingKey.GetPublic().Verify(extendedCommit.VoteExtension, extendedCommit.ExtensionSignature)
+			require.NoError(err)
+			require.True(ok)
+		}
+		return &abci.ResponsePrepareProposal{
+			Txs: req.Txs,
+		}, nil
+	}
+
+	voteExtension := func(_ context.Context, req *abci.RequestExtendVote) (*abci.ResponseExtendVote, error) {
+		return &abci.ResponseExtendVote{
+			VoteExtension: []byte(fmt.Sprintf(expectedExtension, req.Height)),
+		}, nil
+	}
+
+	app.On("Commit", mock.Anything, mock.Anything).Return(&abci.ResponseCommit{}, nil)
+	app.On("PrepareProposal", mock.Anything, mock.Anything).Return(prepareProposalVoteExtChecker)
+	app.On("ProcessProposal", mock.Anything, mock.Anything).Return(&abci.ResponseProcessProposal{Status: abci.ResponseProcessProposal_ACCEPT}, nil)
+	app.On("FinalizeBlock", mock.Anything, mock.Anything).Return(finalizeBlockResponse)
+	app.On("ExtendVote", mock.Anything, mock.Anything).Return(voteExtension)
+	require.NotNil(app)
+
+	require.NoError(node.Start())
+	require.NoError(waitForAtLeastNBlocks(node, 10, Store))
+	require.NoError(node.Stop())
+	app.AssertExpectations(t)
+}
+
 func createAggregatorWithPersistence(ctx context.Context, dbPath string, dalc *da.DAClient, t *testing.T) (Node, *mocks.Application) {
 	t.Helper()
 
@@ -302,6 +360,46 @@ func createAggregatorWithPersistence(ctx context.Context, dbPath string, dalc *d
 	return fullNode, app
 }
 
+func createAggregatorWithApp(ctx context.Context, app abci.Application, voteExtensionEnableHeight int64, t *testing.T) (Node, crypto.PrivKey) {
+	t.Helper()
+
+	key, _, _ := crypto.GenerateEd25519Key(rand.Reader)
+	genesis, genesisValidatorKey := types.GetGenesisWithPrivkey()
+	genesis.ConsensusParams = &cmtypes.ConsensusParams{
+		Block:     cmtypes.DefaultBlockParams(),
+		Evidence:  cmtypes.DefaultEvidenceParams(),
+		Validator: cmtypes.DefaultValidatorParams(),
+		Version:   cmtypes.DefaultVersionParams(),
+		ABCI:      cmtypes.ABCIParams{VoteExtensionsEnableHeight: voteExtensionEnableHeight},
+	}
+	signingKey, err := types.PrivKeyToSigningKey(genesisValidatorKey)
+	require.NoError(t, err)
+
+	node, err := NewNode(
+		ctx,
+		config.NodeConfig{
+			DAAddress:   MockDAAddress,
+			DANamespace: MockDANamespace,
+			Aggregator:  true,
+			BlockManagerConfig: config.BlockManagerConfig{
+				BlockTime:   100 * time.Millisecond,
+				DABlockTime: 300 * time.Millisecond,
+			},
+			Light: false,
+		},
+		key,
+		signingKey,
+		proxy.NewLocalClientCreator(app),
+		genesis,
+		DefaultMetricsProvider(cmconfig.DefaultInstrumentationConfig()),
+		test.NewFileLoggerCustom(t, test.TempLogFileName(t, "")),
+	)
+	require.NoError(t, err)
+	require.NotNil(t, node)
+
+	return node, signingKey
+}
+
 // setupMockApplication initializes a mock application
 func setupMockApplication() *mocks.Application {
 	app := &mocks.Application{}

state/executor.go

@@ -92,7 +92,7 @@ func (e *BlockExecutor) InitChain(genesis *cmtypes.GenesisDoc) (*abci.ResponseIn
 }
 
 // CreateBlock reaps transactions from mempool and builds a block.
-func (e *BlockExecutor) CreateBlock(height uint64, lastCommit *types.Commit, lastHeaderHash types.Hash, state types.State) (*types.Block, error) {
+func (e *BlockExecutor) CreateBlock(height uint64, lastCommit *types.Commit, lastExtendedCommit abci.ExtendedCommitInfo, lastHeaderHash types.Hash, state types.State) (*types.Block, error) {
 	maxBytes := state.ConsensusParams.Block.MaxBytes
 	emptyMaxBytes := maxBytes == -1
 	if emptyMaxBytes {
@@ -140,12 +140,9 @@ func (e *BlockExecutor) CreateBlock(height uint64, lastCommit *types.Commit, las
 	rpp, err := e.proxyApp.PrepareProposal(
 		context.TODO(),
 		&abci.RequestPrepareProposal{
-			MaxTxBytes: maxBytes,
-			Txs:        mempoolTxs.ToSliceOfBytes(),
-			LocalLastCommit: abci.ExtendedCommitInfo{
-				Round: 0,
-				Votes: []abci.ExtendedVoteInfo{},
-			},
+			MaxTxBytes:         maxBytes,
+			Txs:                mempoolTxs.ToSliceOfBytes(),
+			LocalLastCommit:    lastExtendedCommit,
 			Misbehavior:        []abci.Misbehavior{},
 			Height:             int64(block.Height()),
 			Time:               block.Time(),
@@ -241,6 +238,32 @@ func (e *BlockExecutor) ApplyBlock(ctx context.Context, state types.State, block
 	return state, resp, nil
 }
 
+// ExtendVote calls the ExtendVote ABCI method on the proxy app.
+func (e *BlockExecutor) ExtendVote(ctx context.Context, block *types.Block) ([]byte, error) {
+	resp, err := e.proxyApp.ExtendVote(ctx, &abci.RequestExtendVote{
+		Hash:   block.Hash(),
+		Height: int64(block.Height()),
+		Time:   block.Time(),
+		Txs:    block.Data.Txs.ToSliceOfBytes(),
+		ProposedLastCommit: abci.CommitInfo{
+			Votes: []abci.VoteInfo{{
+				Validator: abci.Validator{
+					Address: block.SignedHeader.Validators.GetProposer().Address,
+					Power:   block.SignedHeader.Validators.GetProposer().VotingPower,
+				},
+				BlockIdFlag: cmproto.BlockIDFlagCommit,
+			}},
+		},
+		Misbehavior:        nil,
+		NextValidatorsHash: block.SignedHeader.ValidatorHash,
+		ProposerAddress:    block.SignedHeader.ProposerAddress,
+	})
+	if err != nil {
+		return nil, err
+	}
+	return resp.VoteExtension, nil
+}
+
 // Commit commits the block
 func (e *BlockExecutor) Commit(ctx context.Context, state types.State, block *types.Block, resp *abci.ResponseFinalizeBlock) ([]byte, uint64, error) {
 	appHash, retainHeight, err := e.commit(ctx, state, block, resp)

state/executor_test.go

@@ -60,7 +60,7 @@ func doTestCreateBlock(t *testing.T) {
 	state.ConsensusParams.Block.MaxGas = 100000
 
 	// empty block
-	block, err := executor.CreateBlock(1, &types.Commit{}, []byte{}, state)
+	block, err := executor.CreateBlock(1, &types.Commit{}, abci.ExtendedCommitInfo{}, []byte{}, state)
 	require.NoError(err)
 	require.NotNil(block)
 	assert.Empty(block.Data.Txs)
@@ -69,7 +69,7 @@ func doTestCreateBlock(t *testing.T) {
 	// one small Tx
 	err = mpool.CheckTx([]byte{1, 2, 3, 4}, func(r *abci.ResponseCheckTx) {}, mempool.TxInfo{})
 	require.NoError(err)
-	block, err = executor.CreateBlock(2, &types.Commit{}, []byte{}, state)
+	block, err = executor.CreateBlock(2, &types.Commit{}, abci.ExtendedCommitInfo{}, []byte{}, state)
 	require.NoError(err)
 	require.NotNil(block)
 	assert.Equal(uint64(2), block.Height())
@@ -80,7 +80,7 @@ func doTestCreateBlock(t *testing.T) {
 	require.NoError(err)
 	err = mpool.CheckTx(make([]byte, 100), func(r *abci.ResponseCheckTx) {}, mempool.TxInfo{})
 	require.NoError(err)
-	block, err = executor.CreateBlock(3, &types.Commit{}, []byte{}, state)
+	block, err = executor.CreateBlock(3, &types.Commit{}, abci.ExtendedCommitInfo{}, []byte{}, state)
 	require.NoError(err)
 	require.NotNil(block)
 	assert.Len(block.Data.Txs, 2)
@@ -90,7 +90,7 @@ func doTestCreateBlock(t *testing.T) {
 	executor.maxBytes = 10
 	err = mpool.CheckTx(make([]byte, 10), func(r *abci.ResponseCheckTx) {}, mempool.TxInfo{})
 	require.NoError(err)
-	block, err = executor.CreateBlock(4, &types.Commit{}, []byte{}, state)
+	block, err = executor.CreateBlock(4, &types.Commit{}, abci.ExtendedCommitInfo{}, []byte{}, state)
 	require.NoError(err)
 	require.NotNil(block)
 	assert.Empty(block.Data.Txs)
@@ -171,7 +171,7 @@ func doTestApplyBlock(t *testing.T) {
 
 	err = mpool.CheckTx([]byte{1, 2, 3, 4}, func(r *abci.ResponseCheckTx) {}, mempool.TxInfo{})
 	require.NoError(err)
-	block, err := executor.CreateBlock(1, &types.Commit{Signatures: []types.Signature{types.Signature([]byte{1, 1, 1})}}, []byte{}, state)
+	block, err := executor.CreateBlock(1, &types.Commit{Signatures: []types.Signature{types.Signature([]byte{1, 1, 1})}}, abci.ExtendedCommitInfo{}, []byte{}, state)
 	require.NoError(err)
 	require.NotNil(block)
 	assert.Equal(uint64(1), block.Height())
@@ -201,7 +201,7 @@ func doTestApplyBlock(t *testing.T) {
 	require.NoError(mpool.CheckTx([]byte{5, 6, 7, 8, 9}, func(r *abci.ResponseCheckTx) {}, mempool.TxInfo{}))
 	require.NoError(mpool.CheckTx([]byte{1, 2, 3, 4, 5}, func(r *abci.ResponseCheckTx) {}, mempool.TxInfo{}))
 	require.NoError(mpool.CheckTx(make([]byte, 90), func(r *abci.ResponseCheckTx) {}, mempool.TxInfo{}))
-	block, err = executor.CreateBlock(2, &types.Commit{Signatures: []types.Signature{types.Signature([]byte{1, 1, 1})}}, []byte{}, newState)
+	block, err = executor.CreateBlock(2, &types.Commit{Signatures: []types.Signature{types.Signature([]byte{1, 1, 1})}}, abci.ExtendedCommitInfo{}, []byte{}, newState)
 	require.NoError(err)
 	require.NotNil(block)
 	assert.Equal(uint64(2), block.Height())