rusk-wallet: Allow withdrawing of partial rewards

rusk-wallet/src/bin/command.rs

@@ -200,6 +200,12 @@ pub(crate) enum Command {
         #[arg(short, long)]
         address: Option<Address>,
 
+        /// 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: Option<Dusk>,
+
         /// Max amount of gas for this transaction
         #[arg(short = 'l', long, default_value_t = DEFAULT_LIMIT_CALL)]
         gas_limit: u64,
@@ -452,6 +458,7 @@ impl Command {
             }
             Command::Withdraw {
                 address,
+                reward,
                 gas_limit,
                 gas_price,
             } => {
@@ -462,10 +469,14 @@ impl Command {
                 let tx = match address {
                     Address::Shielded(_) => {
                         wallet.sync().await?;
-                        wallet.phoenix_stake_withdraw(addr_idx, gas).await
+                        wallet
+                            .phoenix_stake_withdraw(addr_idx, reward, gas)
+                            .await
                     }
                     Address::Public(_) => {
-                        wallet.moonlight_stake_withdraw(addr_idx, gas).await
+                        wallet
+                            .moonlight_stake_withdraw(addr_idx, reward, gas)
+                            .await
                     }
                 }?;
 

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,
@@ -415,16 +418,27 @@ fn confirm(cmd: &Command, wallet: &Wallet<WalletFile>) -> anyhow::Result<bool> {
         Command::Withdraw {
             address,
             gas_limit,
+            reward,
             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());
             println!("   > Receive rewards at {}", sender.preview());
+            println!("   > Amount withdrawing {} DUSK", reward);
             println!("   > Max fee = {} DUSK", Dusk::from(max_fee));
             if let Address::Public(_) = sender {
                 println!("   > ALERT: THIS IS A PUBLIC TRANSACTION");

rusk-wallet/src/bin/interactive/command_menu.rs

@@ -95,6 +95,7 @@ pub(crate) async fn online(
     }
 
     let select = select?;
+    let max_withdraw = wallet.get_stake_reward(profile_idx).await?;
 
     let res = match select {
         MenuItem::Transfer => {
@@ -245,6 +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: Some(prompt::request_token_amt(
+                    "withdraw rewards",
+                    max_withdraw,
+                )?),
                 gas_price: prompt::request_gas_price(
                     DEFAULT_PRICE,
                     mempool_gas_prices,

rusk-wallet/src/error.rs

@@ -165,6 +165,9 @@ pub enum Error {
     /// Error while querying archival node
     #[error("Archive node query error: {0}")]
     ArchiveJsonError(String),
+    /// Trying to withdraw more reward than the person has
+    #[error("Trying to withdraw more than existing reward")]
+    NotEnoughReward,
 }
 
 impl From<dusk_bytes::Error> for Error {

rusk-wallet/src/wallet.rs

@@ -623,6 +623,20 @@ impl<F: SecureWalletFile + Debug> Wallet<F> {
 
         Ok(gas_prices)
     }
+
+    /// Get the amount of stake rewards the user has
+    pub async fn get_stake_reward(
+        &self,
+        sender_index: u8,
+    ) -> Result<Dusk, Error> {
+        let available_reward = self
+            .stake_info(sender_index)
+            .await?
+            .ok_or(Error::NotStaked)?
+            .reward;
+
+        Ok(Dusk::from(available_reward))
+    }
 }
 
 /// This structs represent a Note decoded enriched with useful chain information

rusk-wallet/src/wallet/transaction.rs

@@ -472,6 +472,7 @@ impl<F: SecureWalletFile + Debug> Wallet<F> {
     pub async fn phoenix_stake_withdraw(
         &self,
         sender_idx: u8,
+        reward_amt: Option<Dusk>,
         gas: Gas,
     ) -> Result<Transaction, Error> {
         let state = self.state()?;
@@ -488,11 +489,23 @@ impl<F: SecureWalletFile + Debug> Wallet<F> {
 
         let stake_pk = BlsPublicKey::from(&stake_sk);
 
-        let reward_amount = state
+        let available_reward = state
             .fetch_stake(&stake_pk)
             .await?
             .map(|s| s.reward)
-            .unwrap_or(0);
+            .ok_or(Error::NoReward)?;
+
+        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);
@@ -504,7 +517,7 @@ impl<F: SecureWalletFile + Debug> Wallet<F> {
             &stake_owner_sk,
             inputs,
             root,
-            reward_amount,
+            reward_amt_withdrawn,
             gas.limit,
             gas.price,
             chain_id,
@@ -523,6 +536,7 @@ impl<F: SecureWalletFile + Debug> Wallet<F> {
     pub async fn moonlight_stake_withdraw(
         &self,
         sender_idx: u8,
+        reward_amt: Option<Dusk>,
         gas: Gas,
     ) -> Result<Transaction, Error> {
         let mut rng = StdRng::from_entropy();
@@ -532,8 +546,20 @@ impl<F: SecureWalletFile + Debug> Wallet<F> {
         let nonce = state.fetch_account(pk).await?.nonce + 1;
         let chain_id = state.fetch_chain_id().await?;
         let stake_info = state.fetch_stake(pk).await?;
-        let reward = stake_info.map(|s| s.reward).ok_or(Error::NoReward)?;
-        let reward = Dusk::from(reward);
+        let available_reward =
+            stake_info.map(|s| s.reward).ok_or(Error::NoReward)?;
+
+        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);
 
@@ -546,7 +572,7 @@ impl<F: SecureWalletFile + Debug> Wallet<F> {
             &sender_sk,
             &sender_sk,
             &stake_owner_sk,
-            *reward,
+            reward_amt_withdrawn,
             gas.limit,
             gas.price,
             nonce,