From 287640b66980abe722307bb67c0c6208588e76fc Mon Sep 17 00:00:00 2001
From: Zied <26070035+zguesmi@users.noreply.github.com>
Date: Fri, 7 Jun 2024 17:13:03 +0200
Subject: [PATCH] Allow drain only by Hub to keep balances consistent

---
 contracts/beacon/Voucher.sol | 2 +-
 test/beacon/Voucher.test.ts  | 9 ++++++++-
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/contracts/beacon/Voucher.sol b/contracts/beacon/Voucher.sol
index fd6b877..e4e3c0f 100644
--- a/contracts/beacon/Voucher.sol
+++ b/contracts/beacon/Voucher.sol
@@ -273,7 +273,7 @@ contract Voucher is Initializable, IVoucher {
      * Funds are sent to the VoucherHub contract.
      * @param amount amount to be drained
      */
-    function drain(uint256 amount) external onlyExpired {
+    function drain(uint256 amount) external onlyVoucherHub onlyExpired {
         VoucherStorage storage $ = _getVoucherStorage();
         address voucherHub = $._voucherHub;
         // Although transfer function in PoCo always returns true (or reverts),
diff --git a/test/beacon/Voucher.test.ts b/test/beacon/Voucher.test.ts
index c1843ed..daa9c8b 100644
--- a/test/beacon/Voucher.test.ts
+++ b/test/beacon/Voucher.test.ts
@@ -1059,12 +1059,19 @@ describe('Voucher', function () {
             const expirationDate = await voucherAsAnyone.getExpiration();
             await time.setNextBlockTimestamp(expirationDate + 100n); // after expiration
             // Drain
-            await expect(voucherAsAnyone.drain(voucherValue))
+            const voucherHubSigner = await ethers.getImpersonatedSigner(voucherHubAddress);
+            await expect(voucherAsAnyone.connect(voucherHubSigner).drain(voucherValue))
                 .to.emit(iexecPocoInstance, 'Transfer')
                 .withArgs(voucherAddress, voucherHubAddress, voucherValue);
             expect(await iexecPocoInstance.balanceOf(voucherAddress)).to.equal(0);
         });
 
+        it('Should not drain voucher if sender is not authorized', async function () {
+            await expect(voucherAsAnyone.drain(voucherValue)).to.be.revertedWith(
+                'Voucher: sender is not VoucherHub',
+            );
+        });
+
         it('Should not drain voucher if not expired', async function () {
             const voucherHubSigner = await ethers.getImpersonatedSigner(voucherHubAddress);
             await expect(