rusk-wallet: Withdraw all rewards if none specified

Signed-off-by: Daksh <[email protected]>

rusk-wallet/src/bin/command.rs

@@ -202,9 +202,11 @@ pub(crate) enum Command {
         #[arg(short, long)]
         address: Option<Address>,
 
-        /// 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<Dusk>,
 
         /// Max amount of gas for this transaction
         #[arg(short = 'l', long, default_value_t = DEFAULT_LIMIT_CALL)]

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<WalletFile>) -> anyhow::Result<bool> {
+async fn confirm(
+    cmd: &Command,
+    wallet: &Wallet<WalletFile>,
+) -> anyhow::Result<bool> {
     match cmd {
         Command::Transfer {
             sender,
@@ -419,9 +422,18 @@ fn confirm(cmd: &Command, wallet: &Wallet<WalletFile>) -> anyhow::Result<bool> {
             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());

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,

rusk-wallet/src/wallet.rs

@@ -629,12 +629,11 @@ impl<F: SecureWalletFile + Debug> Wallet<F> {
         &self,
         sender_index: u8,
     ) -> Result<Dusk, Error> {
-        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))
     }

rusk-wallet/src/wallet/transaction.rs

@@ -472,7 +472,7 @@ impl<F: SecureWalletFile + Debug> Wallet<F> {
     pub async fn phoenix_stake_withdraw(
         &self,
         sender_idx: u8,
-        reward_amt: Dusk,
+        reward_amt: Option<Dusk>,
         gas: Gas,
     ) -> Result<Transaction, Error> {
         let state = self.state()?;
@@ -495,10 +495,17 @@ impl<F: SecureWalletFile + Debug> Wallet<F> {
             .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<F: SecureWalletFile + Debug> Wallet<F> {
             &stake_owner_sk,
             inputs,
             root,
-            *reward_amt,
+            reward_amt_withdrawn,
             gas.limit,
             gas.price,
             chain_id,
@@ -529,7 +536,7 @@ impl<F: SecureWalletFile + Debug> Wallet<F> {
     pub async fn moonlight_stake_withdraw(
         &self,
         sender_idx: u8,
-        reward_amt: Dusk,
+        reward_amt: Option<Dusk>,
         gas: Gas,
     ) -> Result<Transaction, Error> {
         let mut rng = StdRng::from_entropy();
@@ -542,10 +549,17 @@ impl<F: SecureWalletFile + Debug> Wallet<F> {
         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);