From 88d880bbe4cc9e51a221d22eaa51d1e312f11dd0 Mon Sep 17 00:00:00 2001 From: Pratik Pawar Date: Thu, 16 Jan 2025 13:21:41 +0530 Subject: [PATCH] Add tests for reentrancy in _rentStorage function --- src/IdGateway.sol | 26 +++++++++++++++++--------- 1 file changed, 17 insertions(+), 9 deletions(-) diff --git a/src/IdGateway.sol b/src/IdGateway.sol index b2d44daa..9629611b 100644 --- a/src/IdGateway.sol +++ b/src/IdGateway.sol @@ -168,18 +168,26 @@ contract IdGateway is IIdGateway, Guardians, Signatures, EIP712, Nonces { //////////////////////////////////////////////////////////////*/ function _rentStorage( - uint256 fid, - uint256 extraUnits, - uint256 payment, - address payer + uint256 fid, + uint256 extraUnits, + uint256 payment, + address payer ) internal returns (uint256 overpayment) { - overpayment = storageRegistry.rent{value: payment}(fid, 1 + extraUnits); - - if (overpayment > 0) { - payer.sendNative(overpayment); - } + // Calculate the overpayment before making any external calls + uint256 amountToRent = 1 + extraUnits; + overpayment = payment - storageRegistry.price(amountToRent); + + // Make the external call to rent storage + storageRegistry.rent{value: payment}(fid, amountToRent); + + // Return the overpayment after the external call + if (overpayment > 0) { + (bool success, ) = payer.call{value: overpayment}(""); + require(success, "Transfer failed"); + } } + receive() external payable { if (msg.sender != address(storageRegistry)) revert Unauthorized(); }