High Sequence Number Gas DOS

networks/suzuka/suzuka-client/src/bin/e2e/gas_dos.rs

@@ -69,6 +69,7 @@ pub async fn create_fake_signed_transaction(
 	from_account: &LocalAccount,
 	to_account: AccountAddress,
 	amount: u64,
+	sequence_number: u64,
 ) -> Result<SignedTransaction, anyhow::Error> {
 	let coin_type = "0x1::aptos_coin::AptosCoin";
 	let timeout_secs = 600; // 10 minutes
@@ -89,21 +90,21 @@ pub async fn create_fake_signed_transaction(
 		ChainId::new(chain_id),
 	);
 
-	let raw_txn = transaction_builder
+	let raw_transaction = transaction_builder
 		.sender(from_account.address())
-		.sequence_number(from_account.sequence_number())
+		.sequence_number(sequence_number)
 		.max_gas_amount(max_gas_amount)
 		.gas_unit_price(gas_unit_price)
 		.expiration_timestamp_secs(expiration_time)
 		.chain_id(ChainId::new(chain_id))
 		.build();
 
-	let signed_txn = from_account.sign_transaction(raw_txn);
+	let signed_transaction = from_account.sign_transaction(raw_transaction);
 
-	Ok(signed_txn)
+	Ok(signed_transaction)
 }
 
