Skip to content

Commit

Permalink
refactor: use _handleResponse to handle all responses to reduce cod…
Browse files Browse the repository at this point in the history 
…e size (#54)

* refactor: use _handleResponse to handle all responses to reduce code size

* fix(client): remove unused vars

* fix typo: principle => principal

* refactor: emit DepositResult earlier

* refactor: remove unused function

* doc: clarify comment

* refactor: devide response handling branch by response type

* fix: special calse for deposit-delegate and add add test

---------

Co-authored-by: MaxMustermann2 <[email protected]>
  • Loading branch information
@adu-web3 @MaxMustermann2
adu-web3 and MaxMustermann2 authored Jul 22, 2024
1 parent 8b5ee15 commit 59a2350
Show file tree
Hide file tree
Showing 16 changed files with 1,201 additions and 254 deletions.
807 changes: 806 additions & 1 deletion package-lock.json
View file

Large diffs are not rendered by default.

3 changes: 2 additions & 1 deletion package.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -379,7 +379,8 @@
"@nomicfoundation/hardhat-foundry": "^1.1.1",
"@nomicfoundation/hardhat-toolbox": "^4.0.0",
"hardhat": "^2.19.3",
"mocha": "^10.2.0"
"mocha": "^10.2.0",
"solhint": "^5.0.1"
},
"scripts": {
"test": "mocha"
Expand Down
14 changes: 0 additions & 14 deletions src/core/ClientChainGateway.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -63,16 +63,6 @@ contract ClientChainGateway is

require(owner_ != address(0), "ClientChainGateway: contract owner should not be empty");

_registeredResponseHooks[Action.REQUEST_DEPOSIT] = this.afterReceiveDepositResponse.selector;
_registeredResponseHooks[Action.REQUEST_WITHDRAW_PRINCIPAL_FROM_EXOCORE] =
this.afterReceiveWithdrawPrincipalResponse.selector;
_registeredResponseHooks[Action.REQUEST_DELEGATE_TO] = this.afterReceiveDelegateResponse.selector;
_registeredResponseHooks[Action.REQUEST_UNDELEGATE_FROM] = this.afterReceiveUndelegateResponse.selector;
_registeredResponseHooks[Action.REQUEST_WITHDRAW_REWARD_FROM_EXOCORE] =
this.afterReceiveWithdrawRewardResponse.selector;
_registeredResponseHooks[Action.REQUEST_DEPOSIT_THEN_DELEGATE_TO] =
this.afterReceiveDepositThenDelegateToResponse.selector;

_whiteListFunctionSelectors[Action.REQUEST_ADD_WHITELIST_TOKENS] =
this.afterReceiveAddWhitelistTokensRequest.selector;

Expand Down Expand Up @@ -108,10 +98,6 @@ contract ClientChainGateway is
_unpause();
}

function addWhitelistTokens(address[] calldata) external onlyOwner whenNotPaused {
revert("this function is not supported for client chain, please register on Exocore");
}

