feat: add payment id check

Grateful-protocol · Oct 22, 2024 · c9f1519 · c9f1519
1 parent 52b45a6
commit c9f1519
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 24 deletions.
diff --git a/src/contracts/Grateful.sol b/src/contracts/Grateful.sol
@@ -47,6 +47,9 @@ contract Grateful is IGrateful, Ownable2Step {
   /// @inheritdoc IGrateful
   mapping(address => CustomFee) public override customFees;
 
+  /// @inheritdoc IGrateful
+  mapping(uint256 => bool) public paymentIds;
+
   /*//////////////////////////////////////////////////////////////
                                 MODIFIERS
     //////////////////////////////////////////////////////////////*/
@@ -330,6 +333,11 @@ contract Grateful is IGrateful, Ownable2Step {
       revert Grateful_TransferFailed();
     }
 
+    // Check payment id
+    if (paymentIds[_paymentId]) {
+      revert Grateful_PaymentIdAlreadyUsed();
+    }
+
     // Apply the fee
     uint256 amountWithFee = applyFee(_merchant, _amount);
 
@@ -393,6 +401,8 @@ contract Grateful is IGrateful, Ownable2Step {
       }
     }
 
+    paymentIds[_paymentId] = true;
+
     emit PaymentProcessed(_sender, _merchant, _token, _amount, yieldingFunds[_merchant], _paymentId);
   }
 }
diff --git a/src/interfaces/IGrateful.sol b/src/interfaces/IGrateful.sol
@@ -72,6 +72,9 @@ interface IGrateful {
   /// @notice Thrown when the one-time payment is not found.
   error Grateful_OneTimeNotFound();
 
+  /// @notice Thrown when the payment id has been used.
+  error Grateful_PaymentIdAlreadyUsed();
+
   /*///////////////////////////////////////////////////////////////
                                VARIABLES
     //////////////////////////////////////////////////////////////*/
@@ -127,6 +130,13 @@ interface IGrateful {
     address _merchant
   ) external view returns (bool isSet, uint256 fee);
 
+  /// @notice Returns if a paymentId has been used.
+  /// @param paymentId The payment id.
+  /// @return isUsed If the payment id has been used.
+  function paymentIds(
+    uint256 paymentId
+  ) external view returns (bool isUsed);
+
   /*///////////////////////////////////////////////////////////////
                                 LOGIC
     //////////////////////////////////////////////////////////////*/

diff --git a/test/integration/Grateful.t.sol b/test/integration/Grateful.t.sol
@@ -2,14 +2,19 @@
 pragma solidity 0.8.26;
 
 import {OneTime} from "contracts/OneTime.sol";
-import {IGrateful, IntegrationBase} from "test/integration/IntegrationBase.sol";
+import {IntegrationBase} from "test/integration/IntegrationBase.sol";
 
 contract IntegrationGreeter is IntegrationBase {
-  function test_Payment() public {
-    vm.startPrank(_payer);
-    _usdc.approve(address(_grateful), _amount);
-    _grateful.pay(_merchant, address(_usdc), _amount, _grateful.calculateId(_payer, _merchant, address(_usdc), _amount));
+  function _approveAndPay(address payer, address merchant, uint256 amount) internal {
+    uint256 paymentId = _grateful.calculateId(payer, merchant, address(_usdc), amount);
+    vm.startPrank(payer);
+    _usdc.approve(address(_grateful), amount);
+    _grateful.pay(merchant, address(_usdc), amount, paymentId);
     vm.stopPrank();
+  }
+
+  function test_Payment() public {
+    _approveAndPay(_payer, _merchant, _amount);
 
     assertEq(_usdc.balanceOf(_merchant), _grateful.applyFee(_merchant, _amount));
   }
@@ -22,10 +27,7 @@ contract IntegrationGreeter is IntegrationBase {
 
     assertEq(_grateful.yieldingFunds(_merchant), true);
 
-    vm.startPrank(_payer);
-    _usdc.approve(address(_grateful), _amount);
-    _grateful.pay(_merchant, address(_usdc), _amount, _grateful.calculateId(_payer, _merchant, address(_usdc), _amount));
-    vm.stopPrank();
+    _approveAndPay(_payer, _merchant, _amount);
 
     vm.warp(block.timestamp + 60 days);
 
@@ -93,11 +95,7 @@ contract IntegrationGreeter is IntegrationBase {
     _grateful.setCustomFee(200, _merchant);
 
     // Process payment with custom fee of 2%
-    vm.startPrank(_payer);
-    _usdc.approve(address(_grateful), _amount);
-    uint256 paymentId1 = _grateful.calculateId(_payer, _merchant, address(_usdc), _amount);
-    _grateful.pay(_merchant, address(_usdc), _amount, paymentId1);
-    vm.stopPrank();
+    _approveAndPay(_payer, _merchant, _amount);
 
     // Expected amounts
     uint256 expectedCustomFee = (_amount * 200) / 10_000; // 2% fee
@@ -113,12 +111,11 @@ contract IntegrationGreeter is IntegrationBase {
     vm.prank(_owner);
     _grateful.setCustomFee(0, _merchant);
 
+    // Advance time so calculated paymentId doesn't collide
+    vm.warp(block.timestamp + 1);
+
     // Process payment with custom fee of 0%
-    vm.startPrank(_payer);
-    _usdc.approve(address(_grateful), _amount);
-    uint256 paymentId2 = _grateful.calculateId(_payer, _merchant, address(_usdc), _amount);
-    _grateful.pay(_merchant, address(_usdc), _amount, paymentId2);
-    vm.stopPrank();
+    _approveAndPay(_payer, _merchant, _amount);
 
     // Expected amounts
     uint256 expectedZeroFee = 0; // 0% fee
@@ -140,12 +137,11 @@ contract IntegrationGreeter is IntegrationBase {
     vm.prank(_owner);
     _grateful.unsetCustomFee(_merchant);
 
+    // Advance time so calculated paymentId doesn't collide
+    vm.warp(block.timestamp + 1);
+
     // Process payment after unsetting custom fee
-    vm.startPrank(_payer);
-    _usdc.approve(address(_grateful), _amount);
-    uint256 paymentId3 = _grateful.calculateId(_payer, _merchant, address(_usdc), _amount);
-    _grateful.pay(_merchant, address(_usdc), _amount, paymentId3);
-    vm.stopPrank();
+    _approveAndPay(_payer, _merchant, _amount);
 
     // Expected amounts
     uint256 expectedFeeAfterUnset = (_amount * 100) / 10_000; // 1% fee