Rework makeArbitraryTransaction to avoid reentrancy risk

JoinColony · Oct 14, 2024 · 06c4f1a · 06c4f1a
1 parent e3d9bdd
commit 06c4f1a
Showing 1 changed file with 27 additions and 13 deletions.
diff --git a/contracts/colony/ColonyArbitraryTransaction.sol b/contracts/colony/ColonyArbitraryTransaction.sol
@@ -95,30 +95,48 @@ contract ColonyArbitraryTransaction is ColonyStorage {
     } catch {}
 
     bool res = executeCall(_to, 0, _action);
-
+    assert(res);
     if (sig == APPROVE_SIG) {
-      approveTransactionCleanup(_to, _action);
+      approveTransactionCheck(_to, _action);
     }
 
     emit ArbitraryTransaction(_to, _action, res);
 
     return res;
   }
 
-  function approveTransactionPreparation(address _to, bytes memory _action) internal {
+  function approveTransactionPreparation(address _token, bytes memory _action) internal {
     address spender;
+    uint256 amount;
     assembly {
       spender := mload(add(_action, 0x24))
+      amount := mload(add(_action, 0x44))
     }
-    updateApprovalAmountInternal(_to, spender, false);
+    updateApprovalAmountInternal(_token, spender);
+
+    // Calculate and set the approval we expect to have after the approve is made
+    tokenApprovalTotals[_token] =
+      (tokenApprovalTotals[_token] - tokenApprovals[_token][spender]) +
+      amount;
+    require(
+      fundingPots[1].balance[_token] >= tokenApprovalTotals[_token],
+      "colony-approval-exceeds-balance"
+    );
+
+    tokenApprovals[_token][spender] = amount;
   }
 
-  function approveTransactionCleanup(address _to, bytes memory _action) internal {
+  function approveTransactionCheck(address _token, bytes memory _action) internal view {
     address spender;
+    uint256 amount;
     assembly {
       spender := mload(add(_action, 0x24))
+      amount := mload(add(_action, 0x44))
     }
-    updateApprovalAmountInternal(_to, spender, true);
+
+    uint256 recordedApproval = tokenApprovals[_token][spender];
+    uint256 actualApproval = ERC20Extended(_token).allowance(address(this), spender);
+    require(recordedApproval == actualApproval, "colony-approval-amount-mismatch");
   }
 
   function burnTransactionPreparation(address _to, bytes memory _action) internal {
@@ -148,21 +166,17 @@ contract ColonyArbitraryTransaction is ColonyStorage {
   }
 
   function updateApprovalAmount(address _token, address _spender) public stoppable {
-    updateApprovalAmountInternal(_token, _spender, false);
+    updateApprovalAmountInternal(_token, _spender);
   }
 
-  function updateApprovalAmountInternal(
-    address _token,
-    address _spender,
-    bool _postApproval
-  ) internal {
+  function updateApprovalAmountInternal(address _token, address _spender) internal {
     uint256 recordedApproval = tokenApprovals[_token][_spender];
     uint256 actualApproval = ERC20Extended(_token).allowance(address(this), _spender);
     if (recordedApproval == actualApproval) {
       return;
     }
 
-    if (recordedApproval > actualApproval && !_postApproval) {
+    if (recordedApproval > actualApproval) {
       // They've spend some tokens out of root. Adjust balances accordingly
       // If we are post approval, then they have not spent tokens
       fundingPots[1].balance[_token] =