From 406004a99ad4bdd371610d768e106e3dde46433d Mon Sep 17 00:00:00 2001
From: SylTi <sylvain.tissier@gmail.com>
Date: Thu, 3 Aug 2017 13:21:51 +0200
Subject: [PATCH] refactor HasNoTokens.sol

---
 contracts/ownership/CanReclaimToken.sol | 24 +++++++++++++++++
 contracts/ownership/HasNoTokens.sol     | 14 ++--------
 test/CanReclaimToken.js                 | 35 +++++++++++++++++++++++++
 3 files changed, 61 insertions(+), 12 deletions(-)
 create mode 100644 contracts/ownership/CanReclaimToken.sol
 create mode 100644 test/CanReclaimToken.js

diff --git a/contracts/ownership/CanReclaimToken.sol b/contracts/ownership/CanReclaimToken.sol
new file mode 100644
index 000000000..f057d4394
--- /dev/null
+++ b/contracts/ownership/CanReclaimToken.sol
@@ -0,0 +1,24 @@
+pragma solidity ^0.4.11;
+
+import "./Ownable.sol";
+import "../token/ERC20Basic.sol";
+
+/**
+ * @title Contracts that should be able to recover tokens
+ * @author SylTi
+ * @dev This allow a contract to recover any ERC20 token received in a contract by transfering the balance to the contract owner.
+ * This will prevent any accidental loss of tokens.
+ */
+contract CanReclaimToken is Ownable {
+
+  /**
+   * @dev Reclaim all ERC20Basic compatible tokens
+   * @param tokenAddr address The address of the token contract
+   */
+  function reclaimToken(address tokenAddr) external onlyOwner {
+    ERC20Basic tokenInst = ERC20Basic(tokenAddr);
+    uint256 balance = tokenInst.balanceOf(this);
+    tokenInst.transfer(owner, balance);
+  }
+
+}
\ No newline at end of file
diff --git a/contracts/ownership/HasNoTokens.sol b/contracts/ownership/HasNoTokens.sol
index 85c209e76..3f91ad24a 100644
--- a/contracts/ownership/HasNoTokens.sol
+++ b/contracts/ownership/HasNoTokens.sol
@@ -1,7 +1,6 @@
 pragma solidity ^0.4.11;
 
-import "./Ownable.sol";
-import "../token/ERC20Basic.sol";
+import "./CanReclaimToken.sol";
 
 /**
  * @title Contracts that should not own Tokens
@@ -10,7 +9,7 @@ import "../token/ERC20Basic.sol";
  * Should tokens (any ERC20Basic compatible) end up in the contract, it allows the
  * owner to reclaim the tokens.
  */
-contract HasNoTokens is Ownable {
+contract HasNoTokens is CanReclaimToken {
 
  /**
   * @dev Reject all ERC23 compatible tokens
@@ -22,13 +21,4 @@ contract HasNoTokens is Ownable {
     revert();
   }
 
-  /**
-   * @dev Reclaim all ERC20Basic compatible tokens
-   * @param tokenAddr address The address of the token contract
-   */
-  function reclaimToken(address tokenAddr) external onlyOwner {
-    ERC20Basic tokenInst = ERC20Basic(tokenAddr);
-    uint256 balance = tokenInst.balanceOf(this);
-    tokenInst.transfer(owner, balance);
-  }
 }
diff --git a/test/CanReclaimToken.js b/test/CanReclaimToken.js
new file mode 100644
index 000000000..01378630e
--- /dev/null
+++ b/test/CanReclaimToken.js
@@ -0,0 +1,35 @@
+'use strict';
+import expectThrow from './helpers/expectThrow';
+import toPromise from './helpers/toPromise';
+const CanReclaimToken = artifacts.require('../contracts/ownership/CanReclaimToken.sol');
+const BasicTokenMock = artifacts.require("./helpers/BasicTokenMock.sol");
+
+contract('CanReclaimToken', function(accounts) {
+  let token = null;
+  let canReclaimToken = null;
+
+  beforeEach(async () => {
+    // Create contract and token
+    token = await BasicTokenMock.new(accounts[0], 100);
+    canReclaimToken = await CanReclaimToken.new();
+    // Force token into contract
+    await token.transfer(canReclaimToken.address, 10);
+    const startBalance = await token.balanceOf(canReclaimToken.address);
+    assert.equal(startBalance, 10);
+  });
+
+  it('should allow owner to reclaim tokens', async function() {
+    const ownerStartBalance = await token.balanceOf(accounts[0]);
+    await canReclaimToken.reclaimToken(token.address);
+    const ownerFinalBalance = await token.balanceOf(accounts[0]);
+    const finalBalance = await token.balanceOf(canReclaimToken.address);
+    assert.equal(finalBalance, 0);
+    assert.equal(ownerFinalBalance - ownerStartBalance, 10);
+  });
+
+  it('should allow only owner to reclaim tokens', async function() {
+    await expectThrow(
+      canReclaimToken.reclaimToken(token.address, {from: accounts[1]}),
+    );
+  });
+});