feat: Add transaction cache to prevent duplicate transaction processi…

…ng (#122)

* feat: Add Solana transaction cache to prevent replay attacks

* refactor: Rename type for transaction cache limit

* test: Add check_transaction test after clearing cache

* refactor: Simplify transaction cache update with try_mutate

* refactor: Remove unused API transaction_count

frame/solana/src/lib.rs

@@ -176,8 +176,13 @@ pub mod pallet {
 		#[pallet::no_default_bounds]
 		type GenesisTimestamp: Get<Self::Moment>;
 
+		/// Maximum scan result size in bytes.
 		#[pallet::constant]
 		type ScanResultsLimitBytes: Get<Option<u32>>;
+
+		/// Maximum number of transactions to cache for tracking processed ones.
+		#[pallet::constant]
+		type TransactionCacheLimit: Get<u32>;
 	}
 
 	pub mod config_preludes {
@@ -209,6 +214,8 @@ pub mod pallet {
 			type GenesisTimestamp = ConstU64<1584336540_000>;
 			/// Maximum scan result size in bytes.
 			type ScanResultsLimitBytes = ScanResultsLimitBytes;
+			/// Maximum number of transactions to cache for tracking processed ones.
+			type TransactionCacheLimit = ConstU32<10000>;
 		}
 	}
 
@@ -222,6 +229,8 @@ pub mod pallet {
 	pub enum Error<T> {
 		/// Failed to reallocate account data of this length
 		InvalidRealloc,
+		/// Transaction cache limit reached.
+		CacheLimitReached,
 	}
 
 	#[pallet::storage]
@@ -249,6 +258,16 @@ pub mod pallet {
 		ValueQuery,
 	>;
 
+	#[pallet::storage]
+	#[pallet::getter(fn transaction_cache)]
+	pub type TransactionCache<T: Config> = StorageMap<
+		_,
+		Twox64Concat,
+		T::Hash,
+		BoundedBTreeSet<T::Hash, T::TransactionCacheLimit>,
+		ValueQuery,
+	>;
+
 	#[pallet::genesis_config]
 	#[derive_where(Default)]
 	pub struct GenesisConfig<T: Config> {
@@ -280,10 +299,6 @@ pub mod pallet {
 		}
 	}
 
-	#[pallet::storage]
-	#[pallet::getter(fn transaction_count)]
-	pub type TransactionCount<T> = StorageValue<_, u64, ValueQuery>;
-
 	#[pallet::hooks]
 	impl<T: Config> Hooks<BlockNumberFor<T>> for Pallet<T> {
 		fn on_initialize(now: BlockNumberFor<T>) -> Weight {
@@ -307,12 +322,9 @@ pub mod pallet {
 		fn on_finalize(now: BlockNumberFor<T>) {
 			let max_age = T::BlockhashQueueMaxAge::get();
 			let to_remove = now.saturating_sub(max_age).saturating_sub(One::one());
-			<BlockhashQueue<T>>::remove(<frame_system::Pallet<T>>::block_hash(to_remove));
-
-			let count = frame_system::Pallet::<T>::extrinsic_count() as u64;
-			TransactionCount::<T>::mutate(|total_count| {
-				*total_count = total_count.saturating_add(count)
-			});
+			let blockhash = <frame_system::Pallet<T>>::block_hash(to_remove);
+			<BlockhashQueue<T>>::remove(blockhash);
+			<TransactionCache<T>>::remove(blockhash);
 		}
 	}
 
@@ -385,26 +397,32 @@ pub mod pallet {
 
 			let bank = <Bank<T>>::new(<Slot<T>>::get());
 
-			bank.load_execute_and_commit_sanitized_transaction(sanitized_tx);
+			if bank.load_execute_and_commit_sanitized_transaction(&sanitized_tx).is_ok() {
+				Self::update_transaction_cache(&sanitized_tx)?;
+			}
 
 			Ok(().into())
 		}
 
 		// TODO: unimplemented.
 		fn validate_transaction_in_pool(
 			_fee_payer: Pubkey,
-			_transaction: &Transaction,
+			transaction: &Transaction,
 		) -> TransactionValidity {
 			let mut builder = ValidTransactionBuilder::default();
 
+			Self::check_transaction(transaction)?;
+
 			builder.build()
 		}
 
 		// TODO: unimplemented.
 		fn validate_transaction_in_block(
 			_fee_payer: Pubkey,
-			_transaction: &Transaction,
+			transaction: &Transaction,
 		) -> Result<(), TransactionValidityError> {
+			Self::check_transaction(transaction)?;
+
 			Ok(())
 		}
 
@@ -437,10 +455,6 @@ pub mod pallet {
 			}
 		}
 
-		pub fn get_transaction_count() -> u64 {
-			TransactionCount::<T>::get()
-		}
-
 		pub fn simulate_transaction(
 			sanitized_tx: SanitizedTransaction,
 			enable_cpi_recording: bool,
@@ -524,6 +538,43 @@ pub mod pallet {
 				inner_instructions,
 			}
 		}
+
+		fn update_transaction_cache(sanitized_tx: &SanitizedTransaction) -> Result<(), Error<T>> {
+			let blockhash = T::HashConversion::convert(*sanitized_tx.message().recent_blockhash());
+			let message_hash = T::HashConversion::convert(*sanitized_tx.message_hash());
+
+			<TransactionCache<T>>::try_mutate(blockhash, |cache| cache.try_insert(message_hash))
+				.map_err(|_| Error::<T>::CacheLimitReached)?;
+
+			Ok(())
+		}
+
+		pub(crate) fn check_transaction(
+			transaction: &Transaction,
+		) -> Result<(), InvalidTransaction> {
+			// TODO: Update error code.
+			let sanitized_tx = SanitizedTransaction::try_create(
+				transaction.clone(),
+				MessageHash::Compute,
+				None,
+				SimpleAddressLoader::Disabled,
+				&ReservedAccountKeys::empty_key_set(),
+			)
+			.map_err(|_| InvalidTransaction::Custom(0))?;
+
+			if Self::is_transaction_already_processed(&sanitized_tx) {
+				return Err(InvalidTransaction::Custom(0));
+			}
+
+			Ok(())
+		}
+
+		fn is_transaction_already_processed(sanitized_tx: &SanitizedTransaction) -> bool {
+			let blockhash = T::HashConversion::convert(*sanitized_tx.message().recent_blockhash());
+			let message_hash = T::HashConversion::convert(*sanitized_tx.message_hash());
+
+			<TransactionCache<T>>::get(blockhash).contains(&message_hash)
+		}
 	}
 
 	// TODO: Generalize and move to higher level.

frame/solana/src/runtime/bank.rs

@@ -206,9 +206,9 @@ impl<T: Config> Bank<T> {
 
 	pub fn load_execute_and_commit_sanitized_transaction(
 		&self,
-		sanitized_tx: SanitizedTransaction,
+		sanitized_tx: &SanitizedTransaction,
 	) -> Result<()> {
-		let check_result = self.check_transaction(&sanitized_tx, T::BlockhashQueueMaxAge::get());
+		let check_result = self.check_transaction(sanitized_tx, T::BlockhashQueueMaxAge::get());
 
 		let blockhash = T::HashConversion::convert_back(<frame_system::Pallet<T>>::parent_hash());
 		// FIXME: Update lamports_per_signature.
@@ -240,14 +240,14 @@ impl<T: Config> Bank<T> {
 		let mut sanitized_output =
 			self.transaction_processor.load_and_execute_sanitized_transaction(
 				self,
-				&sanitized_tx,
+				sanitized_tx,
 				check_result,
 				&processing_environment,
 				&processing_config,
 			);
 
 		self.commit_transaction(
-			&sanitized_tx,
+			sanitized_tx,
 			&mut sanitized_output.loaded_transaction,
 			sanitized_output.execution_result.clone(),
 			blockhash,

frame/solana/src/tests.rs

@@ -22,7 +22,7 @@ use frame_support::{
 	sp_runtime::traits::Convert,
 	traits::{
 		fungible::{Inspect, Mutate},
-		Get,
+		Get, OnFinalize,
 	},
 	BoundedVec,
 };
@@ -82,7 +82,7 @@ fn process_transaction(bank: &Bank<Test>, tx: Transaction) -> Result<()> {
 	)
 	.expect("Transaction must be sanitized");
 
-	bank.load_execute_and_commit_sanitized_transaction(sanitized_tx)
+	bank.load_execute_and_commit_sanitized_transaction(&sanitized_tx)
 }
 
 fn mock_deploy_program(program_id: &Pubkey, data: Vec<u8>) {
@@ -272,3 +272,30 @@ fn spl_token_program_should_work() {
 		assert_eq!(state.amount, sol_into_lamports(1_000));
 	});
 }
+
+#[test]
+fn filter_duplicated_transaction() {
+	new_test_ext().execute_with(|| {
+		before_each();
+		let bank = mock_bank();
+
+		let from = Keypair::alice();
+		let to = Keypair::bob();
+		let lamports = 100_000_000;
+
+		let transfer = system_instruction::transfer(&from.pubkey(), &to.pubkey(), lamports);
+		let mut tx = Transaction::new_with_payer(&[transfer], Some(&from.pubkey()));
+
+		let origin = RawOrigin::SolanaTransaction(from.pubkey());
+		let versioned_tx: VersionedTransaction = tx.into();
+		assert!(Pallet::<Test>::check_transaction(&versioned_tx).is_ok());
+
+		assert!(Pallet::<Test>::transact(origin.into(), versioned_tx.clone()).is_ok());
+
+		// A duplicated transaction was submitted, causing an error.
+		assert!(Pallet::<Test>::check_transaction(&versioned_tx).is_err());
+
+		Solana::on_finalize(22);
+		assert!(Pallet::<Test>::check_transaction(&versioned_tx).is_ok());
+	});
+}