feat(staking): add BTC support

omnichain/staking/contracts/Staking.sol

@@ -8,17 +8,25 @@ import "@zetachain/toolkit/contracts/BytesHelperLib.sol";
 
 contract Staking is ERC20, zContract {
     error SenderNotSystemContract();
-    error WrongChain();
-    error NotAuthorizedToClaim();
+    error WrongChain(uint256 chainID);
+    error UnknownAction(uint8 action);
+    error Overflow();
+    error Underflow();
+    error WrongAmount();
+    error NotAuthorized();
+    error NoRewardsToClaim();
 
     SystemContract public immutable systemContract;
     uint256 public immutable chainID;
+    uint256 constant BITCOIN = 18332;
 
-    mapping(address => uint256) public stakes;
-    mapping(address => address) public beneficiaries;
-    mapping(address => uint256) public lastStakeTime;
     uint256 public rewardRate = 1;
 
+    mapping(address => uint256) public stake;
+    mapping(address => bytes) public withdraw;
+    mapping(address => address) public beneficiary;
+    mapping(address => uint256) public lastStakeTime;
+
     constructor(
         string memory name_,
         string memory symbol_,
@@ -29,74 +37,126 @@ contract Staking is ERC20, zContract {
         chainID = chainID_;
     }
 
+    modifier onlySystem() {
+        require(
+            msg.sender == address(systemContract),
+            "Only system contract can call this function"
+        );
+        _;
+    }
+
+    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 onCrossChainCall(
         zContext calldata context,
         address zrc20,
         uint256 amount,
         bytes calldata message
-    ) external override {
-        if (msg.sender != address(systemContract)) {
-            revert SenderNotSystemContract();
+    ) external override onlySystem {
+        if (chainID != context.chainID) {
+            revert WrongChain(context.chainID);
         }
 
-        address acceptedZRC20 = systemContract.gasCoinZRC20ByChainId(chainID);
-        if (zrc20 != acceptedZRC20) revert WrongChain();
-
         address staker = BytesHelperLib.bytesToAddress(context.origin, 0);
-        address beneficiary = abi.decode(message, (address));
 
-        stakeZRC(staker, beneficiary, amount);
+        uint8 action = chainID == BITCOIN
+            ? uint8(message[0])
+            : abi.decode(message, (uint8));
+
+        if (action == 1) {
+            stakeZRC(staker, amount);
+        } else if (action == 2) {
+            unstakeZRC(staker);
+        } else if (action == 3) {
+            setBeneficiary(staker, message);
+        } else if (action == 4) {
+            setWithdraw(staker, message, context.origin);
+        } else {
+            revert UnknownAction(action);
+        }
     }
 
-    function stakeZRC(
-        address staker,
-        address beneficiary,
-        uint256 amount
-    ) internal {
-        stakes[staker] += amount;
-        if (beneficiaries[staker] == address(0)) {
-            beneficiaries[staker] = beneficiary;
-        }
-        lastStakeTime[staker] = block.timestamp;
+    function stakeZRC(address staker, uint256 amount) internal {
+        stake[staker] += amount;
+        if (stake[staker] < amount) revert Overflow();
 
+        lastStakeTime[staker] = block.timestamp;
         updateRewards(staker);
     }
 
     function updateRewards(address staker) internal {
-        uint256 timeDifference = block.timestamp - lastStakeTime[staker];
-        uint256 rewardAmount = timeDifference * stakes[staker] * rewardRate;
+        uint256 rewardAmount = queryRewards(staker);
 
-        _mint(beneficiaries[staker], rewardAmount);
+        _mint(beneficiary[staker], rewardAmount);
         lastStakeTime[staker] = block.timestamp;
     }
 
-    function claimRewards(address staker) external {
-        if (beneficiaries[staker] != msg.sender) {
-            revert NotAuthorizedToClaim();
-        }
+    function unstakeZRC(address staker) internal {
+        uint256 amount = stake[staker];
 
         updateRewards(staker);
-    }
 
-    function unstakeZRC(uint256 amount) external {
-        updateRewards(msg.sender);
+        address zrc20 = systemContract.gasCoinZRC20ByChainId(chainID);
+        (, uint256 gasFee) = IZRC20(zrc20).withdrawGasFee();
 
-        require(stakes[msg.sender] >= amount, "Insufficient staked balance");
+        if (amount < gasFee) revert WrongAmount();
 
-        address zrc20 = systemContract.gasCoinZRC20ByChainId(chainID);
+        bytes memory recipient = withdraw[staker];
 
-        (address gasZRC20, uint256 gasFee) = IZRC20(zrc20).withdrawGasFee();
+        stake[staker] = 0;
 
         IZRC20(zrc20).approve(zrc20, gasFee);
-        IZRC20(zrc20).withdraw(abi.encodePacked(msg.sender), amount - gasFee);
+        IZRC20(zrc20).withdraw(recipient, amount - gasFee);
+
+        if (stake[staker] > amount) revert Underflow();
+
+        lastStakeTime[staker] = block.timestamp;
+    }
+
+    function setBeneficiary(address staker, bytes calldata message) internal {
+        address beneficiaryAddress;
+        if (chainID == BITCOIN) {
+            beneficiaryAddress = BytesHelperLib.bytesToAddress(message, 1);
+        } else {
+            (, beneficiaryAddress) = abi.decode(message, (uint8, address));
+        }
+        beneficiary[staker] = beneficiaryAddress;
+    }
 
-        stakes[msg.sender] -= amount;
-        lastStakeTime[msg.sender] = block.timestamp;
+    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 queryRewards(address account) public view returns (uint256) {
-        uint256 timeDifference = block.timestamp - lastStakeTime[account];
-        uint256 rewardAmount = timeDifference * stakes[account] * rewardRate;
+    function queryRewards(address staker) public view returns (uint256) {
+        uint256 timeDifference = block.timestamp - lastStakeTime[staker];
+        uint256 rewardAmount = timeDifference * stake[staker] * rewardRate;
         return rewardAmount;
     }
+
+    function claimRewards(address staker) external {
+        if (beneficiary[staker] != msg.sender) revert NotAuthorized();
+        uint256 rewardAmount = queryRewards(staker);
+        if (rewardAmount <= 0) revert NoRewardsToClaim();
+        updateRewards(staker);
+    }
 }

omnichain/staking/hardhat.config.ts

@@ -1,8 +1,11 @@
-import "./tasks/interact";
+import "./tasks/stake";
 import "./tasks/deploy";
-import "./tasks/rewards";
 import "./tasks/claim";
 import "./tasks/unstake";
+import "./tasks/beneficiary";
+import "./tasks/withdraw";
+import "./tasks/unstake";
+import "./tasks/address";
 import "@nomicfoundation/hardhat-toolbox";
 import "@zetachain/toolkit/tasks";
 

omnichain/staking/lib/convertToHexAddress.ts

@@ -0,0 +1,13 @@
+import { ethers } from "ethers";
+
+export const convertToHexAddress = (address: string): string => {
+  let addr: string;
+  try {
+    // Check if it's a valid hex address
+    addr = ethers.utils.getAddress(address);
+  } catch (e) {
+    // If not, try to convert it to an address from bech32
+    addr = ("0x" + Buffer.from(address).toString("hex")).slice(0, 42);
+  }
+  return addr;
+};

omnichain/staking/package.json

@@ -26,7 +26,7 @@
     "@types/node": ">=12.0.0",
     "@typescript-eslint/eslint-plugin": "^5.59.9",
     "@typescript-eslint/parser": "^5.59.9",
-    "@zetachain/toolkit": "^2.1.2",
+    "@zetachain/toolkit": "^2.2.2",
     "axios": "^1.3.6",
     "chai": "^4.2.0",
     "dotenv": "^16.0.3",
@@ -48,4 +48,4 @@
     "typechain": "^8.1.0",
     "typescript": ">=4.5.0"
   }
-}
+}

omnichain/staking/tasks/address.ts

@@ -0,0 +1,18 @@
+import { task } from "hardhat/config";
+import { HardhatRuntimeEnvironment } from "hardhat/types";
+import { utils } from "ethers";
+
+const main = async (args: any, hre: HardhatRuntimeEnvironment) => {
+  const dataTypes = ["bytes"];
+  const values = [utils.toUtf8Bytes(args.address)];
+
+  const encodedData = utils.solidityPack(dataTypes, values);
+  console.log(`Encoded: ${encodedData}`);
+  console.log(`context.origin: ${encodedData.slice(0, 42)}`);
+};
+
+task(
+  "address",
+  "Encode a Bitcoin bech32 address to hex",
+  main
+).addPositionalParam("address");

omnichain/staking/tasks/beneficiary.ts

@@ -0,0 +1,33 @@
+import { task } from "hardhat/config";
+import { HardhatRuntimeEnvironment } from "hardhat/types";
+import { parseEther } from "@ethersproject/units";
+import { getAddress } from "@zetachain/protocol-contracts";
+import { prepareData, trackCCTX } from "@zetachain/toolkit/helpers";
+
+const main = async (args: any, hre: HardhatRuntimeEnvironment) => {
+  const [signer] = await hre.ethers.getSigners();
+  console.log(`🔑 Using account: ${signer.address}\n`);
+
+  const data = prepareData(
+    args.contract,
+    ["uint8", "address"],
+    ["3", args.beneficiary]
+  );
+  const to = getAddress("tss", hre.network.name);
+  const value = parseEther("0");
+
+  const tx = await signer.sendTransaction({ data, to, value });
+  console.log(`
+🚀 Successfully broadcasted a token transfer transaction on ${hre.network.name} network.
+📝 Transaction hash: ${tx.hash}
+`);
+  await trackCCTX(tx.hash);
+};
+
+task(
+  "set-beneficiary",
+  "Set the address on ZetaChain which will be allowed to claim staking rewards",
+  main
+)
+  .addParam("contract", "The address of the contract on ZetaChain")
+  .addPositionalParam("beneficiary", "The address of the beneficiary");

omnichain/staking/tasks/claim.ts

@@ -1,14 +1,17 @@
 import { task } from "hardhat/config";
 import { HardhatRuntimeEnvironment } from "hardhat/types";
+import { convertToHexAddress } from "../lib/convertToHexAddress";
 
 const main = async (args: any, hre: HardhatRuntimeEnvironment) => {
   const [signer] = await hre.ethers.getSigners();
   console.log(`🔑 Using account: ${signer.address}\n`);
 
+  const staker = convertToHexAddress(args.staker);
+
   const factory = await hre.ethers.getContractFactory("Staking");
   const contract = factory.attach(args.contract);
 
-  const tx = await contract.claimRewards(args.staker);
+  const tx = await contract.claimRewards(staker);
 
   const receipt = await tx.wait();
 

omnichain/staking/tasks/deploy.ts

@@ -22,20 +22,20 @@ const main = async (args: any, hre: HardhatRuntimeEnvironment) => {
 
   const factory = await hre.ethers.getContractFactory("Staking");
 
-  const chainID = hre.config.networks[args.chain]?.chainId;
-  if (chainID === undefined) {
-    throw new Error(`🚨 Chain ${args.chain} not found in hardhat config.`);
+  let symbol, chainID;
+  if (args.chain === "btc_testnet") {
+    symbol = "BTC";
+    chainID = 18332;
+  } else {
+    const zrc20 = getAddress("zrc20", args.chain);
+    const contract = new hre.ethers.Contract(zrc20, ZRC20.abi, signer);
+    symbol = await contract.symbol();
+    chainID = hre.config.networks[args.chain]?.chainId;
+    if (chainID === undefined) {
+      throw new Error(`🚨 Chain ${args.chain} not found in hardhat config.`);
+    }
   }
 
-  const ZRC20Address = getAddress("zrc20", args.chain);
-  const ZRC20Contract = new hre.ethers.Contract(
-    ZRC20Address,
-    ZRC20.abi,
-    signer
-  );
-
-  const symbol = await ZRC20Contract.symbol();
-
   const contract = await factory.deploy(
     `Staking rewards for ${symbol}`,
     `R${symbol.toUpperCase()}`,
@@ -50,4 +50,7 @@ const main = async (args: any, hre: HardhatRuntimeEnvironment) => {
 `);
 };
 
-task("deploy", "Deploy the contract", main).addParam("chain", "Chain name");
+task("deploy", "Deploy the contract", main).addParam(
+  "chain",
+  "Chain ID (use btc_testnet for Bitcoin Testnet)"
+);

omnichain/staking/tasks/interact.ts → omnichain/staking/tasks/stake.ts

@@ -8,20 +8,18 @@ const main = async (args: any, hre: HardhatRuntimeEnvironment) => {
   const [signer] = await hre.ethers.getSigners();
   console.log(`🔑 Using account: ${signer.address}\n`);
 
-  const data = prepareData(args.contract, ["address"], [args.beneficiary]);
+  const data = prepareData(args.contract, ["uint8"], ["1"]);
   const to = getAddress("tss", hre.network.name);
   const value = parseEther(args.amount);
 
   const tx = await signer.sendTransaction({ data, to, value });
-
   console.log(`
 🚀 Successfully broadcasted a token transfer transaction on ${hre.network.name} network.
 📝 Transaction hash: ${tx.hash}
 `);
   await trackCCTX(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")
-  .addParam("beneficiary");
+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");