Fix re-entrancy in cancelAndMakeLimitOrders

PaintSwap · Feb 8, 2025 · 4ba1e96 · 4ba1e96
1 parent c78ee63
commit 4ba1e96
Show file tree

Hide file tree

Showing 2 changed files with 91 additions and 39 deletions.
diff --git a/contracts/SamWitchOrderBook.sol b/contracts/SamWitchOrderBook.sol
@@ -240,48 +240,20 @@ contract SamWitchOrderBook is UUPSUpgradeable, OwnableUpgradeable, ERC1155Holder
   /// @notice Cancel multiple orders in the order book
   /// @param orderIds Array of order IDs to be cancelled
   /// @param orders Information about the orders so that they can be found in the order book
-  function cancelOrders(uint256[] calldata orderIds, CancelOrder[] calldata orders) public override {
-    require(orderIds.length == orders.length, LengthMismatch());
-
-    address sender = _msgSender();
-
-    uint256 coinsFromUs = 0;
-    uint256 nftsFromUs = 0;
-    uint256 numberOfOrders = orderIds.length;
-    uint256[] memory nftIdsFromUs = new uint256[](numberOfOrders);
-    uint256[] memory nftAmountsFromUs = new uint256[](numberOfOrders);
-    for (uint256 i = 0; i < numberOfOrders; ++i) {
-      CancelOrder calldata cancelOrder = orders[i];
-      (OrderSide side, uint256 tokenId, uint72 price) = (cancelOrder.side, cancelOrder.tokenId, cancelOrder.price);
-
-      if (side == OrderSide.Buy) {
-        uint256 quantity = _cancelOrdersSide(orderIds[i], price, _bidsAtPrice[tokenId][price], _bids[tokenId]);
-        // Send the remaining token back to them
-        coinsFromUs += quantity * price;
-      } else {
-        uint256 quantity = _cancelOrdersSide(orderIds[i], price, _asksAtPrice[tokenId][price], _asks[tokenId]);
-        // Send the remaining NFTs back to them
-        nftIdsFromUs[nftsFromUs] = tokenId;
-        nftAmountsFromUs[nftsFromUs] = quantity;
-        ++nftsFromUs;
-      }
-    }
-
-    emit OrdersCancelled(sender, orderIds);
+  function cancelOrders(uint256[] calldata orderIds, CancelOrder[] calldata orders) external override {
+    (uint256 coinsFromUs, uint256[] memory nftIdsFromUs, uint256[] memory nftAmountsFromUs) = _cancelOrders(
+      orderIds,
+      orders
+    );
 
     // Transfer tokens if there are any to send
     if (coinsFromUs != 0) {
-      _coin.safeTransfer(sender, coinsFromUs);
+      _coin.safeTransfer(_msgSender(), coinsFromUs);
     }
 
     // Send the NFTs
-    if (nftsFromUs != 0) {
-      // shrink the size
-      assembly ("memory-safe") {
-        mstore(nftIdsFromUs, nftsFromUs)
-        mstore(nftAmountsFromUs, nftsFromUs)
-      }
-      _safeBatchTransferNFTsFromUs(sender, nftIdsFromUs, nftAmountsFromUs);
+    if (nftIdsFromUs.length != 0) {
+      _safeBatchTransferNFTsFromUs(_msgSender(), nftIdsFromUs, nftAmountsFromUs);
     }
   }
 
@@ -294,8 +266,21 @@ contract SamWitchOrderBook is UUPSUpgradeable, OwnableUpgradeable, ERC1155Holder
     CancelOrder[] calldata orders,
     LimitOrder[] calldata newOrders
   ) external override {
-    cancelOrders(orderIds, orders);
+    (uint256 coinsFromUs, uint256[] memory nftIdsFromUs, uint256[] memory nftAmountsFromUs) = _cancelOrders(
+      orderIds,
+      orders
+    );
     limitOrders(newOrders);
+
+    // Transfer tokens if there are any to send
+    if (coinsFromUs != 0) {
+      _coin.safeTransfer(_msgSender(), coinsFromUs);
+    }
+
+    // Send the NFTs
+    if (nftIdsFromUs.length != 0) {
+      _safeBatchTransferNFTsFromUs(_msgSender(), nftIdsFromUs, nftAmountsFromUs);
+    }
   }
 
   /// @notice Claim tokens associated with filled or partially filled orders.
@@ -464,6 +449,43 @@ contract SamWitchOrderBook is UUPSUpgradeable, OwnableUpgradeable, ERC1155Holder
     }
   }
 