-pub async fn test_sending_failed_tx() -> Result<(), anyhow::Error> {
+pub async fn test_sending_failed_transaction() -> Result<(), anyhow::Error> {
 	let rest_client = Client::new(NODE_URL.clone());
 	let faucet_client = FaucetClient::new(FAUCET_URL.clone(), NODE_URL.clone());
 	let coin_client = CoinClient::new(&rest_client);
@@ -139,8 +140,14 @@ pub async fn test_sending_failed_tx() -> Result<(), anyhow::Error> {
 	println!("Alice: {:?}", initial_balance);
 
 	// TEST 1: Sending a transaction trying to transfer more coins than Alice has (including gas fees)
-	let transaction =
-		create_fake_signed_transaction(chain_id, &alice, bob.address(), 100_000_000).await?;
+	let transaction = create_fake_signed_transaction(
+		chain_id,
+		&alice,
+		bob.address(),
+		100_000_000,
+		alice.sequence_number(),
+	)
+	.await?;
 
 	let _transaction_will_fail = rest_client
 		.submit(&transaction)
@@ -163,21 +170,92 @@ pub async fn test_sending_failed_tx() -> Result<(), anyhow::Error> {
 	println!("Alice: {:?}", failed_balance);
 	assert!(initial_balance > failed_balance);
 
-	// TEST 2: Sending n transactions with a high sequence number
-	let n = 10;
+	// TEST 2: Sending a transaction with a high sequence number
+	let too_high_sequence_number = alice.sequence_number() + 32 + 2;
+	println!("Alice's sequence number: {}", alice.sequence_number());
+	println!("Too high sequence number: {}", too_high_sequence_number);
 	let mut last_balance = failed_balance;
-	let transaction =
-		create_fake_signed_transaction(chain_id, &alice, bob.address(), 10_000_000).await?;
+	let transaction = create_fake_signed_transaction(
+		chain_id,
+		&alice,
+		bob.address(),
+		100,
+		too_high_sequence_number, // too new tolerance is 32
+	)
+	.await?;
 
 	match rest_client.submit(&transaction).await {
 		Ok(_) => panic!("Transaction should have failed with high sequence number"),
 		Err(e) => match e {
 			suzuka_client::rest_client::error::RestError::Api(aptos_error) => {
 				println!("Transaction failed as expected: {:?}", aptos_error);
-				assert_eq!(
-					aptos_error.error.error_code as u32,
-					suzuka_client::rest_client::aptos_api_types::AptosErrorCode::VmError as u32
-				);
+				assert_eq!(aptos_error.error.error_code as u32, 402); // 402 is used for too old and too new
+			}
+			_ => panic!("Unexpected error: {:?}", e),
+		},
+	}
+
+	// assert that no gas fee charged because the transaction never entered the mempool
+	let failed_balance = coin_client
+		.get_account_balance(&alice.address())
+		.await
+		.context("Failed to get Alice's account balance")?;
+	println!("\n=== After Failed Tx#2 ===");
+	println!("Alice: {:?}", failed_balance);
+	assert!(last_balance == failed_balance);
+
+	// TEST 3: Sending a transaction with a sequence number that won't be accepted by the VM, but would be accepted by the mempool (sequence number cannot be reused)
+	let attack_sequence_number = alice.sequence_number() + 5;
+	let transaction = create_fake_signed_transaction(
+		chain_id,
+		&alice,
+		bob.address(),
+		100,
+		attack_sequence_number,
+	)
+	.await?;
+
+	// transaction should fail in the vm not on the submission
+	let _transaction_will_fail = rest_client
+		.submit(&transaction)
+		.await
+		.context("Failed when waiting for the transaction")?
+		.into_inner();
+	match rest_client.wait_for_signed_transaction(&transaction).await {
+		Ok(_) => panic!("Transaction should have failed"),
+		Err(e) => {
+			println!("Transaction failed as expected: {:?}", e);
+		}
+	}
+
+	// assert gas fee not charged
+	let failed_balance = coin_client
+		.get_account_balance(&alice.address())
+		.await
+		.context("Failed to get Alice's account balance")?;
+	println!("\n=== After Failed Tx#3 ===");
+	println!("Alice: {:?}", failed_balance);
+	assert!(last_balance == failed_balance);
+
+	// transaction using the same sequence number should fail to submit
+	let transaction = create_fake_signed_transaction(
+		chain_id,
+		&alice,
+		bob.address(),
+		100,
+		attack_sequence_number,
+	)
+	.await?;
+
+	match rest_client.submit(&transaction).await {
+		Ok(res) => panic!(
+			"Transaction should have failed with high sequence number. Instead got: {:?}",
+			res
+		),
+		Err(e) => match e {
+			suzuka_client::rest_client::error::RestError::Api(aptos_error) => {
+				println!("Transaction failed as expected: {:?}", aptos_error);
+				assert_eq!(aptos_error.error.error_code as u32, 402); // 402 is used for too old and too new
 			}
 			_ => panic!("Unexpected error: {:?}", e),
 		},
@@ -188,6 +266,6 @@ pub async fn test_sending_failed_tx() -> Result<(), anyhow::Error> {
 
 #[tokio::main]
 async fn main() -> Result<(), anyhow::Error> {
-	test_sending_failed_tx().await?;
+	test_sending_failed_transaction().await?;
 	Ok(())
 }

protocol-units/execution/opt-executor/src/executor/initialization.rs

@@ -132,6 +132,8 @@ impl Executor {
 			&node_config,
 			Arc::clone(&self.transactions_in_flight),
 			maptos_config.load_shedding.max_transactions_in_flight,
+			self.config.mempool.sequence_number_ttl_ms,
+			self.config.mempool.gc_slot_duration_ms,
 		);
 
 		let cx = Context::new(

protocol-units/execution/opt-executor/src/gc_account_sequence_number.rs

@@ -0,0 +1,157 @@
+use aptos_types::account_address::AccountAddress;
+use std::collections::{BTreeMap, HashMap};
+use tracing::info;
+
+pub struct UsedSequenceNumberPool {
+	/// The number of milliseconds a sequence number is valid for.
+	sequence_number_ttl_ms: u64,
+	/// The duration of a garbage collection slot in milliseconds.
+	/// This is used to bin sequence numbers into slots for O(sequence_number_ttl_ms/gc_slot_duration_ms * log sequence_number_ttl_ms/gc_slot_duration_ms) garbage collection.
+	gc_slot_duration_ms: u64,
+	/// The sequence number lifetimes, indexed by slot.
+	sequence_number_lifetimes: BTreeMap<u64, HashMap<AccountAddress, u64>>,
+}
+
+impl UsedSequenceNumberPool {
+	/// Creates a new UsedSequenceNumberPool with a specified garbage collection slot duration.
+	pub(crate) fn new(sequence_number_ttl_ms: u64, gc_slot_duration_ms: u64) -> Self {
+		UsedSequenceNumberPool {
+			sequence_number_ttl_ms,
+			gc_slot_duration_ms,
+			sequence_number_lifetimes: BTreeMap::new(),
+		}
+	}
+
+	/// Gets a sequence number for an account
+	pub(crate) fn get_sequence_number(&self, account: &AccountAddress) -> Option<u64> {
+		// check each slot for the account
+		for (_slot, lifetimes) in self.sequence_number_lifetimes.iter().rev() {
+			// reverse order is better average case because highly-used sequence numbers will be moved up more often
+			match lifetimes.get(account) {
+				Some(sequence_number) => {
+					// check if the sequence number is still valid
+					return Some(*sequence_number);
+				}
+				None => {}
+			}
+		}
+
+		None
+	}
+
+	/// Removes the sequence number for an account.
+	pub(crate) fn remove_sequence_number(&mut self, account_address: &AccountAddress) {
+		// check each slot for the account
+		for (_slot, lifetimes) in self.sequence_number_lifetimes.iter_mut().rev() {
+			if lifetimes.remove(account_address).is_some() {
+				break;
+			}
+		}
+	}
+
+	/// Sets the sequence number for an account.
+	pub(crate) fn set_sequence_number(
+		&mut self,
+		account_address: &AccountAddress,
+		sequence_number: u64,
+		current_time_ms: u64,
+	) {
+		// remove the old sequence number
+		self.remove_sequence_number(account_address);
+
+		// compute the slot for the new lifetime and add accordingly
+		let slot = current_time_ms / self.gc_slot_duration_ms;
+
+		// add the new sequence number
+		self.sequence_number_lifetimes
+			.entry(slot)
+			.or_insert_with(HashMap::new)
+			.insert(*account_address, sequence_number);
+	}
+
+	/// Garbage collects sequence numbers that have expired.
+	/// This should be called periodically.
+	pub(crate) fn gc(&mut self, current_time_ms: u64) {
+		let gc_slot = current_time_ms / self.gc_slot_duration_ms;
+
+		// remove all slots that are too old
+		let slots_to_remove: Vec<u64> = self
+			.sequence_number_lifetimes
+			.range(..gc_slot - self.sequence_number_ttl_ms / self.gc_slot_duration_ms)
+			.map(|(slot, _)| *slot)
+			.collect();
+		for slot in slots_to_remove {
+			println!(
+				"Garbage collecting sequence number slot {} with duration {} timestamp {}",
+				slot,
+				self.gc_slot_duration_ms,
+				slot * self.gc_slot_duration_ms
+			);
+			self.sequence_number_lifetimes.remove(&slot);
+		}
+	}
+}
+
+#[cfg(test)]
+pub mod test {
+
+	use super::*;
+
+	#[test]
+	fn test_inserts() {
+		let mut pool = UsedSequenceNumberPool::new(1000, 100);
+		let account1 = AccountAddress::random();
+		let account2 = AccountAddress::random();
+
+		pool.set_sequence_number(&account1, 1, 0);
+		pool.set_sequence_number(&account2, 2, 0);
+		assert_eq!(pool.get_sequence_number(&account1), Some(1));
+		assert_eq!(pool.get_sequence_number(&account2), Some(2));
+	}
+
+	#[test]
+	fn test_removes() {
+		let mut pool = UsedSequenceNumberPool::new(1000, 100);
+		let account1 = AccountAddress::random();
+		let account2 = AccountAddress::random();
+
+		pool.set_sequence_number(&account1, 1, 0);
+		pool.set_sequence_number(&account2, 2, 0);
+		pool.remove_sequence_number(&account1);
+		assert_eq!(pool.get_sequence_number(&account1), None);
+		assert_eq!(pool.get_sequence_number(&account2), Some(2));
+	}
+
+	#[test]
+	fn test_gc() {
+		let mut pool = UsedSequenceNumberPool::new(1000, 100);
+		let account1 = AccountAddress::random();
+		let account2 = AccountAddress::random();
+
+		pool.set_sequence_number(&account1, 1, 0);
+		pool.set_sequence_number(&account2, 2, 0);
+		pool.gc(1000);
+		assert_eq!(pool.get_sequence_number(&account1), Some(1));
+		assert_eq!(pool.get_sequence_number(&account2), Some(2));
+		pool.gc(2000);
+		assert_eq!(pool.get_sequence_number(&account1), None);
+		assert_eq!(pool.get_sequence_number(&account2), None);
+	}
+
+	#[test]
+	fn test_gc_removes_some_not_all() {
+		let mut pool = UsedSequenceNumberPool::new(1000, 100);
+		let account1 = AccountAddress::random();
+		let account2 = AccountAddress::random();
+
+		pool.set_sequence_number(&account1, 1, 0);
+		pool.set_sequence_number(&account2, 2, 0);
+		pool.gc(1000);
+		assert_eq!(pool.get_sequence_number(&account1), Some(1));
+		assert_eq!(pool.get_sequence_number(&account2), Some(2));
+		pool.set_sequence_number(&account1, 3, 1000);
+		pool.gc(2000);
+		assert_eq!(pool.get_sequence_number(&account1), Some(3));
+		assert_eq!(pool.get_sequence_number(&account2), None);
+	}
+}

protocol-units/execution/opt-executor/src/lib.rs

@@ -2,6 +2,7 @@ pub mod bootstrap;
 pub mod context;
 #[warn(unused_imports)]
 pub mod executor;
+pub mod gc_account_sequence_number;
 pub mod indexer;
 pub mod service;
 pub mod transaction_pipe;

protocol-units/execution/opt-executor/src/service.rs

@@ -120,7 +120,8 @@ mod tests {
 		let service = Service::new(&context);
 		let handle = tokio::spawn(async move { service.run().await });
 
-		let user_transaction = create_signed_transaction(0, &context.config().chain);
+		// this needs to be 1 because the root account should already have a committed genesis transaction
+		let user_transaction = create_signed_transaction(1, &context.config().chain);
 
 		// send transaction to mempool
 		let (req_sender, callback) = oneshot::channel();