blockprod: 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,73 @@
+package replay
+
+import (
+	"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
+)
+
+// 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}
+	}
+	if hasDuplicateAccountKeys(msg.AccountKeys) {
+		return &TransactionError{ErrorType: TransactionErrorAccountLoadedTwice}
+	}
+	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_test.go

@@ -0,0 +1,117 @@
+package replay
+
+import (
+	"testing"
+
+	"github.com/Overclock-Validator/mithril/pkg/features"
+	"github.com/gagliardetto/solana-go"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func accountLockTestKeys(n int) solana.PublicKeySlice {
+	keys := make(solana.PublicKeySlice, n)
+	for i := range keys {
+		keys[i] = solana.NewWallet().PublicKey()
+	}
+	return keys
+}
+
+func accountLockTestTx(keys solana.PublicKeySlice) *solana.Transaction {
+	return &solana.Transaction{Message: solana.Message{AccountKeys: keys}}
+}
+
+func accountLockTestV0Tx(t *testing.T, static solana.PublicKeySlice, lookups int) *solana.Transaction {
+	t.Helper()
+	tx := accountLockTestTx(static)
+	_, err := tx.Message.SetVersion(solana.MessageVersionV0)
+	require.NoError(t, err)
+	indexes := make(solana.Uint8SliceAsNum, lookups)
+	for i := range indexes {
+		indexes[i] = uint8(i)
+	}
+	tx.Message.AddressTableLookups = solana.MessageAddressTableLookupSlice{{
+		AccountKey:      solana.NewWallet().PublicKey(),
+		ReadonlyIndexes: indexes,
+	}}
+	return tx
+}
+
+func lockFeatures(increaseLimit bool) *features.Features {
+	f := features.NewFeaturesDefault()
+	if increaseLimit {
+		f.EnableFeature(features.IncreaseTxAccountLockLimit, 0)
+	}
+	return f
+}
+
+func requireLockError(t *testing.T, want TransactionErrorType, got *TransactionError) {
+	t.Helper()
+	require.NotNil(t, got)
+	assert.Equal(t, want, got.ErrorType, "got %s", got.ErrorType.String())
+}
+
+func TestTransactionAccountLockLimit(t *testing.T) {
+	assert.Equal(t, 64, TransactionAccountLockLimit(nil))
+	assert.Equal(t, 64, TransactionAccountLockLimit(lockFeatures(false)))
+	assert.Equal(t, MaxTxAccountLocks, TransactionAccountLockLimit(lockFeatures(true)))
+}
+
+func TestValidateAccountLocksAcceptsUniqueKeys(t *testing.T) {
+	for _, n := range []int{1, 3, accountLockSetThreshold - 1, accountLockSetThreshold, 64} {
+		assert.Nil(t, ValidateAccountLocks(accountLockTestTx(accountLockTestKeys(n)), lockFeatures(false)), "n=%d", n)
+	}
+}
+
+func TestValidateAccountLocksRejectsDuplicateKeys(t *testing.T) {
+	// Both duplicate-detection strategies: pairwise below the set threshold, set above.
+	for _, n := range []int{2, 3, accountLockSetThreshold - 1, accountLockSetThreshold, 64} {
+		keys := accountLockTestKeys(n)
+		keys[n-1] = keys[0]
+		requireLockError(t, TransactionErrorAccountLoadedTwice, ValidateAccountLocks(accountLockTestTx(keys), lockFeatures(true)))
+	}
+}
+
+func TestValidateAccountLocksLimitFollowsFeature(t *testing.T) {
+	tx := accountLockTestTx(accountLockTestKeys(65))
+	requireLockError(t, TransactionErrorTooManyAccountLocks, ValidateAccountLocks(tx, lockFeatures(false)))
+	assert.Nil(t, ValidateAccountLocks(tx, lockFeatures(true)))
+
+	tx = accountLockTestTx(accountLockTestKeys(MaxTxAccountLocks + 1))
+	requireLockError(t, TransactionErrorTooManyAccountLocks, ValidateAccountLocks(tx, lockFeatures(true)))
+}
+
+func TestValidateAccountLocksTooManyTakesPrecedence(t *testing.T) {
+	keys := accountLockTestKeys(65)
+	keys[64] = keys[0]
+	requireLockError(t, TransactionErrorTooManyAccountLocks, ValidateAccountLocks(accountLockTestTx(keys), lockFeatures(false)))
+}
+
+func TestValidateAccountLocksCountsUnresolvedLookups(t *testing.T) {
+	tx := accountLockTestV0Tx(t, accountLockTestKeys(2), 63)
+	requireLockError(t, TransactionErrorTooManyAccountLocks, ValidateAccountLocks(tx, lockFeatures(false)))
+	assert.Nil(t, ValidateAccountLocks(tx, lockFeatures(true)))
+
+	tx = accountLockTestV0Tx(t, accountLockTestKeys(2), 62)
+	assert.Nil(t, ValidateAccountLocks(tx, lockFeatures(false)))
+}
+
+func TestValidateAccountLocksChecksUnresolvedStaticKeys(t *testing.T) {
+	static := accountLockTestKeys(3)
+	static[2] = static[1]
+	tx := accountLockTestV0Tx(t, static, 1)
+	requireLockError(t, TransactionErrorAccountLoadedTwice, ValidateAccountLocks(tx, lockFeatures(true)))
+}
+
+func TestValidateAccountLocksChecksResolvedLookupKeys(t *testing.T) {
+	static := accountLockTestKeys(2)
+	tx := accountLockTestV0Tx(t, static, 2)
+	table := tx.Message.AddressTableLookups[0].AccountKey
+	// The second looked-up address repeats a static key.
+	addresses := solana.PublicKeySlice{solana.NewWallet().PublicKey(), static[1]}
+	require.NoError(t, tx.Message.SetAddressTables(map[solana.PublicKey]solana.PublicKeySlice{table: addresses}))
+	require.NoError(t, tx.Message.ResolveLookups())
+	require.Len(t, tx.Message.AccountKeys, 4)
+
+	requireLockError(t, TransactionErrorAccountLoadedTwice, ValidateAccountLocks(tx, lockFeatures(true)))
+}