runtime: read the nonce account from parent state when rolling back a failed transaction

pkg/blockprod/bank_nonce_test.go

@@ -0,0 +1,87 @@
+package blockprod
+
+import (
+	"fmt"
+	"math"
+	"testing"
+
+	"github.com/Overclock-Validator/mithril/pkg/accounts"
+	"github.com/Overclock-Validator/mithril/pkg/accountsdb"
+	"github.com/Overclock-Validator/mithril/pkg/addresses"
+	"github.com/Overclock-Validator/mithril/pkg/costmodel"
+	"github.com/Overclock-Validator/mithril/pkg/replay"
+	"github.com/Overclock-Validator/mithril/pkg/sealevel"
+	"github.com/Overclock-Validator/mithril/pkg/tpu/txfixture"
+	"github.com/gagliardetto/solana-go"
+	computebudget "github.com/gagliardetto/solana-go/programs/compute-budget"
+	"github.com/gagliardetto/solana-go/programs/system"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+type parentAccounts map[solana.PublicKey]*accounts.Account
+
+func (p parentAccounts) GetAccount(_ uint64, pubkey solana.PublicKey) (*accounts.Account, error) {
+	if acct, ok := p[pubkey]; ok {
+		return acct.Clone(), nil
+	}
+	return nil, fmt.Errorf("account %s not found", pubkey)
+}
+
+func TestWorkingBankFeesOnlyDurableNonceAdvancesNonceFromParentState(t *testing.T) {
+	env := NewTestEnv(TestEnvConfig{})
+	defer env.Close()
+	nonceKey := solana.PublicKey{0xD7}
+	durableNonce := solana.Hash{0xAA}
+	nonceState := sealevel.NonceStateVersions{
+		Type: sealevel.NonceVersionCurrent,
+		Current: sealevel.NonceData{
+			IsInitialized: true,
+			Authority:     txfixture.PayerPubkey(),
+			DurableNonce:  durableNonce,
+			FeeCalculator: sealevel.FeeCalculator{LamportsPerSignature: 5000},
+		},
+	}
+	nonceData, err := nonceState.Marshal()
+	require.NoError(t, err)
+	// Like a live leader bank, the slot holds no copy of the nonce account;
+	// it is only in parent state.
+	env.SlotCtx.AccountsDb = &accountsdb.AccountsDb{}
+	env.SlotCtx.UnrootedRead = parentAccounts{nonceKey: {
+		Key: nonceKey, Lamports: 1_500_000, Owner: addresses.SystemProgramAddr,
+		Data: nonceData, RentEpoch: math.MaxUint64,
+	}}
+	_, err = env.SlotCtx.GetAccount(nonceKey)
+	require.Error(t, err)
+
+	tx, err := solana.NewTransaction(
+		[]solana.Instruction{
+			system.NewAdvanceNonceAccountInstruction(nonceKey, solana.SysVarRecentBlockHashesPubkey, txfixture.PayerPubkey()).Build(),
+			// A one-byte loaded-data limit makes this a processable fees-only
+			// load failure, which the leader already includes.
+			computebudget.NewSetLoadedAccountsDataSizeLimitInstruction(1).Build(),
+			system.NewTransferInstruction(1, txfixture.PayerPubkey(), txfixture.DestPubkey()).Build(),
+		},
+		durableNonce,
+		solana.TransactionPayer(txfixture.PayerPubkey()),
+	)
+	require.NoError(t, err)
+	payerKey := txfixture.PayerPrivateKey()
+	_, err = tx.Sign(func(solana.PublicKey) *solana.PrivateKey { return &payerKey })
+	require.NoError(t, err)
+	preview := replay.LoadAndExecuteTransaction(replay.LoadAndExecuteTransactionInput{
+		SlotCtx: env.SlotCtx, Transaction: tx, LeanResult: true,
+	})
+	require.NotNil(t, preview.ProcessingResult.TransactionError)
+	require.Equal(t, replay.TransactionErrorMaxLoadedAccountsDataSizeExceeded, preview.ProcessingResult.TransactionError.ErrorType)
+
+	result, reason := env.Bank.ForgeTransaction(tx, 0)
+	require.Equal(t, ForgeAccepted, result)
+	require.Equal(t, costmodel.ExceedNone, reason)
+	require.Contains(t, env.SlotCtx.ModifiedAccts, nonceKey, "replay advances the nonce of a failed durable-nonce transaction")
+	advanced, err := env.SlotCtx.GetAccount(nonceKey)
+	require.NoError(t, err)
+	state, err := sealevel.UnmarshalNonceStateVersions(advanced.Data)
+	require.NoError(t, err)
+	assert.NotEqual(t, durableNonce, solana.Hash(state.State().DurableNonce))
+}

pkg/replay/failed_nonce_test.go

@@ -0,0 +1,119 @@
+package replay
+
+import (
+	"encoding/binary"
+	"errors"
+	"math"
+	"testing"
+
+	"github.com/Overclock-Validator/mithril/pkg/accounts"
+	"github.com/Overclock-Validator/mithril/pkg/addresses"
+	"github.com/Overclock-Validator/mithril/pkg/sealevel"
+	"github.com/Overclock-Validator/mithril/pkg/tpu/txfixture"
+	"github.com/gagliardetto/solana-go"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+type parentStateReader map[solana.PublicKey]*accounts.Account
+
+func (r parentStateReader) GetAccount(_ uint64, pubkey solana.PublicKey) (*accounts.Account, error) {
+	if acct, ok := r[pubkey]; ok {
+		return acct.Clone(), nil
+	}
+	return nil, errors.New("account not found")
+}
+
+func TestHandleFailedTxAdvancesNonceReadFromParentState(t *testing.T) {
+	slotCtx, cleanup := newCommitTestSlotCtx()
+	defer cleanup()
+	previousRecent := sealevel.SysvarCache.RecentBlockHashes.Sysvar
+	emptyRecent := sealevel.SysvarRecentBlockhashes{}
+	sealevel.SysvarCache.RecentBlockHashes.Sysvar = &emptyRecent
+	defer func() { sealevel.SysvarCache.RecentBlockHashes.Sysvar = previousRecent }()
+
+	authority := txfixture.PayerPubkey()
+	nonceKey := solana.PublicKey{0xD6}
+	durableNonce := [32]byte{0xAA}
+	nonceState := sealevel.NonceStateVersions{
+		Type: sealevel.NonceVersionCurrent,
+		Current: sealevel.NonceData{
+			IsInitialized: true,
+			Authority:     authority,
+			DurableNonce:  durableNonce,
+			FeeCalculator: sealevel.FeeCalculator{LamportsPerSignature: 5000},
+		},
+	}
+	nonceData, err := nonceState.Marshal()
+	require.NoError(t, err)
+	// Like a leader bank, the slot has not written the nonce account yet; only
+	// parent state has it.
+	parent := parentStateReader{nonceKey: {
+		Key: nonceKey, Lamports: 1_500_000, Owner: addresses.SystemProgramAddr,
+		Data: nonceData, RentEpoch: math.MaxUint64,
+	}}
+	slotCtx.UnrootedRead = parent
+	_, err = slotCtx.GetAccount(nonceKey)
+	require.Error(t, err)
+	slotCtx.LastBlockhash = [32]byte{0x77}
+
+	tx, err := solana.TransactionFromBytes(txfixture.MustSignedTransferWire(0))
+	require.NoError(t, err)
+	tx.Message.RecentBlockhash = durableNonce
+	data := make([]byte, 4)
+	binary.LittleEndian.PutUint32(data, sealevel.SystemProgramInstrTypeAdvanceNonceAccount)
+	advance := sealevel.Instruction{
+		ProgramId: addresses.SystemProgramAddr,
+		Accounts: []sealevel.AccountMeta{
+			{Pubkey: nonceKey, IsWritable: true},
+			{Pubkey: authority, IsSigner: true},
+		},
+		Data: data,
+	}
+
+	feeInfo, _ := handleFailedTx(slotCtx, tx, []sealevel.Instruction{advance}, &sealevel.ComputeBudgetLimits{},
+		sealevel.InstrErrInvalidArgument, nil)
+	require.NotNil(t, feeInfo)
+	require.Contains(t, slotCtx.ModifiedAccts, nonceKey, "a failed durable-nonce transaction advances its nonce")
+	advanced, err := slotCtx.GetAccount(nonceKey)
+	require.NoError(t, err)
+	state, err := sealevel.UnmarshalNonceStateVersions(advanced.Data)
+	require.NoError(t, err)
+	assert.NotEqual(t, durableNonce, [32]byte(state.State().DurableNonce))
+	assert.Equal(t, uint64(1_500_000), advanced.Lamports)
+	original, err := parent.GetAccount(slotCtx.Slot, nonceKey)
+	require.NoError(t, err)
+	assert.Equal(t, nonceData, original.Data, "parent state is not modified")
+}
+
+func TestHandleFailedTxSkipsNonceMissingEverywhere(t *testing.T) {
+	slotCtx, cleanup := newCommitTestSlotCtx()
+	defer cleanup()
+	previousRecent := sealevel.SysvarCache.RecentBlockHashes.Sysvar
+	emptyRecent := sealevel.SysvarRecentBlockhashes{}
+	sealevel.SysvarCache.RecentBlockHashes.Sysvar = &emptyRecent
+	defer func() { sealevel.SysvarCache.RecentBlockHashes.Sysvar = previousRecent }()
+	slotCtx.UnrootedRead = parentStateReader{}
+	nonceKey := solana.PublicKey{0xD8}
+
+	tx, err := solana.TransactionFromBytes(txfixture.MustSignedTransferWire(0))
+	require.NoError(t, err)
+	tx.Message.RecentBlockhash = [32]byte{0xAA}
+	data := make([]byte, 4)
+	binary.LittleEndian.PutUint32(data, sealevel.SystemProgramInstrTypeAdvanceNonceAccount)
+	advance := sealevel.Instruction{
+		ProgramId: addresses.SystemProgramAddr,
+		Accounts: []sealevel.AccountMeta{
+			{Pubkey: nonceKey, IsWritable: true},
+			{Pubkey: txfixture.PayerPubkey(), IsSigner: true},
+		},
+		Data: data,
+	}
+
+	feeInfo, _ := handleFailedTx(slotCtx, tx, []sealevel.Instruction{advance}, &sealevel.ComputeBudgetLimits{},
+		sealevel.InstrErrInvalidArgument, nil)
+	require.NotNil(t, feeInfo, "the fee is still charged")
+	assert.NotContains(t, slotCtx.ModifiedAccts, nonceKey)
+	_, err = slotCtx.GetAccount(nonceKey)
+	assert.Error(t, err, "no nonce account is created")
+}

pkg/sealevel/blockhash_nonce.go

@@ -54,7 +54,15 @@ func MaybeAdvanceNonceAccountForFailedTx(slotCtx *SlotCtx, tx *solana.Transactio
 	noncePk := instr.Accounts[0].Pubkey
 	nonceAcct, err := slotCtx.GetAccount(noncePk)
 	if err != nil {
-		return solana.PublicKey{}, false
+		// A leader bank starts empty, so an account the slot has not written
+		// yet is read from parent state, as handleFailedTx does for the payer.
+		if slotCtx.UnrootedRead == nil && slotCtx.AccountsDb == nil {
+			return solana.PublicKey{}, false
+		}
+		nonceAcct, err = slotCtx.GetAccountFromAccountsDb(noncePk)
+		if err != nil {
+			return solana.PublicKey{}, false
+		}
 	}
 
 	nonceStateVersions, err := UnmarshalNonceStateVersions(nonceAcct.Data)