From 5a137b107b2f3ef2fa4e2d0881c660777473227c Mon Sep 17 00:00:00 2001
From: Zied <26070035+zguesmi@users.noreply.github.com>
Date: Mon, 10 Jun 2024 11:38:53 +0200
Subject: [PATCH] Check balance before drain

---
 contracts/VoucherHub.sol |  3 +--
 test/VoucherHub.test.ts  | 16 +++++++++++++++-
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/contracts/VoucherHub.sol b/contracts/VoucherHub.sol
index 5375c7b..7bbece2 100644
--- a/contracts/VoucherHub.sol
+++ b/contracts/VoucherHub.sol
@@ -259,9 +259,8 @@ contract VoucherHub is
      * @param voucher address of the expired voucher to drain
      */
     function drainVoucher(address voucher) external {
-        VoucherHubStorage storage $ = _getVoucherHubStorage();
-        require($._isVoucher[voucher], "VoucherHub: unknown voucher");
         uint256 amount = balanceOf(voucher);
+        require(amount > 0, "VoucherHub: nothing to drain");
         _burn(voucher, amount);
         emit VoucherDrained(voucher, amount);
         Voucher(voucher).drain(amount);
diff --git a/test/VoucherHub.test.ts b/test/VoucherHub.test.ts
index f65a4e7..875d596 100644
--- a/test/VoucherHub.test.ts
+++ b/test/VoucherHub.test.ts
@@ -1016,7 +1016,21 @@ describe('VoucherHub', function () {
                 voucherHubWithVoucherManagerSigner.drainVoucher(
                     ethers.Wallet.createRandom().address,
                 ),
-            ).to.be.revertedWith('VoucherHub: unknown voucher');
+            ).to.be.revertedWith('VoucherHub: nothing to drain');
+        });
+
+        it('Should not drain if balance is empty', async function () {
+            // Expire voucher
+            const expirationDate = await voucher.getExpiration();
+            await time.setNextBlockTimestamp(expirationDate + 100n); // after expiration
+            // Drain to empty the voucher from its balance.
+            await voucherHubWithVoucherManagerSigner
+                .drainVoucher(voucherAddress)
+                .then((tx) => tx.wait());
+            // Drain a second time when balance is empty.
+            await expect(
+                voucherHubWithVoucherManagerSigner.drainVoucher(voucherAddress),
+            ).to.be.revertedWith('VoucherHub: nothing to drain');
         });
     });