From 1db5dc773467c8d12b50b8b0dcae422a606f2d4d Mon Sep 17 00:00:00 2001 From: Purva-Chaudhari Date: Mon, 23 Sep 2024 00:40:20 -0400 Subject: [PATCH 1/3] add check for eoa --- smart-wallets/src/token/AssetToken.sol | 23 +++++++++++++++++++++-- 1 file changed, 21 insertions(+), 2 deletions(-) diff --git a/smart-wallets/src/token/AssetToken.sol b/smart-wallets/src/token/AssetToken.sol index 3bf6c0b..888951b 100644 --- a/smart-wallets/src/token/AssetToken.sol +++ b/smart-wallets/src/token/AssetToken.sol @@ -279,13 +279,32 @@ contract AssetToken is YieldDistributionToken, IAssetToken { return _getAssetTokenStorage().totalValue / totalSupply(); } + function isContract(address addr) internal view returns (bool) { + uint32 size; + assembly { + size := extcodesize(addr) + } + return size > 0; + } + /** * @notice Get the available unlocked AssetToken balance of a user * @param user Address of the user to get the available balance of * @return balanceAvailable Available unlocked AssetToken balance of the user */ - function getBalanceAvailable(address user) public view returns (uint256 balanceAvailable) { - return balanceOf(user) - SmartWallet(payable(user)).getBalanceLocked(this); + function getBalanceAvailable(address user) public view returns (uint256) { + if (isContract(user)) { + // User is a contract, safe to call getBalanceLocked + try SmartWallet(payable(user)).getBalanceLocked(this) returns (uint256 lockedBalance) { + return balanceOf(user) - lockedBalance; + } catch { + // Handle the case where the call reverts + return balanceOf(user); // Fallback if contract fails for some reason + } + } else { + // User is an EOA, no locked balance + return balanceOf(user); + } } /// @notice Total yield distributed to all AssetTokens for all users From 0726dc0ae93901bd32f1b202ed27e3966b22f702 Mon Sep 17 00:00:00 2001 From: Purva-Chaudhari Date: Mon, 23 Sep 2024 01:30:14 -0400 Subject: [PATCH 2/3] added revert and low level call --- smart-wallets/src/token/AssetToken.sol | 38 +++++++++++++------------- 1 file changed, 19 insertions(+), 19 deletions(-) diff --git a/smart-wallets/src/token/AssetToken.sol b/smart-wallets/src/token/AssetToken.sol index 888951b..dceb5e5 100644 --- a/smart-wallets/src/token/AssetToken.sol +++ b/smart-wallets/src/token/AssetToken.sol @@ -4,6 +4,7 @@ pragma solidity ^0.8.25; import { ERC20 } from "@openzeppelin/contracts/token/ERC20/ERC20.sol"; import { SmartWallet } from "../SmartWallet.sol"; +import { WalletUtils } from "../WalletUtils.sol"; import { IAssetToken } from "../interfaces/IAssetToken.sol"; import { YieldDistributionToken } from "./YieldDistributionToken.sol"; @@ -13,7 +14,7 @@ import { YieldDistributionToken } from "./YieldDistributionToken.sol"; * @notice ERC20 token that represents a tokenized real world asset * and distributes yield proportionally to token holders */ -contract AssetToken is YieldDistributionToken, IAssetToken { +contract AssetToken is WalletUtils, YieldDistributionToken, IAssetToken { // Storage @@ -86,6 +87,12 @@ contract AssetToken is YieldDistributionToken, IAssetToken { */ error AddressNotWhitelisted(address user); + /** + * @notice Indicates a failure because the user's SmartWallet call failed + * @param user Address of the user whose SmartWallet call failed + */ + error SmartWalletCallFailed(address user); + // Constructor /** @@ -279,31 +286,24 @@ contract AssetToken is YieldDistributionToken, IAssetToken { return _getAssetTokenStorage().totalValue / totalSupply(); } - function isContract(address addr) internal view returns (bool) { - uint32 size; - assembly { - size := extcodesize(addr) - } - return size > 0; - } - /** * @notice Get the available unlocked AssetToken balance of a user + * @dev Perform a low-level call to check for locked balance in the user's SmartWallet. + * Decodes the returned locked balance and subtract it from the user's total balance. + * Reverts if the call fails. * @param user Address of the user to get the available balance of * @return balanceAvailable Available unlocked AssetToken balance of the user */ function getBalanceAvailable(address user) public view returns (uint256) { - if (isContract(user)) { - // User is a contract, safe to call getBalanceLocked - try SmartWallet(payable(user)).getBalanceLocked(this) returns (uint256 lockedBalance) { - return balanceOf(user) - lockedBalance; - } catch { - // Handle the case where the call reverts - return balanceOf(user); // Fallback if contract fails for some reason - } + // Perform a low-level call to check for locked balance in the user's SmartWallet. + (bool success, bytes memory data) = + user.call(abi.encodeWithSignature("getBalanceLocked(address)", address(this))); + + if (success) { + uint256 lockedBalance = abi.decode(data, (uint256)); + return balanceOf(user) - lockedBalance; } else { - // User is an EOA, no locked balance - return balanceOf(user); + revert SmartWalletCallFailed(user); } } From d8736b988097f5ee241f0841a6009edc70d70524 Mon Sep 17 00:00:00 2001 From: Purva-Chaudhari Date: Mon, 23 Sep 2024 01:42:00 -0400 Subject: [PATCH 3/3] switch to try catch --- smart-wallets/src/WalletUtils.sol | 15 +++++++++++++++ smart-wallets/src/token/AssetToken.sol | 20 +++++++++----------- 2 files changed, 24 insertions(+), 11 deletions(-) diff --git a/smart-wallets/src/WalletUtils.sol b/smart-wallets/src/WalletUtils.sol index a6fea03..ea280ea 100644 --- a/smart-wallets/src/WalletUtils.sol +++ b/smart-wallets/src/WalletUtils.sol @@ -23,4 +23,19 @@ contract WalletUtils { _; } + /** + * @notice Checks if an address is a contract. + * @dev This function uses the `extcodesize` opcode to check if the target address contains contract code. + * It returns false for externally owned accounts (EOA) and true for contracts. + * @param addr The address to check. + * @return bool Returns true if the address is a contract, and false if it's an externally owned account (EOA). + */ + function isContract(address addr) internal view returns (bool) { + uint32 size; + assembly { + size := extcodesize(addr) + } + return size > 0; + } + } diff --git a/smart-wallets/src/token/AssetToken.sol b/smart-wallets/src/token/AssetToken.sol index dceb5e5..8f27c39 100644 --- a/smart-wallets/src/token/AssetToken.sol +++ b/smart-wallets/src/token/AssetToken.sol @@ -288,22 +288,20 @@ contract AssetToken is WalletUtils, YieldDistributionToken, IAssetToken { /** * @notice Get the available unlocked AssetToken balance of a user - * @dev Perform a low-level call to check for locked balance in the user's SmartWallet. - * Decodes the returned locked balance and subtract it from the user's total balance. - * Reverts if the call fails. + * @dev Attempts to call `getBalanceLocked` on the user if the user is a contract (SmartWallet). + * Reverts if the SmartWallet call fails. If the user is an EOA, it returns the balance directly. * @param user Address of the user to get the available balance of * @return balanceAvailable Available unlocked AssetToken balance of the user */ function getBalanceAvailable(address user) public view returns (uint256) { - // Perform a low-level call to check for locked balance in the user's SmartWallet. - (bool success, bytes memory data) = - user.call(abi.encodeWithSignature("getBalanceLocked(address)", address(this))); - - if (success) { - uint256 lockedBalance = abi.decode(data, (uint256)); - return balanceOf(user) - lockedBalance; + if (isContract(user)) { + try SmartWallet(payable(user)).getBalanceLocked(this) returns (uint256 lockedBalance) { + return balanceOf(user) - lockedBalance; + } catch { + revert SmartWalletCallFailed(user); + } } else { - revert SmartWalletCallFailed(user); + return balanceOf(user); } }