Skip to content

Commit

Permalink
Add more query options
Browse files Browse the repository at this point in the history
  • Loading branch information
lukema95 committed Sep 25, 2024
1 parent adb6496 commit 28d6765
Show file tree
Hide file tree
Showing 4 changed files with 174 additions and 118 deletions.
3 changes: 2 additions & 1 deletion package.json
Original file line number Diff line number Diff line change
Expand Up @@ -100,6 +100,7 @@
"bech32": "^2.0.0",
"bip39": "^3.1.0",
"bitcoinjs-lib": "^6.1.3",
"bs58": "^6.0.0",
"dotenv": "16.0.3",
"ecpair": "^2.1.0",
"envfile": "^6.18.0",
Expand All @@ -117,4 +118,4 @@
"ws": "^8.17.1"
},
"packageManager": "yarn@1.22.21+sha1.1959a18351b811cdeedbd484a8f86c3cc3bbaf72"
}
}
188 changes: 95 additions & 93 deletions packages/client/src/getBalances.ts
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,7 @@ export const getBalances = async function (
evmAddress,
btcAddress,
solanaAddress,
}: { btcAddress?: string; evmAddress: string; solanaAddress?: string }
}: { btcAddress?: string; evmAddress?: string; solanaAddress?: string }
): Promise<TokenBalance[]> {
let tokens = [];
const supportedChains = await this.getSupportedChains();
Expand Down Expand Up @@ -142,109 +142,111 @@ export const getBalances = async function (

const multicallContexts: Record<string, any[]> = {};

tokens.forEach((token: any) => {
if (token.coin_type === "ERC20" || token.coin_type === "ZRC20") {
if (!multicallContexts[token.chain_name]) {
multicallContexts[token.chain_name] = [];
}
multicallContexts[token.chain_name].push({
callData: new ethers.utils.Interface(
token.coin_type === "ERC20" ? ERC20_ABI.abi : ZRC20.abi
).encodeFunctionData("balanceOf", [evmAddress]),
target: token.contract,
});
}
});

const balances: TokenBalance[] = [];

await Promise.all(
Object.keys(multicallContexts).map(async (chainName) => {
const rpc = await this.getEndpoint("evm", chainName);
const provider = new ethers.providers.StaticJsonRpcProvider(rpc);
const multicallContract = new ethers.Contract(
multicallAddress,
MULTICALL3_ABI,
provider
);

const calls = multicallContexts[chainName];
if (evmAddress) {
tokens.forEach((token: any) => {
if (token.coin_type === "ERC20" || token.coin_type === "ZRC20") {
if (!multicallContexts[token.chain_name]) {
multicallContexts[token.chain_name] = [];
}
multicallContexts[token.chain_name].push({
callData: new ethers.utils.Interface(
token.coin_type === "ERC20" ? ERC20_ABI.abi : ZRC20.abi
).encodeFunctionData("balanceOf", [evmAddress]),
target: token.contract,
});
}
});

try {
const { returnData } = await multicallContract.callStatic.aggregate(
calls
await Promise.all(
Object.keys(multicallContexts).map(async (chainName) => {
const rpc = await this.getEndpoint("evm", chainName);
const provider = new ethers.providers.StaticJsonRpcProvider(rpc);
const multicallContract = new ethers.Contract(
multicallAddress,
MULTICALL3_ABI,
provider
);

returnData.forEach((data: any, index: number) => {
const token = tokens.find(
const calls = multicallContexts[chainName];

try {
const { returnData } = await multicallContract.callStatic.aggregate(
calls
);

returnData.forEach((data: any, index: number) => {
const token = tokens.find(
(t) =>
t.chain_name === chainName &&
(t.coin_type === "ERC20" || t.coin_type === "ZRC20") &&
t.contract === calls[index].target
);
if (token) {
const balance = ethers.utils.defaultAbiCoder.decode(
["uint256"],
data
)[0];
const formattedBalance = formatUnits(balance, token.decimals);
balances.push({ ...token, balance: formattedBalance });
}
});
} catch (error) {
console.error(`Multicall failed for ${chainName}:`, error);
// Fallback to individual calls if multicall fails
for (const token of tokens.filter(
(t) =>
t.chain_name === chainName &&
(t.coin_type === "ERC20" || t.coin_type === "ZRC20") &&
t.contract === calls[index].target
);
if (token) {
const balance = ethers.utils.defaultAbiCoder.decode(
["uint256"],
data
)[0];
const formattedBalance = formatUnits(balance, token.decimals);
balances.push({ ...token, balance: formattedBalance });
(t.coin_type === "ERC20" || t.coin_type === "ZRC20")
)) {
try {
const contract = new ethers.Contract(
token.contract,
token.coin_type === "ERC20" ? ERC20_ABI.abi : ZRC20.abi,
provider
);
const balance = await contract.balanceOf(evmAddress);
const formattedBalance = formatUnits(balance, token.decimals);
balances.push({ ...token, balance: formattedBalance });
} catch (err) {
console.error(
`Failed to get balance for ${token.symbol} on ${chainName}:`,
err
);
}
}
});
} catch (error) {
console.error(`Multicall failed for ${chainName}:`, error);
// Fallback to individual calls if multicall fails
for (const token of tokens.filter(
(t) =>
t.chain_name === chainName &&
(t.coin_type === "ERC20" || t.coin_type === "ZRC20")
)) {
try {
const contract = new ethers.Contract(
token.contract,
token.coin_type === "ERC20" ? ERC20_ABI.abi : ZRC20.abi,
provider
);
const balance = await contract.balanceOf(evmAddress);
}
})
);

await Promise.all(
tokens
.filter(
(token) =>
token.coin_type === "Gas" &&
![
"btc_testnet",
"btc_mainnet",
"solana_mainnet",
"solana_testnet",
"solana_devnet",
].includes(token.chain_name)
)
.map(async (token) => {
const chainLabel = Object.keys(this.getChains()).find(
(key) => this.getChains()[key].chain_id === parseInt(token.chain_id)
);
if (chainLabel) {
const rpc = await this.getEndpoint("evm", chainLabel);
const provider = new ethers.providers.StaticJsonRpcProvider(rpc);
const balance = await provider.getBalance(evmAddress);
const formattedBalance = formatUnits(balance, token.decimals);
balances.push({ ...token, balance: formattedBalance });
} catch (err) {
console.error(
`Failed to get balance for ${token.symbol} on ${chainName}:`,
err
);
}
}
}
})
);

await Promise.all(
tokens
.filter(
(token) =>
token.coin_type === "Gas" &&
![
"btc_testnet",
"btc_mainnet",
"solana_mainnet",
"solana_testnet",
"solana_devnet",
].includes(token.chain_name)
)
.map(async (token) => {
const chainLabel = Object.keys(this.getChains()).find(
(key) => this.getChains()[key].chain_id === parseInt(token.chain_id)
);
if (chainLabel) {
const rpc = await this.getEndpoint("evm", chainLabel);
const provider = new ethers.providers.StaticJsonRpcProvider(rpc);
const balance = await provider.getBalance(evmAddress);
const formattedBalance = formatUnits(balance, token.decimals);
balances.push({ ...token, balance: formattedBalance });
}
})
);
})
);
}

await Promise.all(
tokens
Expand Down
89 changes: 65 additions & 24 deletions packages/tasks/src/balances.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import { Keypair } from "@solana/web3.js";
import bs58 from "bs58";
import * as dotenv from "dotenv";
import { ethers } from "ethers";
import { task } from "hardhat/config";
Expand All @@ -15,10 +16,12 @@ export const walletError = `
To resolve this issue, please follow these steps:
* Set your PRIVATE_KEY environment variable. You can write
it to a .env file in the root of your project like this:
* Set your EVM_PRIVATE_KEY, SOLANA_PRIVATE_KEY, or BTC_PRIVATE_KEY environment variables.
You can write it to a .env file in the root of your project like this:
PRIVATE_KEY=123... (without the 0x prefix)
EVM_PRIVATE_KEY=123... (without the 0x prefix)
BTC_PRIVATE_KEY=123... (without the 0x prefix)
SOLANA_PRIVATE_KEY=123.. (base58 encoded or json array)
Or you can generate a new private key by running:
Expand All @@ -27,9 +30,9 @@ To resolve this issue, please follow these steps:

const balancesError = `
* Alternatively, you can fetch the balance of any address
by using the --address flag:
by using the --evm, --solana, or --bitcoin flag:
npx hardhat balances --address <wallet_address>
npx hardhat balances --evm <evm_address> --solana <solana_address> --bitcoin <bitcoin_address>
`;

const main = async (args: any, hre: HardhatRuntimeEnvironment) => {
Expand All @@ -40,24 +43,60 @@ const main = async (args: any, hre: HardhatRuntimeEnvironment) => {
if (!args.json) {
spinner.start();
}
const pk = process.env.PRIVATE_KEY;
let evmAddress: string;
let btcAddress: any;
let solanaAddress: any;

if (args.address) {
evmAddress = args.address;
} else if (pk) {
evmAddress = new ethers.Wallet(pk).address;
btcAddress = bitcoinAddress(pk, args.mainnet ? "mainnet" : "testnet");
solanaAddress = Keypair.fromSeed(
Buffer.from(pk, "hex")
).publicKey.toString();
} else {
spinner.stop();
console.error(walletError + balancesError);
return process.exit(1);

const evmKey = process.env.EVM_PRIVATE_KEY;
const solanaKey = process.env.SOLANA_PRIVATE_KEY;
const btcKey = process.env.BTC_PRIVATE_KEY;

let evmAddress: string | undefined;
let btcAddress: string | undefined;
let solanaAddress: string | undefined;

if (args.evm) {
evmAddress = args.evm;
}
if (args.solana) {
solanaAddress = args.solana;
}
if (args.bitcoin) {
btcAddress = args.bitcoin;
}

if (!evmAddress && !solanaAddress && !btcAddress) {
if (evmKey) {
evmAddress = new ethers.Wallet(evmKey).address;
}
if (solanaKey) {
try {
if (solanaKey.startsWith("[") && solanaKey.endsWith("]")) {
solanaAddress = Keypair.fromSecretKey(
Uint8Array.from(JSON.parse(solanaKey))
).publicKey.toString();
} else {
solanaAddress = Keypair.fromSecretKey(
bs58.decode(solanaKey)
).publicKey.toString();
}
} catch (error) {
{
spinner.stop();
console.error("Error parsing solanaKey", error);
return process.exit(1);
}
}
}
if (btcKey) {
btcAddress = bitcoinAddress(btcKey, args.mainnet ? "mainnet" : "testnet");
}
if (!solanaKey && !btcKey && !evmKey) {
{
spinner.stop();
console.error(walletError + balancesError);
return process.exit(1);
}
}
}

let balances = (await client.getBalances({
btcAddress,
evmAddress,
Expand Down Expand Up @@ -98,6 +137,8 @@ export const balancesTask = task(
`Fetch native and ZETA token balances`,
main
)
.addOptionalParam("address", `Fetch balances for a specific address`)
.addFlag("json", "Output balances as JSON")
.addFlag("mainnet", "Run the task on mainnet");
.addFlag("mainnet", "Run the task on mainnet")
.addOptionalParam("evm", `Fetch balances for a specific EVM address`)
.addOptionalParam("solana", `Fetch balances for a specific Solana address`)
.addOptionalParam("bitcoin", `Fetch balances for a specific Bitcoin address`);

Check failure on line 144 in packages/tasks/src/balances.ts

View workflow job for this annotation

GitHub Actions / build

Insert `⏎`
12 changes: 12 additions & 0 deletions yarn.lock
Original file line number Diff line number Diff line change
Expand Up @@ -2301,6 +2301,11 @@ base-x@^4.0.0:
resolved "https://registry.yarnpkg.com/base-x/-/base-x-4.0.0.tgz#d0e3b7753450c73f8ad2389b5c018a4af7b2224a"
integrity sha512-FuwxlW4H5kh37X/oW59pwTzzTKRzfrrQwhmyspRM7swOEZcHtDZSCt45U6oKgtuFE+WYPblePMVIPR4RZrh/hw==

base-x@^5.0.0:
version "5.0.0"
resolved "https://registry.yarnpkg.com/base-x/-/base-x-5.0.0.tgz#6d835ceae379130e1a4cb846a70ac4746f28ea9b"
integrity sha512-sMW3VGSX1QWVFA6l8U62MLKz29rRfpTlYdCqLdpLo1/Yd4zZwSbnUaDfciIAowAqvq7YFnWq9hrhdg1KYgc1lQ==

base64-js@^1.3.0, base64-js@^1.3.1:
version "1.5.1"
resolved "https://registry.yarnpkg.com/base64-js/-/base64-js-1.5.1.tgz#1b1b440160a5bf7ad40b650f095963481903930a"
Expand Down Expand Up @@ -2519,6 +2524,13 @@ bs58@^5.0.0:
dependencies:
base-x "^4.0.0"

bs58@^6.0.0:
version "6.0.0"
resolved "https://registry.yarnpkg.com/bs58/-/bs58-6.0.0.tgz#a2cda0130558535dd281a2f8697df79caaf425d8"
integrity sha512-PD0wEnEYg6ijszw/u8s+iI3H17cTymlrwkKhDhPZq+Sokl3AU4htyBFTjAeNAlCCmg0f53g6ih3jATyCKftTfw==
dependencies:
base-x "^5.0.0"

bs58check@<3.0.0, bs58check@^2.1.2:
version "2.1.2"
resolved "https://registry.yarnpkg.com/bs58check/-/bs58check-2.1.2.tgz#53b018291228d82a5aa08e7d796fdafda54aebfc"
Expand Down

0 comments on commit 28d6765

Please sign in to comment.