blockprod, replay: reject transactions that violate account lock rules

pkg/blockprod/bank.go

@@ -151,6 +151,7 @@ const (
 	ForgeDroppedCost
 	ForgeDroppedExecution
 	ForgeDroppedAlreadyProcessed
+	ForgeDroppedAccountLocks
 )
 
 func (r ForgeResult) String() string {
@@ -169,6 +170,8 @@ func (r ForgeResult) String() string {
 		return "dropped_execution"
 	case ForgeDroppedAlreadyProcessed:
 		return "dropped_already_processed"
+	case ForgeDroppedAccountLocks:
+		return "dropped_account_locks"
 	default:
 		return "unknown"
 	}
@@ -222,6 +225,17 @@ func (b *WorkingBank) forgeTransaction(tx *solana.Transaction, wireSize int, pre
 		b.RebateSchedule(wireSize)
 		return ForgeDroppedParse, costmodel.ExceedNone
 	}
+	// A transaction that repeats an account or exceeds the account lock limit
+	// can never be processed, and a block containing one is invalid. The check
+	// is static, so do it before the bank lock.
+	var lockFeats *features.Features
+	if b.slotCtx != nil {
+		lockFeats = b.slotCtx.Features
+	}
+	if txErr := replay.ValidateAccountLocks(tx, lockFeats); txErr != nil {
+		b.RebateSchedule(wireSize)
+		return ForgeDroppedAccountLocks, costmodel.ExceedNone
+	}
 	ancestorAlreadyProcessed := false
 	var ancestorStatusErr error
 	if b.ancestorStatuses == nil {

pkg/blockprod/bank_test.go

@@ -243,6 +243,79 @@ func TestWorkingBankRejectsUnsupportedV1WithoutCharging(t *testing.T) {
 	assert.Zero(t, env.Bank.TxFeeAccumulator().TotalFees)
 }
 
+func resignWithPayer(t *testing.T, tx *solana.Transaction) {
+	t.Helper()
+	tx.Signatures = nil
+	payerKey := txfixture.PayerPrivateKey()
+	_, err := tx.Sign(func(key solana.PublicKey) *solana.PrivateKey {
+		if key.Equals(txfixture.PayerPubkey()) {
+			return &payerKey
+		}
+		return nil
+	})
+	require.NoError(t, err)
+}
+
+func assertForgeDroppedWithoutCharging(t *testing.T, env *TestEnv, tx *solana.Transaction, want ForgeResult) {
+	t.Helper()
+	wire, err := tx.MarshalBinary()
+	require.NoError(t, err)
+	payerBefore, err := env.SlotCtx.GetAccount(txfixture.PayerPubkey())
+	require.NoError(t, err)
+
+	result, reason := env.Bank.ForgeTransaction(tx, len(wire))
+	assert.Equal(t, want, result)
+	assert.Equal(t, costmodel.ExceedNone, reason)
+	payerAfter, err := env.SlotCtx.GetAccount(txfixture.PayerPubkey())
+	require.NoError(t, err)
+	assert.Equal(t, payerBefore.Lamports, payerAfter.Lamports)
+	assert.Empty(t, env.Bank.ForgedTransactions())
+	assert.Zero(t, env.Bank.TxFeeAccumulator().TotalFees)
+}
+
+func TestWorkingBankRejectsDuplicateAccountKeyWithoutCharging(t *testing.T) {
+	env := NewTestEnv(TestEnvConfig{})
+	defer env.Close()
+
+	// [payer, dest, system] -> [payer, dest, dest, system]: dest is named
+	// twice, which the account lock rules forbid.
+	tx := mustSignedTransfer(t, 1)
+	keys := tx.Message.AccountKeys
+	require.Len(t, keys, 3)
+	tx.Message.AccountKeys = solana.PublicKeySlice{keys[0], keys[1], keys[1], keys[2]}
+	for i := range tx.Message.Instructions {
+		if tx.Message.Instructions[i].ProgramIDIndex == 2 {
+			tx.Message.Instructions[i].ProgramIDIndex = 3
+		}
+	}
+	resignWithPayer(t, tx)
+
+	assertForgeDroppedWithoutCharging(t, env, tx, ForgeDroppedAccountLocks)
+}
+
+func TestWorkingBankRejectsTooManyAccountLocksWithoutCharging(t *testing.T) {
+	env := NewTestEnv(TestEnvConfig{})
+	defer env.Close()
+	require.False(t, env.SlotCtx.Features.IsActive(features.IncreaseTxAccountLockLimit))
+
+	// Three static keys plus 62 looked-up addresses name 65 accounts, one
+	// more than the lock limit without IncreaseTxAccountLockLimit.
+	tx := mustSignedTransfer(t, 1)
+	_, err := tx.Message.SetVersion(solana.MessageVersionV0)
+	require.NoError(t, err)
+	indexes := make(solana.Uint8SliceAsNum, 62)
+	for i := range indexes {
+		indexes[i] = uint8(i)
+	}
+	tx.Message.AddressTableLookups = solana.MessageAddressTableLookupSlice{{
+		AccountKey:      solana.NewWallet().PublicKey(),
+		ReadonlyIndexes: indexes,
+	}}
+	resignWithPayer(t, tx)
+
+	assertForgeDroppedWithoutCharging(t, env, tx, ForgeDroppedAccountLocks)
+}
+
 func TestWorkingBankAcceptsTransferThatDrainsPayerToZero(t *testing.T) {
 	env := NewTestEnv(TestEnvConfig{})
 	defer env.Close()

pkg/blockprod/scheduler/scheduler.go

@@ -31,6 +31,7 @@ type Stats struct {
 	EvictedCapacity         uint64
 	DroppedExpired          uint64
 	DroppedAlreadyProcessed uint64
+	DroppedAccountLocks     uint64
 	Accepted                uint64
 	DroppedNoBank           uint64
 	DroppedCost             uint64
@@ -362,6 +363,8 @@ func (s *Scheduler) recordForge(result blockprod.ForgeResult, reason costmodel.E
 		s.stats.DroppedExecution++
 	case blockprod.ForgeDroppedAlreadyProcessed:
 		s.stats.DroppedAlreadyProcessed++
+	case blockprod.ForgeDroppedAccountLocks:
+		s.stats.DroppedAccountLocks++
 	default:
 		s.stats.DroppedNoBank++
 	}

pkg/forge/forge.go

@@ -24,6 +24,7 @@ type Stats struct {
 	DroppedCost             uint64
 	DroppedExecution        uint64
 	DroppedAlreadyProcessed uint64
+	DroppedAccountLocks     uint64
 	DroppedBlockCost        uint64
 	DroppedAccountCost      uint64
 	DroppedAllocCost        uint64
@@ -79,6 +80,8 @@ func (s *Sink) Receive(pkt packet.Packet) {
 		s.stats.DroppedExecution++
 	case blockprod.ForgeDroppedAlreadyProcessed:
 		s.stats.DroppedAlreadyProcessed++
+	case blockprod.ForgeDroppedAccountLocks:
+		s.stats.DroppedAccountLocks++
 	default:
 		s.stats.DroppedNoBank++
 	}

pkg/replay/account_locks.go

@@ -0,0 +1,80 @@
+package replay
+
+import (
+	"errors"
+
+	"github.com/Overclock-Validator/mithril/pkg/features"
+	"github.com/gagliardetto/solana-go"
+)
+
+const (
+	// MaxTxAccountLocks is the protocol's account lock limit once
+	// IncreaseTxAccountLockLimit is active.
+	MaxTxAccountLocks = 128
+	// legacyTxAccountLockLimit applies before IncreaseTxAccountLockLimit.
+	legacyTxAccountLockLimit = 64
+	// accountLockSetThreshold is the key count from which a set beats a
+	// pairwise scan for duplicate detection.
+	accountLockSetThreshold = 32
+)
+
+var (
+	TxErrTooManyAccountLocks = errors.New("TxErrTooManyAccountLocks")
+	TxErrAccountLoadedTwice  = errors.New("TxErrAccountLoadedTwice")
+)
+
+// TransactionAccountLockLimit returns how many accounts a transaction may
+// lock under f.
+func TransactionAccountLockLimit(f *features.Features) int {
+	if f != nil && f.IsActive(features.IncreaseTxAccountLockLimit) {
+		return MaxTxAccountLocks
+	}
+	return legacyTxAccountLockLimit
+}
+
+// ValidateAccountLocks applies the Solana protocol's account lock rules: a
+// transaction naming more accounts than the lock limit fails with
+// TooManyAccountLocks, otherwise one naming any account twice fails with
+// AccountLoadedTwice. Such a transaction can never be processed or charged,
+// and a block containing one is invalid. solana-go only rejects duplicate
+// keys for v1 messages, so legacy and v0 messages must be checked here.
+//
+// The account count includes address-table lookups. The keys of a v0 message
+// whose lookups are not resolved yet are unknown, so only its static keys are
+// checked for duplicates; check the resolved message again once its loaded
+// addresses are appended.
+func ValidateAccountLocks(tx *solana.Transaction, f *features.Features) *TransactionError {
+	msg := &tx.Message
+	numAccounts := len(msg.AccountKeys)
+	if msg.GetVersion() == solana.MessageVersionV0 && !msg.IsResolved() {
+		numAccounts += msg.AddressTableLookups.NumLookups()
+	}
+	if numAccounts > TransactionAccountLockLimit(f) {
+		return &TransactionError{ErrorType: TransactionErrorTooManyAccountLocks, InstructionError: TxErrTooManyAccountLocks}
+	}
+	if hasDuplicateAccountKeys(msg.AccountKeys) {
+		return &TransactionError{ErrorType: TransactionErrorAccountLoadedTwice, InstructionError: TxErrAccountLoadedTwice}
+	}
+	return nil
+}
+
+func hasDuplicateAccountKeys(keys solana.PublicKeySlice) bool {
+	if len(keys) >= accountLockSetThreshold {
+		seen := make(map[solana.PublicKey]struct{}, len(keys))
+		for _, key := range keys {
+			if _, ok := seen[key]; ok {
+				return true
+			}
+			seen[key] = struct{}{}
+		}
+		return false
+	}
+	for i := 1; i < len(keys); i++ {
+		for j := 0; j < i; j++ {
+			if keys[i] == keys[j] {
+				return true
+			}
+		}
+	}
+	return false
+}

pkg/replay/account_locks_replay_test.go

@@ -0,0 +1,102 @@
+package replay
+
+import (
+	"sync"
+	"testing"
+
+	"github.com/Overclock-Validator/mithril/pkg/tpu/txfixture"
+	"github.com/gagliardetto/solana-go"
+	"github.com/stretchr/testify/require"
+)
+
+func resignTransferWithPayer(t *testing.T, tx *solana.Transaction) {
+	t.Helper()
+	tx.Signatures = nil
+	payerKey := txfixture.PayerPrivateKey()
+	_, err := tx.Sign(func(key solana.PublicKey) *solana.PrivateKey {
+		if key.Equals(txfixture.PayerPubkey()) {
+			return &payerKey
+		}
+		return nil
+	})
+	require.NoError(t, err)
+}
+
+// duplicateAccountKeyTransfer names the destination twice:
+// [payer, dest, system] -> [payer, dest, dest, system].
+func duplicateAccountKeyTransfer(t *testing.T, seq uint64) *solana.Transaction {
+	t.Helper()
+	tx, err := solana.TransactionFromBytes(txfixture.MustSignedTransferWire(seq))
+	require.NoError(t, err)
+	keys := tx.Message.AccountKeys
+	require.Len(t, keys, 3)
+	tx.Message.AccountKeys = solana.PublicKeySlice{keys[0], keys[1], keys[1], keys[2]}
+	for i := range tx.Message.Instructions {
+		if tx.Message.Instructions[i].ProgramIDIndex == 2 {
+			tx.Message.Instructions[i].ProgramIDIndex = 3
+		}
+	}
+	resignTransferWithPayer(t, tx)
+	return tx
+}
+
+// tooManyAccountLocksTransfer names 3 static keys plus 62 looked-up addresses,
+// one account more than the lock limit without IncreaseTxAccountLockLimit.
+func tooManyAccountLocksTransfer(t *testing.T, seq uint64) *solana.Transaction {
+	t.Helper()
+	tx, err := solana.TransactionFromBytes(txfixture.MustSignedTransferWire(seq))
+	require.NoError(t, err)
+	_, err = tx.Message.SetVersion(solana.MessageVersionV0)
+	require.NoError(t, err)
+	indexes := make(solana.Uint8SliceAsNum, 62)
+	for i := range indexes {
+		indexes[i] = uint8(i)
+	}
+	tx.Message.AddressTableLookups = solana.MessageAddressTableLookupSlice{{
+		AccountKey:      solana.NewWallet().PublicKey(),
+		ReadonlyIndexes: indexes,
+	}}
+	resignTransferWithPayer(t, tx)
+	return tx
+}
+
+func TestProcessTransactionRejectsAccountLockViolationsWithoutCharging(t *testing.T) {
+	cases := []struct {
+		name string
+		tx   func(*testing.T) *solana.Transaction
+		want error
+	}{
+		{"duplicate account key", func(t *testing.T) *solana.Transaction { return duplicateAccountKeyTransfer(t, 11) }, TxErrAccountLoadedTwice},
+		{"too many account locks", func(t *testing.T) *solana.Transaction { return tooManyAccountLocksTransfer(t, 12) }, TxErrTooManyAccountLocks},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			slotCtx, cleanup := newCommitTestSlotCtx()
+			defer cleanup()
+			payerBefore, err := slotCtx.GetAccount(txfixture.PayerPubkey())
+			require.NoError(t, err)
+
+			var sigverify sync.WaitGroup
+			feeInfo, cu, err := ProcessTransaction(slotCtx, &sigverify, tc.tx(t), nil, nil, nil, false)
+			sigverify.Wait()
+
+			require.ErrorIs(t, err, tc.want)
+			require.Nil(t, feeInfo, "a transaction that cannot be processed must not return a fee")
+			require.Zero(t, cu)
+			payerAfter, err := slotCtx.GetAccount(txfixture.PayerPubkey())
+			require.NoError(t, err)
+			require.Equal(t, payerBefore.Lamports, payerAfter.Lamports)
+		})
+	}
+}
+
+func TestStreamingGroupRejectsDuplicateAccountKeys(t *testing.T) {
+	for _, workers := range []int{0, 4} {
+		env := newGroupExecutionEnv(t, workers, 10_000_000)
+		defer env.cleanup()
+		txs := transferTransactions(t, 2, 1)
+		txs[1] = duplicateAccountKeyTransfer(t, 3)
+		err := env.exec.executeTransactionGroup(txs, nil, false)
+		require.ErrorIs(t, err, TxErrAccountLoadedTwice, "workers=%d", workers)
+	}
+}