Final Audit Fixes (#20)

src/FoldCaptiveStaking.sol

@@ -43,6 +43,7 @@ contract FoldCaptiveStaking is Owned(msg.sender) {
     error AlreadyInitialized();
     error NotInitialized();
     error ZeroLiquidity();
+    error NotEnoughLiquidity();
     error WithdrawFailed();
     error DepositCapReached();
     error DepositAmountBelowMinimum();
@@ -60,13 +61,11 @@ contract FoldCaptiveStaking is Owned(msg.sender) {
         positionManager = INonfungiblePositionManager(_positionManager);
         POOL = IUniswapV3Pool(_pool);
 
-        token0 = ERC20(POOL.token0());
-        token1 = ERC20(POOL.token1());
+        token0 = ERC20(IUniswapV3Pool(_pool).token0());
+        token1 = ERC20(IUniswapV3Pool(_pool).token1());
 
         WETH9 = WETH(payable(_weth));
         FOLD = ERC20(_fold);
-
-        initialized = false;
     }
 
     /// @notice Initialize the contract by minting a small initial liquidity position
@@ -293,6 +292,7 @@ contract FoldCaptiveStaking is Owned(msg.sender) {
     /// @param liquidity The amount of liquidity to withdraw
     function withdraw(uint128 liquidity) external isInitialized {
         if (block.timestamp < depositTimeStamp[msg.sender] + COOLDOWN_PERIOD) revert WithdrawalCooldownPeriodNotMet();
+        if (liquidity > liquidityUnderManagement) revert NotEnoughLiquidity();
 
         collectFees();
         collectRewards();

test/UnitTests.t.sol

@@ -6,6 +6,19 @@ import "test/interfaces/ISwapRouter.sol";
 contract UnitTests is BaseCaptiveTest {
     ISwapRouter public router = ISwapRouter(0xE592427A0AEce92De3Edee1F18E0157C05861564);
 
+    /// @dev Test the constructor that you cannot supply a zero address
+    function testNoZeroAddressInitialization() public {
+        vm.expectRevert(FoldCaptiveStaking.ZeroAddress.selector);
+        new FoldCaptiveStaking(address(0), address(0), address(0), address(0));
+    }
+
+    /// @dev Test that you cannot initialize the same contract twice 
+    function testNoDoubleInitialization() public {
+        vm.startPrank(foldCaptiveStaking.owner());
+        vm.expectRevert(FoldCaptiveStaking.AlreadyInitialized.selector);
+        foldCaptiveStaking.initialize();
+    }
+
     /// @dev Ensure that balances and state variables are updated correctly.
     function testAddLiquidity() public {
         fold.transfer(User01, 1_000 ether);