From b2c296ab92f76235a37f6a03f9bb1aa3ffb6d656 Mon Sep 17 00:00:00 2001 From: Tarrence van As Date: Fri, 4 Oct 2024 11:31:15 -0400 Subject: [PATCH] Manage nonce internally --- .github/workflows/test.yml | 2 +- packages/account_sdk/src/controller.rs | 119 ++++++++++++++------ packages/account_sdk/src/controller_test.rs | 87 +++++++++++++- packages/account_sdk/src/session.rs | 8 +- 4 files changed, 172 insertions(+), 44 deletions(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 54a8b38c7..087e999db 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -14,7 +14,7 @@ jobs: rust: runs-on: ubuntu-latest container: - image: ghcr.io/dojoengine/dojo-dev:v1.0.0-alpha.4 + image: ghcr.io/dojoengine/dojo-dev:v1.0.0-alpha.11 steps: - uses: actions/checkout@v4 - run: git config --system --add safe.directory '*' diff --git a/packages/account_sdk/src/controller.rs b/packages/account_sdk/src/controller.rs index 741ebb55f..1f0da5aa0 100644 --- a/packages/account_sdk/src/controller.rs +++ b/packages/account_sdk/src/controller.rs @@ -47,6 +47,7 @@ pub struct Controller { contract: Option>>, pub factory: ControllerFactory, pub(crate) storage: Storage, + nonce: Felt, } impl Controller { @@ -85,6 +86,7 @@ impl Controller { contract: None, factory, storage: Storage::default(), + nonce: Felt::ZERO, }; let contract = Box::new(abigen::controller::Controller::new( @@ -184,43 +186,68 @@ impl Controller { return self.execute_from_outside(calls).await; } - let result = self.execute_v1(calls.clone()).max_fee(max_fee).send().await; - - match result { - Ok(tx_result) => { - // Update is_registered to true after successful execution with a session - if let Some((key, metadata)) = - self.session_metadata(&Policy::from_calls(&calls), None) - { - if !metadata.is_registered { - let mut updated_metadata = metadata; - updated_metadata.is_registered = true; - self.storage.set_session(&key, updated_metadata)?; + let mut retry_count = 0; + let max_retries = 1; + + loop { + let nonce = self.get_nonce().await?; + let result = self + .execute_v1(calls.clone()) + .nonce(nonce) + .max_fee(max_fee) + .send() + .await; + + match result { + Ok(tx_result) => { + // Update nonce + self.nonce = self.nonce + 1; + + // Update is_registered to true after successful execution with a session + if let Some((key, metadata)) = + self.session_metadata(&Policy::from_calls(&calls), None) + { + if !metadata.is_registered { + let mut updated_metadata = metadata; + updated_metadata.is_registered = true; + self.storage.set_session(&key, updated_metadata)?; + } } + return Ok(tx_result); } - Ok(tx_result) - } - Err(e) => { - if let AccountError::Provider(ProviderError::StarknetError( - StarknetError::TransactionExecutionError(data), - )) = &e - { - if data - .execution_error - .contains(&format!("{:x} is not deployed.", self.address)) - { - let balance = self.eth_balance().await?; - let mut fee_estimate = self.deploy().estimate_fee().await?; - fee_estimate.overall_fee += WEBAUTHN_GAS * fee_estimate.gas_price; - Err(ControllerError::NotDeployed { - fee_estimate: Box::new(fee_estimate), - balance, - }) - } else { - Err(ControllerError::AccountError(e)) + Err(e) => { + match &e { + AccountError::Provider(ProviderError::StarknetError( + StarknetError::TransactionExecutionError(data), + )) if data + .execution_error + .contains(&format!("{:x} is not deployed.", self.address)) => + { + let balance = self.eth_balance().await?; + let mut fee_estimate = self.deploy().estimate_fee().await?; + fee_estimate.overall_fee += WEBAUTHN_GAS * fee_estimate.gas_price; + return Err(ControllerError::NotDeployed { + fee_estimate: Box::new(fee_estimate), + balance, + }); + } + AccountError::Provider(ProviderError::StarknetError( + StarknetError::InvalidTransactionNonce, + )) => { + if retry_count < max_retries { + // Refetch nonce from the provider + let new_nonce = self + .provider + .get_nonce(self.block_id(), self.address()) + .await?; + self.nonce = new_nonce; + retry_count += 1; + continue; + } + } + _ => {} } - } else { - Err(ControllerError::AccountError(e)) + return Err(ControllerError::AccountError(e)); } } } @@ -265,6 +292,26 @@ impl Controller { let signature = self.owner.sign(&hash).await?; Ok(Vec::::cairo_serialize(&vec![signature])) } + + async fn get_nonce(&self) -> Result { + let current_nonce = self.nonce; + + if current_nonce == Felt::ZERO { + match self + .provider + .get_nonce(self.block_id(), self.address()) + .await + { + Ok(nonce) => Ok(nonce), + Err(ProviderError::StarknetError(StarknetError::ContractNotFound)) => { + Ok(Felt::ZERO) + } + Err(e) => Err(e), + } + } else { + Ok(current_nonce) + } + } } impl_account!(Controller, |account: &Controller, context| { @@ -289,9 +336,7 @@ impl ConnectedAccount for Controller { } async fn get_nonce(&self) -> Result { - self.provider - .get_nonce(self.block_id(), self.address()) - .await + self.get_nonce().await } } diff --git a/packages/account_sdk/src/controller_test.rs b/packages/account_sdk/src/controller_test.rs index a2c97ae75..9784f3f42 100644 --- a/packages/account_sdk/src/controller_test.rs +++ b/packages/account_sdk/src/controller_test.rs @@ -1,11 +1,16 @@ use std::time::Duration; use crate::{ + abigen::erc_20::Erc20, artifacts::{Version, CONTROLLERS}, controller::Controller, signers::Signer, - tests::{runners::katana::KatanaRunner, transaction_waiter::TransactionWaiter}, + tests::{ + account::FEE_TOKEN_ADDRESS, runners::katana::KatanaRunner, + transaction_waiter::TransactionWaiter, + }, }; +use cainome::cairo_serde::{ContractAddress, U256}; use starknet::{accounts::Account, macros::felt, providers::Provider, signers::SigningKey}; use starknet_crypto::Felt; @@ -71,3 +76,83 @@ async fn test_deploy_controller() { "Deployed address doesn't match expected address" ); } + +#[tokio::test] +async fn test_controller_nonce_mismatch_recovery() { + let username = "testuser".to_string(); + let signer = Signer::new_starknet_random(); + + let runner = KatanaRunner::load(); + let mut controller1 = runner + .deploy_controller(username.clone(), signer.clone(), Version::LATEST) + .await; + + let chain_id = runner.client().chain_id().await.unwrap(); + + // Create the second controller with the same credentials and address + let mut controller2 = Controller::new( + "app_id".to_string(), + username.clone(), + controller1.class_hash, + runner.rpc_url.clone(), + signer.clone(), + controller1.address, + chain_id, + ); + + // Send a transaction using controller1 + let recipient = ContractAddress(felt!("0x18301129")); + let amount = U256 { low: 0, high: 0 }; + let erc20 = Erc20::new(*FEE_TOKEN_ADDRESS, &controller1); + + let tx1 = erc20.transfer_getcall(&recipient, &amount); + let max_fee = controller1 + .estimate_invoke_fee(vec![tx1.clone()]) + .await + .unwrap(); + let res = Controller::execute(&mut controller1, vec![tx1.clone()], max_fee.overall_fee) + .await + .unwrap(); + + TransactionWaiter::new(res.transaction_hash, runner.client()) + .wait() + .await + .unwrap(); + + // Now send a transaction using controller2, which should have a stale nonce + let erc20 = Erc20::new(*FEE_TOKEN_ADDRESS, &controller2); + let tx2 = erc20.transfer_getcall(&recipient, &amount); + + let tx2_result = + Controller::execute(&mut controller2, vec![tx2.clone()], max_fee.overall_fee).await; + + // Verify that it succeeds after recovering from nonce mismatch + assert!( + tx2_result.is_ok(), + "Controller did not recover from nonce mismatch: {:?}", + tx2_result.err() + ); + + TransactionWaiter::new(tx2_result.unwrap().transaction_hash, runner.client()) + .wait() + .await + .unwrap(); + + let res = Controller::execute(&mut controller1, vec![tx1], max_fee.overall_fee) + .await + .unwrap(); + + TransactionWaiter::new(res.transaction_hash, runner.client()) + .wait() + .await + .unwrap(); + + let tx2_result = Controller::execute(&mut controller2, vec![tx2], max_fee.overall_fee).await; + + // Verify that it succeeds after recovering from nonce mismatch + assert!( + tx2_result.is_ok(), + "Controller did not recover from nonce mismatch: {:?}", + tx2_result.err() + ); +} diff --git a/packages/account_sdk/src/session.rs b/packages/account_sdk/src/session.rs index 04f23c7cc..c3bb0dc40 100644 --- a/packages/account_sdk/src/session.rs +++ b/packages/account_sdk/src/session.rs @@ -90,12 +90,10 @@ impl Controller { }), )?; - let txn = self + let call = self .contract() - .register_session(&session.raw(), &self.owner_guid()) - .max_fee(max_fee) - .send() - .await?; + .register_session_getcall(&session.raw(), &self.owner_guid()); + let txn = self.execute(vec![call], max_fee).await?; self.storage.set_session( &Selectors::session(&self.address, &self.app_id, &self.chain_id),