fix(drive): charge fees for unshield and shielded withdrawal

packages/rs-drive-abci/src/execution/validation/state_transition/processor/traits/shielded_proof.rs

@@ -68,11 +68,15 @@ impl StateTransitionHasShieldedProofValidationV0 for StateTransition {
 /// The minimum fee is computed dynamically based on the number of actions:
 ///   min_fee = proof_verification_fee + num_actions × (processing_fee + storage_fee)
 ///
-/// The fee is derived from the public `value_balance` field (no ZK proof execution needed):
-/// - ShieldedTransfer: fee = value_balance
-/// - Unshield: fee = value_balance - amount
-/// - ShieldedWithdrawal: fee = value_balance - amount
-/// - Shield: fee paid by transparent address inputs (skipped here)
+/// The amount checked against `min_fee` is derived from public fields (no ZK proof
+/// execution needed):
+/// - ShieldedTransfer: the whole `value_balance` is the fee, so we check `value_balance >= min_fee`.
+/// - Unshield / ShieldedWithdrawal: `unshielding_amount` is the TOTAL value leaving the
+///   shielded pool (recipient/net amount + fee). The fee actually charged at execution time
+///   is `compute_minimum_shielded_fee` (carved out of `unshielding_amount`), and the
+///   recipient/net receives `unshielding_amount - fee`. Requiring `unshielding_amount >= min_fee`
+///   here guarantees that net amount is non-negative.
+/// - Shield: fee paid by transparent address inputs (skipped here).
 pub(crate) trait StateTransitionShieldedMinimumFeeValidationV0 {
     fn validate_minimum_shielded_fee(
         &self,

packages/rs-drive-abci/src/execution/validation/state_transition/state_transitions/shielded_withdrawal/transform_into_action/v0/mod.rs

@@ -75,8 +75,8 @@ impl ShieldedWithdrawalStateTransitionTransformIntoActionValidationV0
         }
 
         // Verify the pool has sufficient balance for the withdrawal.
-        let unshielding_amount = match self {
-            ShieldedWithdrawalTransition::V0(v0) => v0.unshielding_amount,
+        let (unshielding_amount, num_actions) = match self {
+            ShieldedWithdrawalTransition::V0(v0) => (v0.unshielding_amount, v0.actions.len()),
         };
 
         if current_total_balance < unshielding_amount {
@@ -126,13 +126,21 @@ impl ShieldedWithdrawalStateTransitionTransformIntoActionValidationV0
         )?;
         execution_context.add_operation(ValidationOperation::PrecalculatedOperation(fee));
 
+        // The fee charged to the shielded pool is the minimum shielded fee computed
+        // from the same `num_actions` that `validate_minimum_shielded_fee` enforced
+        // `unshielding_amount >=` against. Because that check passed, the net amount
+        // withdrawn to Core (`unshielding_amount - fee_amount`) is guaranteed to be
+        // non-negative.
+        let fee_amount = dpp::shielded::compute_minimum_shielded_fee(num_actions, platform_version);
+
         // Build the action, which includes creating the withdrawal document
         let creation_time_ms = block_info.time_ms;
 
         let result = ShieldedWithdrawalTransitionAction::try_from_transition(
             self,
             current_total_balance,
             creation_time_ms,
+            fee_amount,
         );
 
         Ok(result.map(|action| action.into()))

packages/rs-drive-abci/src/execution/validation/state_transition/state_transitions/unshield/transform_into_action/v0/mod.rs

@@ -103,8 +103,8 @@ impl UnshieldStateTransitionTransformIntoActionValidationV0 for UnshieldTransiti
         execution_context.add_operation(ValidationOperation::PrecalculatedOperation(fee));
 
         // Verify the pool has sufficient balance for the unshield amount
-        let amount = match self {
-            UnshieldTransition::V0(v0) => v0.unshielding_amount,
+        let (amount, num_actions) = match self {
+            UnshieldTransition::V0(v0) => (v0.unshielding_amount, v0.actions.len()),
         };
 
         if current_total_balance < amount {
@@ -121,7 +121,14 @@ impl UnshieldStateTransitionTransformIntoActionValidationV0 for UnshieldTransiti
             ));
         }
 
-        let result = UnshieldTransitionAction::try_from_transition(self, current_total_balance);
+        // The fee charged to the shielded pool is the minimum shielded fee computed
+        // from the same `num_actions` that `validate_minimum_shielded_fee` enforced
+        // `unshielding_amount >=` against. Because that check passed, the net recipient
+        // amount (`unshielding_amount - fee_amount`) is guaranteed to be non-negative.
+        let fee_amount = dpp::shielded::compute_minimum_shielded_fee(num_actions, platform_version);
+
+        let result =
+            UnshieldTransitionAction::try_from_transition(self, current_total_balance, fee_amount);
 
         Ok(result.map(|action| action.into()))
     }

packages/rs-drive/src/state_transition_action/action_convert_to_operations/shielded/shielded_withdrawal_transition.rs

@@ -32,20 +32,17 @@ impl DriveHighLevelOperationConverter for ShieldedWithdrawalTransitionAction {
                     // 2. Insert change notes into CommitmentTree
                     insert_notes(&mut ops, &v0.notes);
 
-                    // 3. Update total balance: subtract withdrawal amount + fee (both leave the pool)
-                    let total_deduction =
-                        v0.amount.checked_add(v0.fee_amount).ok_or_else(|| {
-                            Error::Drive(DriveError::CorruptedDriveState(
-                                "overflow when adding shielded_withdrawal amount and fee"
-                                    .to_string(),
-                            ))
-                        })?;
+                    // 3. Update total balance: the pool decreases by the full `amount`
+                    //    (= unshielding_amount). Of that, `amount - fee_amount` leaves the
+                    //    platform to Core (see RemoveFromSystemCredits below) and
+                    //    `fee_amount` stays in-platform, flowing to the fee pools at block
+                    //    finalization, so credits are conserved.
                     let new_total_balance =
                         v0.current_total_balance
-                            .checked_sub(total_deduction)
+                            .checked_sub(v0.amount)
                             .ok_or_else(|| {
                                 Error::Drive(DriveError::CorruptedDriveState(
-                                "shielded pool total balance underflow when subtracting shielded_withdrawal amount and fee"
+                                "shielded pool total balance underflow when subtracting shielded_withdrawal amount"
                                     .to_string(),
                             ))
                             })?;
@@ -64,9 +61,22 @@ impl DriveHighLevelOperationConverter for ShieldedWithdrawalTransitionAction {
                         },
                     ));
 
