Merge pull request #120 from DyadStablecoin/fix/H-01

Fix/h 01
DyadStablecoin · Oct 23, 2024 · 2bec494 · 2bec494
2 parents d5ce96a + 6f7753f
commit 2bec494
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 2 deletions.
diff --git a/src/periphery/AtomicSwapExtension.sol b/src/periphery/AtomicSwapExtension.sol
@@ -1,24 +1,29 @@
 // SPDX-License-Identifier: MIT
 pragma solidity ^0.8.27;
 
+import {IERC721} from "forge-std/interfaces/IERC721.sol";
 import {IExtension, IAfterWithdrawHook} from "../interfaces/IExtension.sol";
 import {IVaultManager} from "../interfaces/IVaultManager.sol";
 import {IVault} from "../interfaces/IVault.sol";
 import {DyadHooks} from "../core/DyadHooks.sol";
 import {ERC20} from "solmate/tokens/ERC20.sol";
 import {SafeTransferLib} from "solmate/utils/SafeTransferLib.sol";
 
+
 contract AtomicSwapExtension is IExtension, IAfterWithdrawHook {
     using SafeTransferLib for ERC20;
 
     error SwapFailed();
     error OnlyVaultManager();
+    error NotOwnerOfNote();
 
     address public constant AUGUSTUS_6_2 = 0x6A000F20005980200259B80c5102003040001068;
     IVaultManager public immutable vaultManager;
+    IERC721 public immutable dnft;
 
-    constructor(address _vaultManager) {
+    constructor(address _vaultManager, address _dnft) {
         vaultManager = IVaultManager(_vaultManager);
+        dnft = IERC721(_dnft);
     }
 
     function name() external pure returns (string memory) {
@@ -41,6 +46,8 @@ contract AtomicSwapExtension is IExtension, IAfterWithdrawHook {
         uint256 toAmount,
         bytes calldata swapData
     ) external {
+        if (dnft.ownerOf(noteId) != msg.sender) revert NotOwnerOfNote();
+
         // tstore swap data
         assembly {
             let size := calldatasize()

diff --git a/test/fork/v5/AtomicSwapExtension.t.sol b/test/fork/v5/AtomicSwapExtension.t.sol
@@ -10,7 +10,7 @@ contract AtomicSwapExtensionTest is Test, Parameters {
     address public USER = 0xd0953aC488190BbD59240E0A59F697e32e303532;
 
     function setUp() public {
-        atomicSwapExtension = new AtomicSwapExtension(MAINNET_V2_VAULT_MANAGER);
+        atomicSwapExtension = new AtomicSwapExtension(MAINNET_V2_VAULT_MANAGER, MAINNET_DNFT);
         VaultManagerV5 vaultManager = VaultManagerV5(MAINNET_V2_VAULT_MANAGER);
 
         vm.createSelectFork(vm.envString("ETH_RPC_URL"), 20917569);