feat: Test and fix logic for cancelling subscriptions and withdrawing…

… funds from subscriptions

src/SafeSubscriptions.sol

@@ -20,6 +20,7 @@ contract SafeSubscriptions is EIP712, ISafeSubscriptions {
     Safe private s_safe;
     uint256 private s_nonce;
     mapping(bytes32 subscriptionDataHash => Subscription subscription) private s_subscriptions;
+    mapping(bytes32 subscriptionDataHash => bool isCancelled) private s_isCancelled;
 
     constructor(address _safe) EIP712("Safe Subscriptions", "1") {
         s_safe = Safe(payable(_safe));
@@ -32,6 +33,7 @@ contract SafeSubscriptions is EIP712, ISafeSubscriptions {
         bytes memory _signatures
     )
         external
+        returns (bytes32)
     {
         if (
             _subscription.serviceProvider == address(0) || _subscription.token == address(0)
@@ -59,6 +61,8 @@ contract SafeSubscriptions is EIP712, ISafeSubscriptions {
         s_nonce++;
 
         emit SubscriptionCreated(_subscription);
+
+        return subscriptionDataHash;
     }
 
     function cancelSubscription(
@@ -85,7 +89,7 @@ contract SafeSubscriptions is EIP712, ISafeSubscriptions {
         _checkDeadline(_deadline);
         _checkNonce(_nonce);
 
-        delete s_subscriptions[_subscriptionDataHash];
+        s_isCancelled[_subscriptionDataHash] = true;
         s_nonce++;
 
         emit SubscriptionCancelled(_subscriptionDataHash);
@@ -94,20 +98,23 @@ contract SafeSubscriptions is EIP712, ISafeSubscriptions {
     function withdrawFromSubscription(bytes32 _subscriptionDataHash) external {
         Subscription memory subscription = s_subscriptions[_subscriptionDataHash];
 
+        if (subscription.serviceProvider == address(0)) revert SubscriptionDoesNotExist(_subscriptionDataHash);
+        if (s_isCancelled[_subscriptionDataHash]) revert SubscriptionRevoked();
         if (subscription.startingTimestamp > block.timestamp) {
             revert SubscriptionHasNotStartedYet(_subscriptionDataHash);
         }
 
         uint256 roundsToClaim;
         if (subscription.isRecurring) {
-            roundsToClaim = (block.timestamp - subscription.startingTimestamp / subscription.duration)
+            roundsToClaim = ((block.timestamp - subscription.startingTimestamp) / subscription.duration)
                 - subscription.roundsClaimedSoFar;
         } else {
-            roundsToClaim = (block.timestamp - subscription.startingTimestamp / subscription.duration);
+            roundsToClaim = (block.timestamp - subscription.startingTimestamp) / subscription.duration;
             if (roundsToClaim > subscription.rounds) roundsToClaim = subscription.rounds;
             roundsToClaim -= subscription.roundsClaimedSoFar;
         }
         uint256 amountToWithdraw = roundsToClaim * subscription.amount;
+        if (amountToWithdraw == 0) revert ZeroAmountToWithdraw();
         s_subscriptions[_subscriptionDataHash].roundsClaimedSoFar += roundsToClaim;
 
         bool success;
@@ -171,10 +178,14 @@ contract SafeSubscriptions is EIP712, ISafeSubscriptions {
         return s_nonce + 1;
     }
 
-    function getSbscriptionData(bytes32 _subscriptionDataHash) external view returns (Subscription memory) {
+    function getSubscriptionData(bytes32 _subscriptionDataHash) external view returns (Subscription memory) {
         return s_subscriptions[_subscriptionDataHash];
     }
 
+    function isSubscriptionCancelled(bytes32 _subscriptionDataHash) external view returns (bool) {
+        return s_isCancelled[_subscriptionDataHash];
+    }
+
     function getEncodedSubscriptionDataAndHash(
         Subscription memory _subscription,
         uint256 _deadline,

src/interfaces/ISafeSubscriptions.sol

@@ -21,8 +21,10 @@ interface ISafeSubscriptions {
     error InvalidSubscription();
     error SubscriptionAlreadyExists(bytes32 subscriptionDataHash);
     error SubscriptionDoesNotExist(bytes32 subscriptionDataHash);
+    error SubscriptionRevoked();
     error SubscriptionHasNotStartedYet(bytes32 subscriptionDataHash);
+    error ZeroAmountToWithdraw();
     error TransactionFailed();
     error DeadlinePassed(uint256 deadline, uint256 currentTimestamp);
-    error InvalidNonce(uint256 givenNonce, uint256 _expectedNonce);
+    error InvalidNonce(uint256 givenNonce, uint256 expectedNonce);
 }

test/unit/CancelSubscription.t.sol

@@ -0,0 +1,72 @@
+// SPDX-License-Identifier: MIT
+pragma solidity 0.8.24;
+
+import { ISafeSubscriptions } from "../../src/interfaces/ISafeSubscriptions.sol";
+
+import { GlobalHelper } from "../utils/GlobalHelper.sol";
+
+contract CancelSubscriptionTest is GlobalHelper {
+    function test_cancellingSubscriptionFailsIfSubscriptionDoesNotExist() public {
+        (ISafeSubscriptions.Subscription memory subscriptionData,,) = _getTestCreateSubscriptionData();
+        bytes32 subscriptionDataHash = safeSubscriptions.getSubscriptionDataHash(subscriptionData);
+
+        vm.expectRevert(
+            abi.encodeWithSelector(ISafeSubscriptions.SubscriptionDoesNotExist.selector, subscriptionDataHash)
+        );
+        safeSubscriptions.cancelSubscription(subscriptionDataHash, 0, 0, "");
+    }
+
+    function test_cancellingSubscriptionFailsIfDeadlineHasPassed() public {
+        bytes32 subscriptionDataHash = _createTestSubscription();
+        ISafeSubscriptions.Subscription memory subscriptionData =
+            safeSubscriptions.getSubscriptionData(subscriptionDataHash);
+
+        _warpBy(delay);
+
+        uint256 deadline = block.timestamp - 1;
+        uint256 nonce = safeSubscriptions.getNextNonce();
+        bytes memory signatures = _getSignaturesForSubscriptionOperation(subscriptionData, deadline, nonce);
+
+        vm.expectRevert(abi.encodeWithSelector(ISafeSubscriptions.DeadlinePassed.selector, deadline, block.timestamp));
+        safeSubscriptions.cancelSubscription(subscriptionDataHash, deadline, nonce, signatures);
+    }
+
+    function test_cancellingSubscriptionFailsIfInvalidNonceIsPassed() public {
+        bytes32 subscriptionDataHash = _createTestSubscription();
+        ISafeSubscriptions.Subscription memory subscriptionData =
+            safeSubscriptions.getSubscriptionData(subscriptionDataHash);
+        uint256 deadline = block.timestamp + delay;
+        uint256 nonce = safeSubscriptions.getNextNonce() + 1;
+        bytes memory signatures = _getSignaturesForSubscriptionOperation(subscriptionData, deadline, nonce);
+
+        vm.expectRevert(abi.encodeWithSelector(ISafeSubscriptions.InvalidNonce.selector, nonce, nonce - 1));
+        safeSubscriptions.cancelSubscription(subscriptionDataHash, deadline, nonce, signatures);
+    }
+
+    function test_cancellingSubscriptionSucceeds() public {
+        bytes32 subscriptionDataHash = _createTestSubscription();
+        ISafeSubscriptions.Subscription memory subscriptionData =
+            safeSubscriptions.getSubscriptionData(subscriptionDataHash);
+        uint256 deadline = block.timestamp + delay;
+        uint256 nonce = safeSubscriptions.getNextNonce();
+        bytes memory signatures = _getSignaturesForSubscriptionOperation(subscriptionData, deadline, nonce);
+
+        safeSubscriptions.cancelSubscription(subscriptionDataHash, deadline, nonce, signatures);
+
+        assertEq(safeSubscriptions.getNextNonce(), nonce + 1);
+        assertTrue(safeSubscriptions.isSubscriptionCancelled(subscriptionDataHash));
+    }
+
+    function test_cancellingSubscriptionEmitsEvent() public {
+        bytes32 subscriptionDataHash = _createTestSubscription();
+        ISafeSubscriptions.Subscription memory subscriptionData =
+            safeSubscriptions.getSubscriptionData(subscriptionDataHash);
+        uint256 deadline = block.timestamp + delay;
+        uint256 nonce = safeSubscriptions.getNextNonce();
+        bytes memory signatures = _getSignaturesForSubscriptionOperation(subscriptionData, deadline, nonce);
+
+        vm.expectEmit(true, true, true, true);
+        emit ISafeSubscriptions.SubscriptionCancelled(subscriptionDataHash);
+        safeSubscriptions.cancelSubscription(subscriptionDataHash, deadline, nonce, signatures);
+    }
+}

test/unit/CreateSubscription.t.sol

@@ -82,7 +82,7 @@ contract CreateSubscriptionTest is GlobalHelper {
         (ISafeSubscriptions.Subscription memory subscription, uint256 deadline, uint256 nonce) =
             _getTestCreateSubscriptionData();
         subscription.startingTimestamp = block.timestamp + 2 * delay;
-        bytes memory signatures = _getSignaturesForSubscriptionCreation(subscription, deadline, nonce);
+        bytes memory signatures = _getSignaturesForSubscriptionOperation(subscription, deadline, nonce);
 
         _warpBy(delay + 1);
 
@@ -94,7 +94,7 @@ contract CreateSubscriptionTest is GlobalHelper {
         (ISafeSubscriptions.Subscription memory subscription, uint256 deadline, uint256 nonce) =
             _getTestCreateSubscriptionData();
         nonce = 10;
-        bytes memory signatures = _getSignaturesForSubscriptionCreation(subscription, deadline, nonce);
+        bytes memory signatures = _getSignaturesForSubscriptionOperation(subscription, deadline, nonce);
 
         vm.expectRevert(
             abi.encodeWithSelector(ISafeSubscriptions.InvalidNonce.selector, nonce, safeSubscriptions.getNextNonce())
@@ -105,12 +105,12 @@ contract CreateSubscriptionTest is GlobalHelper {
     function test_creatingANewSubscriptionSucceeds() public {
         (ISafeSubscriptions.Subscription memory subscription, uint256 deadline, uint256 nonce) =
             _getTestCreateSubscriptionData();
-        bytes memory signatures = _getSignaturesForSubscriptionCreation(subscription, deadline, nonce);
+        bytes memory signatures = _getSignaturesForSubscriptionOperation(subscription, deadline, nonce);
 
-        safeSubscriptions.createSubscription(subscription, deadline, nonce, signatures);
+        bytes32 subscriptionDataHash = safeSubscriptions.createSubscription(subscription, deadline, nonce, signatures);
 
         ISafeSubscriptions.Subscription memory subscriptionData =
-            safeSubscriptions.getSbscriptionData(safeSubscriptions.getSubscriptionDataHash(subscription));
+            safeSubscriptions.getSubscriptionData(subscriptionDataHash);
         assertEq(subscriptionData.serviceProvider, subscription.serviceProvider);
         assertEq(subscriptionData.token, subscription.token);
         assertEq(subscriptionData.amount, subscription.amount);
@@ -127,7 +127,7 @@ contract CreateSubscriptionTest is GlobalHelper {
     function test_creatingANewSubscriptionEmitsEvent() public {
         (ISafeSubscriptions.Subscription memory subscription, uint256 deadline, uint256 nonce) =
             _getTestCreateSubscriptionData();
-        bytes memory signatures = _getSignaturesForSubscriptionCreation(subscription, deadline, nonce);
+        bytes memory signatures = _getSignaturesForSubscriptionOperation(subscription, deadline, nonce);
 
         vm.expectEmit(true, true, true, true);
         emit ISafeSubscriptions.SubscriptionCreated(subscription);