// implementation of ITokenWhitelister
function getWhitelistedTokensCount() external view returns (uint256) {
return whitelistTokens.length;
Expand Down
226 changes: 121 additions & 105 deletions src/core/ClientGatewayLzReceiver.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,27 +36,7 @@ abstract contract ClientGatewayLzReceiver is PausableUpgradeable, OAppReceiverUp

Action act = Action(uint8(payload[0]));
if (act == Action.RESPOND) {
uint64 requestId = uint64(bytes8(payload[1:9]));

Action requestAct = _registeredRequestActions[requestId];
bytes4 hookSelector = _registeredResponseHooks[requestAct];
if (hookSelector == bytes4(0)) {
revert UnsupportedResponse(act);
}

bytes memory requestPayload = _registeredRequests[requestId];
if (requestPayload.length == 0) {
revert UnexpectedResponse(requestId);
}

(bool success, bytes memory reason) =
address(this).call(abi.encodePacked(hookSelector, abi.encode(requestPayload, payload[9:])));
if (!success) {
revert RequestOrResponseExecuteFailed(act, _origin.nonce, reason);
}

delete _registeredRequestActions[requestId];
delete _registeredRequests[requestId];
_handleResponse(payload);
} else {
bytes4 selector_ = _whiteListFunctionSelectors[act];
if (selector_ == bytes4(0)) {
Expand All @@ -81,121 +61,157 @@ abstract contract ClientGatewayLzReceiver is PausableUpgradeable, OAppReceiverUp
return inboundNonce[srcEid][sender] + 1;
}

function afterReceiveDepositResponse(bytes memory requestPayload, bytes calldata responsePayload)
public
onlyCalledFromThis
{
(address token, address depositor, uint256 amount) = abi.decode(requestPayload, (address, address, uint256));

bool success = (uint8(bytes1(responsePayload[0])) == 1);
uint256 lastlyUpdatedPrincipalBalance = uint256(bytes32(responsePayload[1:]));

if (!success) {
revert DepositShouldNotFailOnExocore(token, depositor);
}

if (token == VIRTUAL_STAKED_ETH_ADDRESS) {
IExoCapsule capsule = _getCapsule(depositor);
capsule.updatePrincipalBalance(lastlyUpdatedPrincipalBalance);
// Though this function makes external calls to contract Vault or ExoCapsule, we just update their state variables
// and don't make
// calls to other contracts that do not belong to Exocore.
// And (success, updatedBalance) would be updated according to response message.
// slither-disable-next-line reentrancy-no-eth
function _handleResponse(bytes calldata response) internal {
(uint64 requestId, Action requestAct, bytes memory cachedRequest) = _getCachedRequestForResponse(response);

bool success = false;
uint256 updatedBalance;

if (_expectBasicResponse(requestAct)) {
success = _decodeBasicResponse(response);
} else if (_expectBalanceResponse(requestAct)) {
(address token, address staker,, uint256 amount) = _decodeCachedRequest(requestAct, cachedRequest);
(success, updatedBalance) = _decodeBalanceResponse(response);

if (_isPrincipalType(requestAct)) {
// we assume deposit request must always be successful, thus we should always update balance for deposit
// request
// Notice: Action.REQUEST_DEPOSIT_THEN_DELEGATE_TO is a special operation that is not atomic, since
// deposit should always be successful while delegate could fail for some cases
if (success || _isDeposit(requestAct)) {
_updatePrincipalAssetState(requestAct, token, staker, amount, updatedBalance);
}
} else {
// otherwise this is an operation aimed at reward since Action.REQUEST_WITHDRAW_REWARD_FROM_EXOCORE is
// the only asset operation request that deals with reward instead of principal
if (success) {
IVault vault = _getVault(token);

vault.updateRewardBalance(staker, updatedBalance);
if (_isWithdrawal(requestAct)) {
vault.updateWithdrawableBalance(staker, 0, amount);
}
}
}
} else {
IVault vault = _getVault(token);
vault.updatePrincipalBalance(depositor, lastlyUpdatedPrincipalBalance);
revert UnsupportedResponse(requestAct);
}

emit DepositResult(success, token, depositor, amount);
delete _registeredRequestActions[requestId];
delete _registeredRequests[requestId];

emit RequestFinished(requestAct, requestId, success);
}

function afterReceiveWithdrawPrincipalResponse(bytes memory requestPayload, bytes calldata responsePayload)
public
onlyCalledFromThis
{
(address token, address withdrawer, uint256 unlockPrincipalAmount) =
abi.decode(requestPayload, (address, address, uint256));
function _getCachedRequestForResponse(bytes calldata response) internal returns (uint64, Action, bytes memory) {
uint64 requestId = uint64(bytes8(response[1:9]));

bool success = (uint8(bytes1(responsePayload[0])) == 1);
uint256 lastlyUpdatedPrincipalBalance = uint256(bytes32(responsePayload[1:33]));
bytes memory cachedRequest = _registeredRequests[requestId];
if (cachedRequest.length == 0) {
revert UnexpectedResponse(requestId);
}
Action requestAct = _registeredRequestActions[requestId];

if (!success) {
emit WithdrawFailedOnExocore(token, withdrawer);
} else {
if (token == VIRTUAL_STAKED_ETH_ADDRESS) {
IExoCapsule capsule = _getCapsule(withdrawer);
return (requestId, requestAct, cachedRequest);
}

capsule.updatePrincipalBalance(lastlyUpdatedPrincipalBalance);
capsule.updateWithdrawableBalance(unlockPrincipalAmount);
} else {
IVault vault = _getVault(token);
function _isAssetOperationRequest(Action action) internal pure returns (bool) {
return action == Action.REQUEST_DEPOSIT || action == Action.REQUEST_WITHDRAW_PRINCIPAL_FROM_EXOCORE
|| action == Action.REQUEST_WITHDRAW_REWARD_FROM_EXOCORE;
}

vault.updatePrincipalBalance(withdrawer, lastlyUpdatedPrincipalBalance);
vault.updateWithdrawableBalance(withdrawer, unlockPrincipalAmount, 0);
}
function _isStakingOperationRequest(Action action) internal pure returns (bool) {
return action == Action.REQUEST_DELEGATE_TO || action == Action.REQUEST_UNDELEGATE_FROM
|| action == Action.REQUEST_DEPOSIT_THEN_DELEGATE_TO;
}

emit WithdrawPrincipalResult(success, token, withdrawer, unlockPrincipalAmount);
}
// Basic response only includes reqeust execution status, no other informations like balance update
// and it is typically the response of a staking only operations.
function _expectBasicResponse(Action action) internal pure returns (bool) {
return action == Action.REQUEST_DELEGATE_TO || action == Action.REQUEST_UNDELEGATE_FROM;
}

function afterReceiveWithdrawRewardResponse(bytes memory requestPayload, bytes calldata responsePayload)
public
onlyCalledFromThis
{
(address token, address withdrawer, uint256 unlockRewardAmount) =
abi.decode(requestPayload, (address, address, uint256));
// Balance response includes not only request execution status, but also the balance update informations,
// so it is typically the response of an asset operation.
function _expectBalanceResponse(Action action) internal pure returns (bool) {
return action == Action.REQUEST_DEPOSIT || action == Action.REQUEST_WITHDRAW_PRINCIPAL_FROM_EXOCORE
|| action == Action.REQUEST_WITHDRAW_REWARD_FROM_EXOCORE || action == Action.REQUEST_DEPOSIT_THEN_DELEGATE_TO;
}

bool success = (uint8(bytes1(responsePayload[0])) == 1);
uint256 lastlyUpdatedRewardBalance = uint256(bytes32(responsePayload[1:33]));
if (success) {
IVault vault = _getVault(token);
function _isPrincipalType(Action action) internal pure returns (bool) {
return action == Action.REQUEST_DEPOSIT || action == Action.REQUEST_WITHDRAW_PRINCIPAL_FROM_EXOCORE
|| action == Action.REQUEST_DEPOSIT_THEN_DELEGATE_TO;
}

vault.updateRewardBalance(withdrawer, lastlyUpdatedRewardBalance);
vault.updateWithdrawableBalance(withdrawer, 0, unlockRewardAmount);
}
function _isWithdrawal(Action action) internal pure returns (bool) {
return action == Action.REQUEST_WITHDRAW_PRINCIPAL_FROM_EXOCORE
|| action == Action.REQUEST_WITHDRAW_REWARD_FROM_EXOCORE;
}

emit WithdrawRewardResult(success, token, withdrawer, unlockRewardAmount);
function _isDeposit(Action action) internal pure returns (bool) {
return action == Action.REQUEST_DEPOSIT || action == Action.REQUEST_DEPOSIT_THEN_DELEGATE_TO;
}

function afterReceiveDelegateResponse(bytes memory requestPayload, bytes calldata responsePayload)
public
onlyCalledFromThis
function _decodeCachedRequest(Action requestAct, bytes memory cachedRequest)
internal
pure
returns (address token, address staker, string memory operator, uint256 amount)
{
(address token, address delegator, string memory operator, uint256 amount) =
abi.decode(requestPayload, (address, address, string, uint256));

bool success = (uint8(bytes1(responsePayload[0])) == 1);
if (_isAssetOperationRequest(requestAct)) {
(token, staker, amount) = abi.decode(cachedRequest, (address, address, uint256));
} else if (_isStakingOperationRequest(requestAct)) {
(token, staker, operator, amount) = abi.decode(cachedRequest, (address, address, string, uint256));
} else {
revert UnsupportedRequest(requestAct);
}

emit DelegateResult(success, delegator, operator, token, amount);
return (token, staker, operator, amount);
}

function afterReceiveUndelegateResponse(bytes memory requestPayload, bytes calldata responsePayload)
public
onlyCalledFromThis
function _decodeBalanceResponse(bytes calldata response)
internal
pure
returns (bool success, uint256 updatedBalance)
{
(address token, address undelegator, string memory operator, uint256 amount) =
abi.decode(requestPayload, (address, address, string, uint256));

bool success = (uint8(bytes1(responsePayload[0])) == 1);
success = (uint8(bytes1(response[9])) == 1);
updatedBalance = uint256(bytes32(response[10:]));

emit UndelegateResult(success, undelegator, operator, token, amount);
return (success, updatedBalance);
}

function afterReceiveDepositThenDelegateToResponse(bytes memory requestPayload, bytes calldata responsePayload)
public
onlyCalledFromThis
{
(address token, address delegator, string memory operator, uint256 amount) =
abi.decode(requestPayload, (address, address, string, uint256));
function _decodeBasicResponse(bytes calldata response) internal pure returns (bool success) {
success = (uint8(bytes1(response[9])) == 1);

bool delegateSuccess = (uint8(bytes1(responsePayload[0])) == 1);
uint256 lastlyUpdatedPrincipalBalance = uint256(bytes32(responsePayload[1:]));
return success;
}

function _updatePrincipalAssetState(
Action requestAct,
address token,
address staker,
uint256 amount,
uint256 updatedBalance
) internal {
if (token == VIRTUAL_STAKED_ETH_ADDRESS) {
IExoCapsule capsule = _getCapsule(delegator);
capsule.updatePrincipalBalance(lastlyUpdatedPrincipalBalance);
IExoCapsule capsule = _getCapsule(staker);

capsule.updatePrincipalBalance(updatedBalance);
if (_isWithdrawal(requestAct)) {
capsule.updateWithdrawableBalance(amount);
}
} else {
IVault vault = _getVault(token);
vault.updatePrincipalBalance(delegator, lastlyUpdatedPrincipalBalance);
}

emit DepositThenDelegateResult(delegateSuccess, delegator, operator, token, amount);
vault.updatePrincipalBalance(staker, updatedBalance);
if (_isWithdrawal(requestAct)) {
vault.updateWithdrawableBalance(staker, amount, 0);
}
}
}

// Though `_deployVault` would make external call to newly created `Vault` contract and initialize it,
Expand Down
Loading

0 comments on commit 59a2350

Please sign in to comment.