Allow drain only by Hub to keep balances consistent

iExecBlockchainComputing · Jun 7, 2024 · 287640b · 287640b
1 parent 3e0d4d9
commit 287640b
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 2 deletions.
diff --git 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
@@ -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(