diff --git a/omnichain/staking/contracts/Staking.sol b/omnichain/staking/contracts/Staking.sol index ff8b996f..00b5304f 100644 --- a/omnichain/staking/contracts/Staking.sol +++ b/omnichain/staking/contracts/Staking.sol @@ -13,7 +13,6 @@ contract Staking is ERC20, zContract { uint256 public rewardRate = 1; - error SenderNotSystemContract(); error WrongChain(uint256 chainID); error UnknownAction(uint8 action); error Overflow(); @@ -22,10 +21,9 @@ contract Staking is ERC20, zContract { error NotAuthorized(); error NoRewardsToClaim(); - mapping(address => uint256) public stake; - mapping(address => bytes) public withdraw; - mapping(address => address) public beneficiary; - mapping(address => uint256) public lastStakeTime; + mapping(bytes => uint256) public stakes; + mapping(bytes => address) public beneficiary; + mapping(bytes => uint256) public lastStakeTime; constructor( string memory name_, @@ -55,48 +53,50 @@ contract Staking is ERC20, zContract { revert WrongChain(context.chainID); } - address staker = BytesHelperLib.bytesToAddress(context.origin, 0); - uint8 action = chainID == BITCOIN ? uint8(message[0]) : abi.decode(message, (uint8)); if (action == 1) { - stakeZRC(staker, amount); + stake(context.origin, amount, message); } else if (action == 2) { - unstakeZRC(staker); + unstake(context.origin); } else if (action == 3) { - setBeneficiary(staker, message); - } else if (action == 4) { - setWithdraw(staker, message, context.origin); + updateBeneficiary(context.origin, message); } else { revert UnknownAction(action); } } - function stakeZRC(address staker, uint256 amount) internal { - stake[staker] += amount; - if (stake[staker] < amount) revert Overflow(); + function stake( + bytes memory staker, + uint256 amount, + bytes calldata message + ) internal { + updateBeneficiary(staker, message); + + stakes[staker] += amount; + if (stakes[staker] < amount) revert Overflow(); lastStakeTime[staker] = block.timestamp; updateRewards(staker); } - function updateRewards(address staker) internal { + function updateRewards(bytes memory staker) internal { uint256 rewardAmount = queryRewards(staker); _mint(beneficiary[staker], rewardAmount); lastStakeTime[staker] = block.timestamp; } - function queryRewards(address staker) public view returns (uint256) { + function queryRewards(bytes memory staker) public view returns (uint256) { uint256 timeDifference = block.timestamp - lastStakeTime[staker]; - uint256 rewardAmount = timeDifference * stake[staker] * rewardRate; + uint256 rewardAmount = timeDifference * stakes[staker] * rewardRate; return rewardAmount; } - function unstakeZRC(address staker) internal { - uint256 amount = stake[staker]; + function unstake(bytes memory staker) internal { + uint256 amount = stakes[staker]; updateRewards(staker); @@ -105,19 +105,20 @@ contract Staking is ERC20, zContract { if (amount < gasFee) revert WrongAmount(); - bytes memory recipient = withdraw[staker]; - - stake[staker] = 0; + stakes[staker] = 0; IZRC20(zrc20).approve(zrc20, gasFee); - IZRC20(zrc20).withdraw(recipient, amount - gasFee); + IZRC20(zrc20).withdraw(staker, amount - gasFee); - if (stake[staker] > amount) revert Underflow(); + if (stakes[staker] > amount) revert Underflow(); lastStakeTime[staker] = block.timestamp; } - function setBeneficiary(address staker, bytes calldata message) internal { + function updateBeneficiary( + bytes memory staker, + bytes calldata message + ) internal { address beneficiaryAddress; if (chainID == BITCOIN) { beneficiaryAddress = BytesHelperLib.bytesToAddress(message, 1); @@ -127,33 +128,7 @@ contract Staking is ERC20, zContract { beneficiary[staker] = beneficiaryAddress; } - function setWithdraw( - address staker, - bytes calldata message, - bytes memory origin - ) internal { - bytes memory withdrawAddress; - if (chainID == BITCOIN) { - withdrawAddress = bytesToBech32Bytes(message, 1); - } else { - withdrawAddress = origin; - } - withdraw[staker] = withdrawAddress; - } - - function bytesToBech32Bytes( - bytes calldata data, - uint256 offset - ) internal pure returns (bytes memory) { - bytes memory bech32Bytes = new bytes(42); - for (uint i = 0; i < 42; i++) { - bech32Bytes[i] = data[i + offset]; - } - - return bech32Bytes; - } - - function claimRewards(address staker) external { + function claimRewards(bytes memory staker) external { if (beneficiary[staker] != msg.sender) revert NotAuthorized(); uint256 rewardAmount = queryRewards(staker); if (rewardAmount <= 0) revert NoRewardsToClaim(); diff --git a/omnichain/staking/tasks/beneficiary.ts b/omnichain/staking/tasks/beneficiary.ts index a837792f..87214573 100644 --- a/omnichain/staking/tasks/beneficiary.ts +++ b/omnichain/staking/tasks/beneficiary.ts @@ -28,7 +28,7 @@ const main = async (args: any, hre: HardhatRuntimeEnvironment) => { }; task( - "set-beneficiary", + "update-beneficiary", "Set the address on ZetaChain which will be allowed to claim staking rewards", main ) diff --git a/omnichain/staking/tasks/interact.ts b/omnichain/staking/tasks/interact.ts deleted file mode 100644 index 23764bdd..00000000 --- a/omnichain/staking/tasks/interact.ts +++ /dev/null @@ -1,34 +0,0 @@ -import { task } from "hardhat/config"; -import { HardhatRuntimeEnvironment } from "hardhat/types"; -import { parseEther } from "@ethersproject/units"; -import { getAddress } from "@zetachain/protocol-contracts"; -import { prepareData } from "@zetachain/toolkit/helpers"; - -const main = async (args: any, hre: HardhatRuntimeEnvironment) => { - const [signer] = await hre.ethers.getSigners(); - - const data = prepareData( - args.contract, - [], - [] - ); - const to = getAddress("tss", hre.network.name); - const value = parseEther(args.amount); - - const tx = await signer.sendTransaction({ data, to, value }); - - if (args.json) { - console.log(JSON.stringify(tx, null, 2)); - } else { - console.log(`🔑 Using account: ${signer.address}\n`); - - console.log(`🚀 Successfully broadcasted a token transfer transaction on ${hre.network.name} network. -📝 Transaction hash: ${tx.hash} -`); - } -}; - -task("interact", "Interact with the contract", main) - .addParam("contract", "The address of the withdraw contract on ZetaChain") - .addParam("amount", "Amount of tokens to send") - .addFlag("json", "Output in JSON") diff --git a/omnichain/staking/tasks/stake.ts b/omnichain/staking/tasks/stake.ts index 3a0666d0..8f6bed64 100644 --- a/omnichain/staking/tasks/stake.ts +++ b/omnichain/staking/tasks/stake.ts @@ -7,8 +7,12 @@ import { prepareData, trackCCTX } from "@zetachain/toolkit/helpers"; const main = async (args: any, hre: HardhatRuntimeEnvironment) => { const [signer] = await hre.ethers.getSigners(); - const data = prepareData(args.contract, ["uint8"], ["1"]); - const to = getAddress("tss", hre.network.name); + const data = prepareData( + args.contract, + ["uint8", "address"], + ["1", args.beneficiary] + ); + const to = getAddress("tss", hre.network.name as any); const value = parseEther(args.amount); const tx = await signer.sendTransaction({ data, to, value }); @@ -26,4 +30,5 @@ const main = async (args: any, hre: HardhatRuntimeEnvironment) => { task("stake", "Deposit tokens to ZetaChain and stake them", main) .addParam("contract", "The address of the contract on ZetaChain") .addParam("amount", "Amount of tokens to send") + .addParam("beneficiary", "Beneficiary") .addFlag("json", "Output in JSON");