Fix the BalanceOfPartition audit item (#679)

* fix balance of partition function * fix balance of Partition function * return balance in the parent implementation of the getTokensByPartition()
PolymathNetwork · Jun 13, 2019 · 2706d20 · 2706d20
1 parent 95338da
commit 2706d20
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 5 deletions.
diff --git a/contracts/modules/TransferManager/TransferManager.sol b/contracts/modules/TransferManager/TransferManager.sol
@@ -20,9 +20,14 @@ contract TransferManager is ITransferManager, Module {
 
     /**
      * @notice return the amount of tokens for a given user as per the partition
+     * @dev returning the balance of token holder against the UNLOCKED partition. 
+     * This condition is valid only when the base contract doesn't implement the
+     * `getTokensByPartition()` function.  
      */
-    function getTokensByPartition(bytes32 /*_partition*/, address /*_tokenHolder*/, uint256 /*_additionalBalance*/) external view returns(uint256) {
-        return 0;
+    function getTokensByPartition(bytes32 _partition, address _tokenHolder, uint256 /*_additionalBalance*/) external view returns(uint256) {
+        if (_partition == UNLOCKED)
+            return ISecurityToken(securityToken).balanceOf(_tokenHolder);
+        return uint256(0);
     }
 
 }
diff --git a/contracts/tokens/SecurityToken.sol b/contracts/tokens/SecurityToken.sol
@@ -486,13 +486,20 @@ contract SecurityToken is ERC20, ReentrancyGuard, SecurityTokenStorage, IERC1594
         return _balanceOfByPartition(_partition, _tokenHolder, 0);
     }
 
-    function _balanceOfByPartition(bytes32 _partition, address _tokenHolder, uint256 _additionalBalance) internal view returns(uint256 max) {
+    function _balanceOfByPartition(bytes32 _partition, address _tokenHolder, uint256 _additionalBalance) internal view returns(uint256 partitionBalance) {
         address[] memory tms = modules[TRANSFER_KEY];
         uint256 amount;
         for (uint256 i = 0; i < tms.length; i++) {
             amount = ITransferManager(tms[i]).getTokensByPartition(_partition, _tokenHolder, _additionalBalance);
-            if (max < amount) {
-                max = amount;
+            // In UNLOCKED partition we are returning the minimum of all the unlocked balances
+            if (_partition == UNLOCKED) {
+                if (amount < partitionBalance || i == 0)
+                    partitionBalance = amount;
+            }
+            // In locked partition we are returning the maximum of all the Locked balances
+            else {
+                if (partitionBalance < amount)
+                    partitionBalance = amount;
             }
         }
     }