Skip to content

Commit

Permalink
feat: add asset cap functionality (#242)
Browse files Browse the repository at this point in the history 
* feat: add asset cap functionality

* fix: unit tests

* fix: arbi fork tests

* fix: arbi fork tests

* chore: return account address and emit Repay event

* test: add assert
  • Loading branch information
@r0ohafza
r0ohafza authored Jan 23, 2023
1 parent 6c21849 commit 06c35d2
Show file tree
Hide file tree
Showing 39 changed files with 365 additions and 149 deletions.
2 changes: 1 addition & 1 deletion lib/forge-std
Submodule forge-std updated 43 files
+92 −0 .github/workflows/ci.yml
+0 −26 .github/workflows/tests.yml
+1 −1 .gitignore
+1 −1 LICENSE-APACHE
+1 −1 LICENSE-MIT
+8 −4 README.md
+21 −0 foundry.toml
+1 −1 lib/ds-test
+16 −0 package.json
+33 −0 src/Base.sol
+94 −0 src/InvariantTest.sol
+20 −33 src/Script.sol
+209 −0 src/StdAssertions.sol
+189 −0 src/StdChains.sol
+565 −0 src/StdCheats.sol
+15 −0 src/StdError.sol
+179 −0 src/StdJson.sol
+43 −0 src/StdMath.sol
+327 −0 src/StdStorage.sol
+123 −0 src/StdUtils.sol
+26 −777 src/Test.sol
+334 −154 src/Vm.sol
+14 −6 src/console2.sol
+105 −0 src/interfaces/IERC1155.sol
+12 −0 src/interfaces/IERC165.sol
+43 −0 src/interfaces/IERC20.sol
+190 −0 src/interfaces/IERC4626.sol
+164 −0 src/interfaces/IERC721.sol
+0 −12 src/test/Script.t.sol
+0 −226 src/test/StdCheats.t.sol
+0 −200 src/test/StdMath.t.sol
+65 −69 test/StdAssertions.t.sol
+123 −0 test/StdChains.t.sol
+305 −0 test/StdCheats.t.sol
+5 −11 test/StdError.t.sol
+197 −0 test/StdMath.t.sol
+46 −84 test/StdStorage.t.sol
+191 −0 test/StdUtils.t.sol
+10 −0 test/compilation/CompilationScript.sol
+10 −0 test/compilation/CompilationScriptBase.sol
+10 −0 test/compilation/CompilationTest.sol
+10 −0 test/compilation/CompilationTestBase.sol
+187 −0 test/fixtures/broadcast.log.json
160 changes: 91 additions & 69 deletions src/core/AccountManager.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,9 @@ contract AccountManager is ReentrancyGuard, Pausable, IAccountManager {
/// @notice Mapping of collateral enabled tokens
mapping(address => bool) public isCollateralAllowed;

/// @notice Number of assets a sentiment account can hold - 1
uint256 public assetCap;

/* -------------------------------------------------------------------------- */
/* CUSTOM MODIFIERS */
/* -------------------------------------------------------------------------- */
Expand Down Expand Up @@ -76,10 +79,11 @@ contract AccountManager is ReentrancyGuard, Pausable, IAccountManager {

/// @notice Initializes external dependencies
function initDep() external adminOnly {
riskEngine = IRiskEngine(registry.getAddress('RISK_ENGINE'));
controller = IControllerFacade(registry.getAddress('CONTROLLER'));
accountFactory =
IAccountFactory(registry.getAddress('ACCOUNT_FACTORY'));
riskEngine = IRiskEngine(registry.getAddress("RISK_ENGINE"));
controller = IControllerFacade(registry.getAddress("CONTROLLER"));
accountFactory = IAccountFactory(
registry.getAddress("ACCOUNT_FACTORY")
);
}

/**
Expand All @@ -89,10 +93,15 @@ contract AccountManager is ReentrancyGuard, Pausable, IAccountManager {
Emits AccountAssigned(account, owner) event
@param owner Owner of the newly opened account
*/
function openAccount(address owner) external nonReentrant whenNotPaused {
function openAccount(address owner)
external
nonReentrant
whenNotPaused
returns (address)
{
if (owner == address(0)) revert Errors.ZeroAddress();
address account;
uint length = inactiveAccountsOf[owner].length;
uint256 length = inactiveAccountsOf[owner].length;
if (length == 0) {
account = accountFactory.create(address(this));
IAccount(account).init(address(this));
Expand All @@ -104,6 +113,7 @@ contract AccountManager is ReentrancyGuard, Pausable, IAccountManager {
}
IAccount(account).activate();
emit AccountAssigned(account, owner);
return account;
}

/**
Expand All @@ -112,7 +122,11 @@ contract AccountManager is ReentrancyGuard, Pausable, IAccountManager {
Emits AccountClosed(account, owner) event
@param _account Address of account to be closed
*/
function closeAccount(address _account) public nonReentrant onlyOwner(_account) {
function closeAccount(address _account)
public
nonReentrant
onlyOwner(_account)
{
IAccount account = IAccount(_account);
if (account.activationBlock() == block.number)
revert Errors.AccountDeactivationFailure();
Expand Down Expand Up @@ -145,12 +159,12 @@ contract AccountManager is ReentrancyGuard, Pausable, IAccountManager {
@param account Address of account
@param amt Amount of Eth to withdraw
*/
function withdrawEth(address account, uint amt)
function withdrawEth(address account, uint256 amt)
external
nonReentrant
onlyOwner(account)
{
if(!riskEngine.isWithdrawAllowed(account, address(0), amt))
if (!riskEngine.isWithdrawAllowed(account, address(0), amt))
revert Errors.RiskThresholdBreached();
account.withdrawEth(msg.sender, amt);
}
Expand All @@ -163,16 +177,18 @@ contract AccountManager is ReentrancyGuard, Pausable, IAccountManager {
@param token Address of token
@param amt Amount of token to deposit
*/
function deposit(address account, address token, uint amt)
external
nonReentrant
whenNotPaused
onlyOwner(account)
{
function deposit(
address account,
address token,
uint256 amt
) external nonReentrant whenNotPaused onlyOwner(account) {
if (!isCollateralAllowed[token])
revert Errors.CollateralTypeRestricted();
if (IAccount(account).hasAsset(token) == false)
if (IAccount(account).hasAsset(token) == false) {
if (IAccount(account).getAssets().length > assetCap)
revert Errors.MaxAssetCap();
IAccount(account).addAsset(token);
}
token.safeTransferFrom(msg.sender, account, amt);
}

Expand All @@ -185,16 +201,15 @@ contract AccountManager is ReentrancyGuard, Pausable, IAccountManager {
@param token Address of token
@param amt Amount of token to withdraw
*/
function withdraw(address account, address token, uint amt)
external
nonReentrant
onlyOwner(account)
{
function withdraw(
address account,
address token,
uint256 amt
) external nonReentrant onlyOwner(account) {
if (!riskEngine.isWithdrawAllowed(account, token, amt))
revert Errors.RiskThresholdBreached();
account.withdraw(msg.sender, token, amt);
if (token.balanceOf(account) == 0)
IAccount(account).removeAsset(token);
if (token.balanceOf(account) == 0) IAccount(account).removeAsset(token);
}

/**
Expand All @@ -206,16 +221,18 @@ contract AccountManager is ReentrancyGuard, Pausable, IAccountManager {
@param token Address of token
@param amt Amount of token to borrow
*/
function borrow(address account, address token, uint amt)
external
nonReentrant
whenNotPaused
onlyOwner(account)
{
function borrow(
address account,
address token,
uint256 amt
) external nonReentrant whenNotPaused onlyOwner(account) {
if (registry.LTokenFor(token) == address(0))
revert Errors.LTokenUnavailable();
if (IAccount(account).hasAsset(token) == false)
if (IAccount(account).hasAsset(token) == false) {
if (IAccount(account).getAssets().length > assetCap)
revert Errors.MaxAssetCap();
IAccount(account).addAsset(token);
}
if (ILToken(registry.LTokenFor(token)).lendTo(account, amt))
IAccount(account).addBorrow(token);
if (!riskEngine.isAccountHealthy(account))
Expand All @@ -231,11 +248,11 @@ contract AccountManager is ReentrancyGuard, Pausable, IAccountManager {
@param token Address of token
@param amt Amount of token to borrow
*/
function repay(address account, address token, uint amt)
public
nonReentrant
onlyOwner(account)
{
function repay(
address account,
address token,
uint256 amt
) public nonReentrant onlyOwner(account) {
_repay(account, token, amt);
}

Expand Down Expand Up @@ -265,13 +282,9 @@ contract AccountManager is ReentrancyGuard, Pausable, IAccountManager {
address account,
address token,
address spender,
uint amt
)
external
nonReentrant
onlyOwner(account)
{
if(address(controller.controllerFor(spender)) == address(0))
uint256 amt
) external nonReentrant onlyOwner(account) {
if (address(controller.controllerFor(spender)) == address(0))
revert Errors.FunctionCallRestricted();
account.safeApprove(token, spender, amt);
}
Expand All @@ -289,24 +302,25 @@ contract AccountManager is ReentrancyGuard, Pausable, IAccountManager {
function exec(
address account,
address target,
uint amt,
uint256 amt,
bytes calldata data
)
external
nonReentrant
onlyOwner(account)
{
) external nonReentrant onlyOwner(account) {
bool isAllowed;
address[] memory tokensIn;
address[] memory tokensOut;
(isAllowed, tokensIn, tokensOut) =
controller.canCall(target, (amt > 0), data);
(isAllowed, tokensIn, tokensOut) = controller.canCall(
target,
(amt > 0),
data
);
if (!isAllowed) revert Errors.FunctionCallRestricted();
_updateTokensIn(account, tokensIn);
(bool success,) = IAccount(account).exec(target, amt, data);
(bool success, ) = IAccount(account).exec(target, amt, data);
if (!success)
revert Errors.AccountInteractionFailure(account, target, amt, data);
_updateTokensOut(account, tokensOut);
if (IAccount(account).getAssets().length > assetCap + 1)
revert Errors.MaxAssetCap();
if (!riskEngine.isAccountHealthy(account))
revert Errors.RiskThresholdBreached();
}
Expand All @@ -317,8 +331,8 @@ contract AccountManager is ReentrancyGuard, Pausable, IAccountManager {
*/
function settle(address account) external nonReentrant onlyOwner(account) {
address[] memory borrows = IAccount(account).getBorrows();
for (uint i; i < borrows.length; i++) {
_repay(account, borrows[i], type(uint).max);
for (uint256 i; i < borrows.length; i++) {
_repay(account, borrows[i], type(uint256).max);
}
}

Expand All @@ -327,9 +341,7 @@ contract AccountManager is ReentrancyGuard, Pausable, IAccountManager {
@param user Address of user
@return address[] List of inactive accounts
*/
function getInactiveAccountsOf(
address user
)
function getInactiveAccountsOf(address user)
external
view
returns (address[] memory)
Expand All @@ -344,8 +356,8 @@ contract AccountManager is ReentrancyGuard, Pausable, IAccountManager {
function _updateTokensIn(address account, address[] memory tokensIn)
internal
{
uint tokensInLen = tokensIn.length;
for(uint i; i < tokensInLen; ++i) {
uint256 tokensInLen = tokensIn.length;
for (uint256 i; i < tokensInLen; ++i) {
if (IAccount(account).hasAsset(tokensIn[i]) == false)
IAccount(account).addAsset(tokensIn[i]);
}
Expand All @@ -354,8 +366,8 @@ contract AccountManager is ReentrancyGuard, Pausable, IAccountManager {
function _updateTokensOut(address account, address[] memory tokensOut)
internal
{
uint tokensOutLen = tokensOut.length;
for(uint i; i < tokensOutLen; ++i) {
uint256 tokensOutLen = tokensOut.length;
for (uint256 i; i < tokensOutLen; ++i) {
if (tokensOut[i].balanceOf(account) == 0)
IAccount(account).removeAsset(tokensOut[i]);
}
Expand All @@ -364,29 +376,31 @@ contract AccountManager is ReentrancyGuard, Pausable, IAccountManager {
function _liquidate(address _account) internal {
IAccount account = IAccount(_account);
address[] memory accountBorrows = account.getBorrows();
uint borrowLen = accountBorrows.length;
uint256 borrowLen = accountBorrows.length;

ILToken LToken;
uint amt;
uint256 amt;

for(uint i; i < borrowLen; ++i) {
for (uint256 i; i < borrowLen; ++i) {
address token = accountBorrows[i];
LToken = ILToken(registry.LTokenFor(token));
LToken.updateState();
amt = LToken.getBorrowBalance(_account);
token.safeTransferFrom(msg.sender, address(LToken), amt);
LToken.collectFrom(_account, amt);
account.removeBorrow(token);
emit Repay(_account, msg.sender, token, amt);
}
account.sweepTo(msg.sender);
}

function _repay(address account, address token, uint amt)
internal
{
function _repay(
address account,
address token,
uint256 amt
) internal {
ILToken LToken = ILToken(registry.LTokenFor(token));
if (address(LToken) == address(0))
revert Errors.LTokenUnavailable();
if (address(LToken) == address(0)) revert Errors.LTokenUnavailable();
LToken.updateState();
if (amt == type(uint256).max) amt = LToken.getBorrowBalance(account);
account.withdraw(address(LToken), token, amt);
Expand All @@ -408,4 +422,12 @@ contract AccountManager is ReentrancyGuard, Pausable, IAccountManager {
function toggleCollateralStatus(address token) external adminOnly {
isCollateralAllowed[token] = !isCollateralAllowed[token];
}
}

/**
@notice Set asset cap
@param cap Number of assets
*/
function setAssetCap(uint256 cap) external adminOnly {
assetCap = cap - 1;
}
}
Loading

0 comments on commit 06c35d2

Please sign in to comment.