fee-abstraction: clear residual allowance after eager fee collection

packages/fee-abstraction/src/storage.rs

@@ -204,6 +204,25 @@ pub fn collect_fee(
 
     token_client.transfer_from(&e.current_contract_address(), user, fee_recipient, &fee_amount);
 
+    // In `Eager` mode the approval is scoped to this single collection, so the
+    // unspent `max_fee_amount - fee_amount` is consumed back to the user to
+    // leave no residual allowance for the target invocation to abuse. In `Lazy`
+    // mode the remaining allowance is intentional and left untouched.
+    //
+    // The self transfer still moves the amount out of and back into the same
+    // balance, so the token rejects it when the user cannot cover it. Only what
+    // the balance allows is consumed, and the allowance left over is by
+    // definition not spendable by the user's current balance.
+    if let FeeAbstractionApproval::Eager = approval {
+        let residual = max_fee_amount - fee_amount;
+        if residual > 0 {
+            let consumable = residual.min(token_client.balance(user));
+            if consumable > 0 {
+                token_client.transfer_from(&e.current_contract_address(), user, user, &consumable);
+            }
+        }
+    }
+
     emit_fee_collected(e, user, fee_recipient, fee_token, fee_amount);
 }
 

packages/fee-abstraction/src/test.rs

@@ -71,16 +71,98 @@ fn collect_fee_with_eager_approval_overwrites_allowance() {
     });
 
     let events = e.events().all();
-    // approval, transfer and collect fee
-    assert_eq!(events.events().len(), 3);
+    // approval, fee transfer, residual refund and collect fee
+    assert_eq!(events.events().len(), 4);
 
+    // the unspent allowance is consumed back to the user in eager mode
     let allowance = token_client.allowance(&user, &contract_address);
-    assert_eq!(allowance, 30);
+    assert_eq!(allowance, 0);
 
     let balance = token_client.balance(&recipient);
     assert_eq!(balance, 20);
 }
 
+#[test]
+fn collect_fee_with_eager_approval_consumes_only_the_affordable_residual() {
+    let e = Env::default();
+    e.mock_all_auths_allowing_non_root_auth();
+
+    let contract_address = e.register(MockContract, ());
+    let user = Address::generate(&e);
+    let token_address = e.register(MockToken, (user.clone(),));
+    let recipient = Address::generate(&e);
+    let other = Address::generate(&e);
+
+    let max_fee_amount = 50;
+
+    let token_client = TokenClient::new(&e, &token_address);
+    // the user keeps 25, which covers the fee but not the whole residual
+    token_client.transfer(&user, &other, &975);
+
+    e.as_contract(&contract_address, || {
+        // approve 50, spend 20
+        collect_fee(
+            &e,
+            &token_address,
+            20,
+            max_fee_amount,
+            100,
+            &user,
+            &recipient,
+            FeeAbstractionApproval::Eager,
+        );
+    });
+
+    // only the 5 left after the fee are consumed, the rest of the allowance stays
+    let allowance = token_client.allowance(&user, &contract_address);
+    assert_eq!(allowance, 25);
+
+    assert_eq!(token_client.balance(&user), 5);
+    assert_eq!(token_client.balance(&recipient), 20);
+}
+
+#[test]
+fn collect_fee_with_eager_approval_and_no_residual_balance() {
+    let e = Env::default();
+    e.mock_all_auths_allowing_non_root_auth();
+
+    let contract_address = e.register(MockContract, ());
+    let user = Address::generate(&e);
+    let token_address = e.register(MockToken, (user.clone(),));
+    let recipient = Address::generate(&e);
+    let other = Address::generate(&e);
+
+    let max_fee_amount = 50;
+
+    let token_client = TokenClient::new(&e, &token_address);
+    // the user keeps exactly the fee, so nothing is left to consume
+    token_client.transfer(&user, &other, &980);
+
+    e.as_contract(&contract_address, || {
+        // approve 50, spend 20
+        collect_fee(
+            &e,
+            &token_address,
+            20,
+            max_fee_amount,
+            100,
+            &user,
+            &recipient,
+            FeeAbstractionApproval::Eager,
+        );
+    });
+
+    let events = e.events().all();
+    // approval, fee transfer and collect fee, without a residual refund
+    assert_eq!(events.events().len(), 3);
+
+    let allowance = token_client.allowance(&user, &contract_address);
+    assert_eq!(allowance, 30);
+
+    assert_eq!(token_client.balance(&user), 0);
+    assert_eq!(token_client.balance(&recipient), 20);
+}
+
 #[test]
 fn collect_fee_with_lazy_approval_no_previous() {
     let e = Env::default();