+  function _cancelOrders(
+    uint256[] calldata orderIds,
+    CancelOrder[] calldata orders
+  ) private returns (uint256 coinsFromUs, uint256[] memory nftIdsFromUs, uint256[] memory nftAmountsFromUs) {
+    require(orderIds.length == orders.length, LengthMismatch());
+
+    address sender = _msgSender();
+
+    uint256 nftsFromUs = 0;
+    uint256 numberOfOrders = orderIds.length;
+    nftIdsFromUs = new uint256[](numberOfOrders);
+    nftAmountsFromUs = new uint256[](numberOfOrders);
+    for (uint256 i = 0; i < numberOfOrders; ++i) {
+      CancelOrder calldata cancelOrder = orders[i];
+      (OrderSide side, uint256 tokenId, uint72 price) = (cancelOrder.side, cancelOrder.tokenId, cancelOrder.price);
+
+      if (side == OrderSide.Buy) {
+        uint256 quantity = _cancelOrdersSide(orderIds[i], price, _bidsAtPrice[tokenId][price], _bids[tokenId]);
+        // Send the remaining token back to them
+        coinsFromUs += quantity * price;
+      } else {
+        uint256 quantity = _cancelOrdersSide(orderIds[i], price, _asksAtPrice[tokenId][price], _asks[tokenId]);
+        // Send the remaining NFTs back to them
+        nftIdsFromUs[nftsFromUs] = tokenId;
+        nftAmountsFromUs[nftsFromUs] = quantity;
+        ++nftsFromUs;
+      }
+    }
+
+    emit OrdersCancelled(sender, orderIds);
+
+    assembly ("memory-safe") {
+      mstore(nftIdsFromUs, nftsFromUs)
+      mstore(nftAmountsFromUs, nftsFromUs)
+    }
+  }
+
   function _takeFromOrderBookSide(
     uint256 tokenId,
     uint72 price,

diff --git a/test/SamWitchOrderBook.ts b/test/SamWitchOrderBook.ts
@@ -3276,7 +3276,7 @@ describe("SamWitchOrderBook", function () {
   });
 
   describe("Edit orders", function () {
-    it("Cancels and makes a new order", async function () {
+    it("Cancels a buy order and makes a new one", async function () {
       const {orderBook, tokenId, tick, owner} = await loadFixture(deployContractsFixture);
 
       // Set up order book
@@ -3291,7 +3291,7 @@ describe("SamWitchOrderBook", function () {
         },
       ]);
 
-      const newOrder = {side: OrderSide.Buy, tokenId, price: price + 1 * tick, quantity: quantity + 2};
+      const newOrder = {side: OrderSide.Buy, tokenId, price: price + tick, quantity: quantity + 2};
       await orderBook.cancelAndMakeLimitOrders([1], [{side: OrderSide.Buy, tokenId, price}], [newOrder]);
 
       expect(await getNextOrderId(orderBook)).to.eq(3);
@@ -3305,6 +3305,36 @@ describe("SamWitchOrderBook", function () {
         orderId,
       ]);
     });
+
+    it("Cancels a sell order and makes a new order", async function () {
+      const {orderBook, tokenId, tick, owner} = await loadFixture(deployContractsFixture);
+
+      // Set up order book
+      const price = 100;
+      const quantity = 10;
+      await orderBook.limitOrders([
+        {
+          side: OrderSide.Sell,
+          tokenId,
+          price,
+          quantity,
+        },
+      ]);
+
+      const newOrder = {side: OrderSide.Sell, tokenId, price: price - tick, quantity: quantity + 2};
+      await orderBook.cancelAndMakeLimitOrders([1], [{side: OrderSide.Sell, tokenId, price}], [newOrder]);
+
+      expect(await getNextOrderId(orderBook)).to.eq(3);
+
+      expect((await orderBook.allOrdersAtPrice(OrderSide.Sell, tokenId, price)).length).to.eq(0);
+      expect((await orderBook.allOrdersAtPrice(OrderSide.Sell, tokenId, price - tick)).length).to.eq(1);
+      const orderId = 2;
+      expect((await orderBook.allOrdersAtPrice(OrderSide.Sell, tokenId, price - tick))[0]).to.deep.eq([
+        owner.address,
+        quantity + 2,
+        orderId,
+      ]);
+    });
   });
 
   it("Max coins price", async function () {