From 57c592fb302443f7160d3aeb9c31e7351918292f Mon Sep 17 00:00:00 2001
From: Jorge Izquierdo <jorge@aragon.one>
Date: Wed, 18 Nov 2020 11:54:48 +0100
Subject: [PATCH] Token: restricted transferability

---
 .../govern-token/contracts/GovernToken.sol    | 37 +++++++++++++++++--
 .../contracts/GovernTokenFactory.sol          |  7 +++-
 2 files changed, 38 insertions(+), 6 deletions(-)

diff --git a/packages/govern-token/contracts/GovernToken.sol b/packages/govern-token/contracts/GovernToken.sol
index 3fbaacaa8..952f0abce 100644
--- a/packages/govern-token/contracts/GovernToken.sol
+++ b/packages/govern-token/contracts/GovernToken.sol
@@ -37,6 +37,9 @@ contract GovernToken is IERC20, Initializable {
     mapping (address => uint256) override public balanceOf;
     mapping (address => mapping (address => uint256)) override public allowance;
 
+    bool public transfersRestricted;
+    mapping (address => bool) internal isWhitelisted;
+
     // ERC-2612, ERC-3009 state
     mapping (address => uint256) public nonces;
     mapping (address => mapping (bytes32 => bool)) public authorizationState;
@@ -45,18 +48,26 @@ contract GovernToken is IERC20, Initializable {
     event Transfer(address indexed from, address indexed to, uint256 value);
     event AuthorizationUsed(address indexed authorizer, bytes32 indexed nonce);
     event ChangeMinter(address indexed minter);
+    event RestrictedTransfers(bool transfersRestricted);
+    event SetWhitelisted(address indexed entity, bool indexed whitelisted);
 
     modifier onlyMinter {
         require(msg.sender == minter, "token: not minter");
         _;
     }
 
-    constructor(address _initialMinter, string memory _name, string memory _symbol, uint8 _decimals) public {
-        initialize(_initialMinter, _name, _symbol, _decimals);
+    modifier canTransact(address entity, string memory inOrOut) {
+        require(!transfersRestricted || isWhitelisted[entity], string(abi.encodePacked("token: bouncer ", inOrOut)));
+        _;
+    }
+
+    constructor(address _initialMinter, bool _transfersRestricted, string memory _name, string memory _symbol, uint8 _decimals) public {
+        initialize(_initialMinter, _transfersRestricted,_name, _symbol, _decimals);
     }
 
-    function initialize(address _initialMinter, string memory _name, string memory _symbol, uint8 _decimals) public onlyInit("token") {
+    function initialize(address _initialMinter, bool _transfersRestricted, string memory _name, string memory _symbol, uint8 _decimals) public onlyInit("token") {
         _changeMinter(_initialMinter);
+        _setTransfersRestricted(_transfersRestricted);
         name = _name;
         symbol = _symbol;
         decimals = _decimals;
@@ -80,6 +91,11 @@ contract GovernToken is IERC20, Initializable {
         emit ChangeMinter(newMinter);
     }
 
+    function _setTransfersRestricted(bool _transfersRestricted) internal {
+        transfersRestricted = _transfersRestricted;
+        emit RestrictedTransfers(_transfersRestricted);
+    }
+
     function _mint(address to, uint256 value) internal {
         totalSupply = totalSupply.add(value);
         balanceOf[to] = balanceOf[to].add(value);
@@ -98,7 +114,11 @@ contract GovernToken is IERC20, Initializable {
         emit Approval(owner, spender, value);
     }
 
-    function _transfer(address from, address to, uint256 value) private {
+    function _transfer(address from, address to, uint256 value)
+        canTransact(from, "in")
+        canTransact(to,  "out")
+        private
+    {
         require(to != address(this) && to != address(0), "token: bad to");
 
         // Balance is implicitly checked with SafeMath's underflow protection
@@ -132,6 +152,15 @@ contract GovernToken is IERC20, Initializable {
         _changeMinter(newMinter);
     }
 
+    function setTransfersRestricted(bool _transfersRestricted) external onlyMinter {
+        _setTransfersRestricted(_transfersRestricted);
+    }
+
+    function setWhitelisted(address _entity, bool _whitelisted) external onlyMinter {
+        isWhitelisted[_entity] = _whitelisted;
+        emit SetWhitelisted(_entity, _whitelisted);
+    }
+
     function burn(uint256 value) external returns (bool) {
         _burn(msg.sender, value);
         return true;
diff --git a/packages/govern-token/contracts/GovernTokenFactory.sol b/packages/govern-token/contracts/GovernTokenFactory.sol
index 1039681a6..8bb8f39e7 100644
--- a/packages/govern-token/contracts/GovernTokenFactory.sol
+++ b/packages/govern-token/contracts/GovernTokenFactory.sol
@@ -37,11 +37,12 @@ contract GovernTokenFactory {
         GovernMinter minter
     ) {
         if (!_useProxies) {
-            (token, minter) = _deployContracts(_initialMinter, _tokenName, _tokenSymbol, _tokenDecimals);
+            (token, minter) = _deployContracts(_initialMinter, false, _tokenName, _tokenSymbol, _tokenDecimals);
         } else {
             token = GovernToken(tokenBase.clone(abi.encodeWithSelector(
                 token.initialize.selector,
                 address(this),
+                false,
                 _tokenName,
                 _tokenSymbol,
                 _tokenDecimals
@@ -77,6 +78,7 @@ contract GovernTokenFactory {
         
         (GovernToken token, GovernMinter minter) = _deployContracts(
             address(this),
+            false,
             "GovernToken base",
             "GTB",
             0
@@ -94,6 +96,7 @@ contract GovernTokenFactory {
 
     function _deployContracts(
         address _initialMinter,
+        bool _transfersRestricted,
         string memory _tokenName,
         string memory _tokenSymbol,
         uint8 _tokenDecimals
@@ -101,7 +104,7 @@ contract GovernTokenFactory {
         GovernToken token,
         GovernMinter minter
     ) {
-        token = new GovernToken(address(this), _tokenName, _tokenSymbol, _tokenDecimals);
+        token = new GovernToken(address(this), _transfersRestricted, _tokenName, _tokenSymbol, _tokenDecimals);
         minter = new GovernMinter(GovernToken(token), address(_initialMinter), MerkleDistributor(distributorBase));
     }
 }