-                    // 5. Remove credits from system (they leave the system to Core)
+                    // 5. Remove credits from the system: only the NET amount
+                    //    (`amount - fee_amount`) actually leaves the platform to Core. The
+                    //    fee stays in-platform (moves from the shielded pool sum tree to the
+                    //    fee pools sum tree), so the `total_credits_in_platform` counter must
+                    //    only drop by the net. Validation guarantees `amount >= fee_amount`;
+                    //    we still guard the subtraction defensively.
+                    let net_withdrawal_amount =
+                        v0.amount.checked_sub(v0.fee_amount).ok_or_else(|| {
+                            Error::Drive(DriveError::CorruptedDriveState(
+                                "shielded_withdrawal fee exceeds withdrawal amount".to_string(),
+                            ))
+                        })?;
                     ops.push(DriveOperation::SystemOperation(
-                        SystemOperationType::RemoveFromSystemCredits { amount: v0.amount },
+                        SystemOperationType::RemoveFromSystemCredits {
+                            amount: net_withdrawal_amount,
+                        },
                     ));
 
                     Ok(ops)
@@ -175,12 +185,12 @@ mod tests {
             other => panic!("expected InsertNote, got {:?}", other),
         }
 
-        // Verify UpdateTotalBalance = 10000 - 3000 - 500 = 6500
+        // Verify UpdateTotalBalance = 10000 - 3000 (amount only) = 7000
         match &ops[2] {
             DriveOperation::ShieldedPoolOperation(
                 ShieldedPoolOperationType::UpdateTotalBalance { new_total_balance },
             ) => {
-                assert_eq!(*new_total_balance, 6500);
+                assert_eq!(*new_total_balance, 7000);
             }
             other => panic!("expected UpdateTotalBalance, got {:?}", other),
         }
@@ -198,19 +208,21 @@ mod tests {
             other => panic!("expected AddWithdrawalDocument, got {:?}", other),
         }
 
-        // Verify RemoveFromSystemCredits amount = 3000
+        // Verify RemoveFromSystemCredits amount = NET (3000 - 500) = 2500
         match &ops[4] {
             DriveOperation::SystemOperation(SystemOperationType::RemoveFromSystemCredits {
                 amount,
             }) => {
-                assert_eq!(*amount, 3000);
+                assert_eq!(*amount, 2500);
             }
             other => panic!("expected RemoveFromSystemCredits, got {:?}", other),
         }
     }
 
     #[test]
