diff --git a/contracts/Collateral.sol b/contracts/Collateral.sol
index 765c58b..5799e3b 100644
--- a/contracts/Collateral.sol
+++ b/contracts/Collateral.sol
@@ -4,7 +4,7 @@ pragma solidity ^0.8.0;
 import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 import "./AccountLocks.sol";
 
-contract Collateral is AccountLocks {
+abstract contract Collateral is AccountLocks {
   IERC20 public immutable token;
   CollateralFunds private funds;
 
@@ -39,8 +39,14 @@ contract Collateral is AccountLocks {
     add(msg.sender, amount);
   }
 
-  function withdraw() public collateralInvariant {
+  // TODO: remove AccountLocks
+  function isWithdrawAllowed() internal virtual returns (bool) {
     _unlockAccount();
+    return true;
+  }
+
+  function withdraw() public collateralInvariant {
+    require(isWithdrawAllowed(), "Account locked");
     uint256 amount = balanceOf(msg.sender);
     funds.withdrawn += amount;
     subtract(msg.sender, amount);
diff --git a/contracts/Marketplace.sol b/contracts/Marketplace.sol
index 4308ac0..5df30ab 100644
--- a/contracts/Marketplace.sol
+++ b/contracts/Marketplace.sol
@@ -50,6 +50,10 @@ contract Marketplace is Collateral, Proofs {
     return _toSlotIds(slotsPerHost[msg.sender].filter(isActive));
   }
 
+  function isWithdrawAllowed() internal view override returns (bool) {
+    return slotsPerHost[msg.sender].length() == 0;
+  }
+
   function _equals(RequestId a, RequestId b) internal pure returns (bool) {
     return RequestId.unwrap(a) == RequestId.unwrap(b);
   }
@@ -178,7 +182,10 @@ contract Marketplace is Collateral, Proofs {
     private
     marketplaceInvariant
   {
-    require(_isFinished(requestId), "Contract not ended");
+    require(
+      _isFinished(requestId) || _isCancelled(requestId),
+      "Contract not ended"
+    );
     RequestContext storage context = _context(requestId);
     Request storage request = _request(requestId);
     context.state = RequestState.Finished;
diff --git a/contracts/Storage.sol b/contracts/Storage.sol
index 2736983..6f76e82 100644
--- a/contracts/Storage.sol
+++ b/contracts/Storage.sol
@@ -5,7 +5,7 @@ import "./Marketplace.sol";
 import "./Proofs.sol";
 import "./Collateral.sol";
 
-contract Storage is Collateral, Marketplace {
+contract Storage is Marketplace {
   uint256 public collateralAmount;
   uint256 public slashMisses;
   uint256 public slashPercentage;
diff --git a/test/Storage.test.js b/test/Storage.test.js
index 668fa7e..ebd3a58 100644
--- a/test/Storage.test.js
+++ b/test/Storage.test.js
@@ -77,6 +77,7 @@ describe("Storage", function () {
     it("unlocks the host collateral", async function () {
       await storage.fillSlot(slot.request, slot.index, proof)
       await waitUntilFinished(storage, slot.request)
+      await storage.freeSlot(slotId(slot))
       await expect(storage.withdraw()).not.to.be.reverted
     })
   })