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()

Author: satyamakgec

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);
     }
 
 }

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;
             }
         }
     }