-    fn test_balance_decreases_by_amount_plus_fee() {
+    fn test_balance_decreases_by_amount_only() {
+        // The pool decrements by exactly `amount` (= unshielding_amount); the fee is
+        // carved out of that amount, not charged on top of it.
         let action = make_action();
         let epoch = Epoch::new(0).unwrap();
         let platform_version = PlatformVersion::latest();
@@ -232,7 +244,7 @@ mod tests {
             DriveOperation::ShieldedPoolOperation(
                 ShieldedPoolOperationType::UpdateTotalBalance { new_total_balance },
             ) => {
-                assert_eq!(*new_total_balance, 6500); // 10000 - 3000 - 500
+                assert_eq!(*new_total_balance, 7000); // 10000 - 3000 (amount only)
             }
             _ => unreachable!(),
         }
@@ -278,7 +290,9 @@ mod tests {
     }
 
     #[test]
-    fn test_removes_from_system_credits() {
+    fn test_removes_net_from_system_credits() {
+        // Only the NET amount (amount - fee) leaves the platform to Core; the fee
+        // stays in-platform and flows to the fee pools.
         let action = make_action();
         let epoch = Epoch::new(0).unwrap();
         let platform_version = PlatformVersion::latest();
@@ -291,23 +305,91 @@ mod tests {
             DriveOperation::SystemOperation(SystemOperationType::RemoveFromSystemCredits {
                 amount,
             }) => {
-                assert_eq!(*amount, 3000);
+                assert_eq!(*amount, 2500); // 3000 amount - 500 fee
             }
             other => panic!("expected RemoveFromSystemCredits, got {:?}", other),
         }
     }
 
     #[test]
-    fn test_underflow_returns_error() {
+    fn test_fee_amount_is_non_zero() {
+        // Regression guard for the fee-bypass bug: withdrawal must charge a non-zero fee.
+        let action = make_action();
+        match action {
+            ShieldedWithdrawalTransitionAction::V0(v0) => assert!(v0.fee_amount > 0),
+        }
+    }
+
+    #[test]
+    fn test_conservation_pool_minus_amount_system_minus_net() {
+        // Conservation at the operation level:
+        //   - shielded pool sum tree:        -amount
+        //   - RemoveFromSystemCredits (counter): -(amount - fee)
+        // The fee (= amount - net) is reconciled into the fee pools sum tree at block
+        // finalization, so both the sum-tree total and the counter ultimately drop by
+        // exactly the net (amount - fee).
+        let amount = 3000u64;
+        let fee = 500u64;
+        let action = make_action();
+        let epoch = Epoch::new(0).unwrap();
+        let platform_version = PlatformVersion::latest();
+
+        let ops = action
+            .into_high_level_drive_operations(&epoch, platform_version)
+            .expect("expected operations");
+
+        let mut shielded_delta: i128 = 0;
+        let mut system_credit_delta: i128 = 0;
+        for op in &ops {
+            match op {
+                DriveOperation::ShieldedPoolOperation(
+                    ShieldedPoolOperationType::UpdateTotalBalance { new_total_balance },
+                ) => shielded_delta = *new_total_balance as i128 - 10000i128,
+                DriveOperation::SystemOperation(SystemOperationType::RemoveFromSystemCredits {
+                    amount,
+                }) => system_credit_delta = -(*amount as i128),
+                _ => {}
+            }
+        }
+
+        assert_eq!(shielded_delta, -(amount as i128));
+        assert_eq!(system_credit_delta, -((amount - fee) as i128));
+    }
+
+    #[test]
+    fn test_pool_underflow_returns_error() {
+        // Pool has less than `amount`; pool decrement must error.
         let action = ShieldedWithdrawalTransitionAction::V0(ShieldedWithdrawalTransitionActionV0 {
             amount: 5000,
             notes: vec![],
             anchor: [0x00; 32],
             core_fee_per_byte: 1,
             pooling: Pooling::Never,
             output_script: CoreScript::from_bytes(vec![]),
-            fee_amount: 6000,
-            current_total_balance: 10000, // 5000 + 6000 > 10000
+            fee_amount: 500,
+            current_total_balance: 4000, // 4000 < 5000 (amount)
+            prepared_withdrawal_document: make_document(),
+        });
+        let epoch = Epoch::new(0).unwrap();
+        let platform_version = PlatformVersion::latest();
+
+        let result = action.into_high_level_drive_operations(&epoch, platform_version);
+        assert!(result.is_err());
+    }
+
+    #[test]
+    fn test_fee_exceeds_amount_returns_error() {
+        // Defensive guard: if fee somehow exceeds amount, the net-withdrawal
+        // subtraction in RemoveFromSystemCredits must error rather than underflow.
+        let action = ShieldedWithdrawalTransitionAction::V0(ShieldedWithdrawalTransitionActionV0 {
+            amount: 100,
+            notes: vec![],
+            anchor: [0x00; 32],
+            core_fee_per_byte: 1,
+            pooling: Pooling::Never,
+            output_script: CoreScript::from_bytes(vec![]),
+            fee_amount: 500, // fee > amount
+            current_total_balance: 10000,
             prepared_withdrawal_document: make_document(),
         });
         let epoch = Epoch::new(0).unwrap();