From d7770d7a56a9eedfc6451f14b69dd402f44a4117 Mon Sep 17 00:00:00 2001 From: Daksh <41485688+Daksh14@users.noreply.github.com> Date: Tue, 26 Nov 2024 02:47:09 -0500 Subject: [PATCH] rusk-wallet: Withdraw all rewards if none specified Signed-off-by: Daksh <41485688+Daksh14@users.noreply.github.com> --- rusk-wallet/src/bin/command.rs | 6 ++-- rusk-wallet/src/bin/interactive.rs | 20 ++++++++--- .../src/bin/interactive/command_menu.rs | 4 +-- rusk-wallet/src/wallet.rs | 11 +++--- rusk-wallet/src/wallet/transaction.rs | 36 +++++++++++++------ 5 files changed, 52 insertions(+), 25 deletions(-) diff --git a/rusk-wallet/src/bin/command.rs b/rusk-wallet/src/bin/command.rs index a86818bd7..eabfabb6d 100644 --- a/rusk-wallet/src/bin/command.rs +++ b/rusk-wallet/src/bin/command.rs @@ -202,9 +202,11 @@ pub(crate) enum Command { #[arg(short, long)] address: Option
, - /// Amount of rewards to withdraw from the stake contract + /// Amount of rewards to withdraw from the stake contract. If the + /// reward is not provided, all the rewards are withdrawn at + /// once #[arg(short, long)] - reward: Dusk, + reward: Option, /// Max amount of gas for this transaction #[arg(short = 'l', long, default_value_t = DEFAULT_LIMIT_CALL)] diff --git a/rusk-wallet/src/bin/interactive.rs b/rusk-wallet/src/bin/interactive.rs index 7f45e325b..2421a565f 100644 --- a/rusk-wallet/src/bin/interactive.rs +++ b/rusk-wallet/src/bin/interactive.rs @@ -84,7 +84,7 @@ pub(crate) async fn run_loop( match op { Ok(ProfileOp::Run(cmd)) => { // request confirmation before running - if confirm(&cmd, wallet)? { + if confirm(&cmd, wallet).await? { // run command prompt::hide_cursor()?; let res = cmd.run(wallet, settings).await?; @@ -346,7 +346,10 @@ async fn menu_wallet( } /// Request user confirmation for a transfer transaction -fn confirm(cmd: &Command, wallet: &Wallet) -> anyhow::Result { +async fn confirm( + cmd: &Command, + wallet: &Wallet, +) -> anyhow::Result { match cmd { Command::Transfer { sender, @@ -419,9 +422,18 @@ fn confirm(cmd: &Command, wallet: &Wallet) -> anyhow::Result { gas_price, } => { let sender = address.as_ref().ok_or(Error::BadAddress)?; + let sender_index = wallet.find_index(sender)?; let max_fee = gas_limit * gas_price; - let withdraw_from = - wallet.public_address(wallet.find_index(sender)?)?; + let withdraw_from = wallet.public_address(sender_index)?; + + let total_rewards = wallet.get_stake_reward(sender_index).await?; + + // withdraw all rewards if no amt specified + let reward = if let Some(withdraw_reward) = reward { + withdraw_reward + } else { + &total_rewards + }; println!(" > Pay with {}", sender.preview()); println!(" > Withdraw rewards from {}", withdraw_from.preview()); diff --git a/rusk-wallet/src/bin/interactive/command_menu.rs b/rusk-wallet/src/bin/interactive/command_menu.rs index c5702a07f..6407e3ebe 100644 --- a/rusk-wallet/src/bin/interactive/command_menu.rs +++ b/rusk-wallet/src/bin/interactive/command_menu.rs @@ -246,10 +246,10 @@ pub(crate) async fn online( ProfileOp::Run(Box::new(Command::Withdraw { address: Some(addr), gas_limit: prompt::request_gas_limit(gas::DEFAULT_LIMIT_CALL)?, - reward: prompt::request_token_amt( + reward: Some(prompt::request_token_amt( "withdraw rewards", max_withdraw, - )?, + )?), gas_price: prompt::request_gas_price( DEFAULT_PRICE, mempool_gas_prices, diff --git a/rusk-wallet/src/wallet.rs b/rusk-wallet/src/wallet.rs index b1f8b74e4..c841a985e 100644 --- a/rusk-wallet/src/wallet.rs +++ b/rusk-wallet/src/wallet.rs @@ -629,12 +629,11 @@ impl Wallet { &self, sender_index: u8, ) -> Result { - let state = self.state()?; - let pk = self.public_key(sender_index)?; - - let stake_info = state.fetch_stake(pk).await?; - let available_reward = - stake_info.map(|s| s.reward).ok_or(Error::NoReward)?; + let available_reward = self + .stake_info(sender_index) + .await? + .ok_or(Error::NotStaked)? + .reward; Ok(Dusk::from(available_reward)) } diff --git a/rusk-wallet/src/wallet/transaction.rs b/rusk-wallet/src/wallet/transaction.rs index 60b29817a..825d1098e 100644 --- a/rusk-wallet/src/wallet/transaction.rs +++ b/rusk-wallet/src/wallet/transaction.rs @@ -472,7 +472,7 @@ impl Wallet { pub async fn phoenix_stake_withdraw( &self, sender_idx: u8, - reward_amt: Dusk, + reward_amt: Option, gas: Gas, ) -> Result { let state = self.state()?; @@ -495,10 +495,17 @@ impl Wallet { .map(|s| s.reward) .ok_or(Error::NoReward)?; - // throw error if we try to withdraw more than available - if reward_amt > available_reward { - return Err(Error::NotEnoughReward); - } + let reward_amt_withdrawn = if let Some(reward_amt) = reward_amt { + // throw error if we try to withdraw more than available + if reward_amt > available_reward { + return Err(Error::NotEnoughReward); + } + + *reward_amt + } else { + // withdraw all the reward if no amt specified to withdraw + available_reward + }; let stake_owner_idx = self.find_stake_owner_idx(&stake_pk).await?; let mut stake_owner_sk = self.derive_bls_sk(stake_owner_idx); @@ -510,7 +517,7 @@ impl Wallet { &stake_owner_sk, inputs, root, - *reward_amt, + reward_amt_withdrawn, gas.limit, gas.price, chain_id, @@ -529,7 +536,7 @@ impl Wallet { pub async fn moonlight_stake_withdraw( &self, sender_idx: u8, - reward_amt: Dusk, + reward_amt: Option, gas: Gas, ) -> Result { let mut rng = StdRng::from_entropy(); @@ -542,10 +549,17 @@ impl Wallet { let available_reward = stake_info.map(|s| s.reward).ok_or(Error::NoReward)?; - // throw error if we try to withdraw more than available - if reward_amt > available_reward { - return Err(Error::NotEnoughReward); - } + let reward_amt_withdrawn = if let Some(reward_amt) = reward_amt { + // throw error if we try to withdraw more than available + if reward_amt > available_reward { + return Err(Error::NotEnoughReward); + } + + *reward_amt + } else { + // withdraw all the reward if no amt specified to withdraw + available_reward + }; let mut sender_sk = self.derive_bls_sk(sender_idx);