From 62b98556db6bbfffa828f5c00998dc4370388d96 Mon Sep 17 00:00:00 2001 From: 7i7o Date: Tue, 15 Oct 2024 14:20:46 +0200 Subject: [PATCH 1/4] add tokenBalance api endpoint --- src/api/background.ts | 3 +++ src/api/foreground.ts | 3 +++ src/api/modules/token_balance/index.ts | 11 ++++++++ .../token_balance/token_balance.background.ts | 27 +++++++++++++++++++ .../token_balance/token_balance.foreground.ts | 6 +++++ 5 files changed, 50 insertions(+) create mode 100644 src/api/modules/token_balance/index.ts create mode 100644 src/api/modules/token_balance/token_balance.background.ts create mode 100644 src/api/modules/token_balance/token_balance.foreground.ts diff --git a/src/api/background.ts b/src/api/background.ts index 0be5da8e5..2b5df2526 100644 --- a/src/api/background.ts +++ b/src/api/background.ts @@ -45,6 +45,8 @@ import subscriptionModule from "./modules/subscription"; import subscription from "./modules/subscription/subscription.background"; import userTokensModule from "./modules/user_tokens"; import userTokens from "./modules/user_tokens/user_tokens.background"; +import tokenBalanceModule from "./modules/token_balance"; +import tokenBalance from "./modules/token_balance/token_balance.background"; /** Background modules */ const modules: BackgroundModule[] = [ @@ -69,6 +71,7 @@ const modules: BackgroundModule[] = [ { ...signDataItemModule, function: signDataItem }, { ...subscriptionModule, function: subscription }, { ...userTokensModule, function: userTokens }, + { ...tokenBalanceModule, function: tokenBalance }, { ...batchSignDataItemModule, function: batchSignDataItem } ]; diff --git a/src/api/foreground.ts b/src/api/foreground.ts index cd0ca12b2..bd079a7b7 100644 --- a/src/api/foreground.ts +++ b/src/api/foreground.ts @@ -65,6 +65,8 @@ import signDataItem, { } from "./modules/sign_data_item/sign_data_item.foreground"; import userTokensModule from "./modules/user_tokens"; import userTokens from "./modules/user_tokens/user_tokens.foreground"; +import tokenBalanceModule from "./modules/token_balance"; +import tokenBalance from "./modules/token_balance/token_balance.foreground"; /** Foreground modules */ const modules: ForegroundModule[] = [ @@ -101,6 +103,7 @@ const modules: ForegroundModule[] = [ }, { ...subscriptionModule, function: subscription }, { ...userTokensModule, function: userTokens }, + { ...tokenBalanceModule, function: tokenBalance }, { ...batchSignDataItemModule, function: batchSignDataItem, diff --git a/src/api/modules/token_balance/index.ts b/src/api/modules/token_balance/index.ts new file mode 100644 index 000000000..9e341fac4 --- /dev/null +++ b/src/api/modules/token_balance/index.ts @@ -0,0 +1,11 @@ +import type { PermissionType } from "~applications/permissions"; +import type { ModuleProperties } from "~api/module"; + +const permissions: PermissionType[] = ["ACCESS_TOKENS"]; + +const tokenBalanceModule: ModuleProperties = { + functionName: "tokenBalance", + permissions +}; + +export default tokenBalanceModule; diff --git a/src/api/modules/token_balance/token_balance.background.ts b/src/api/modules/token_balance/token_balance.background.ts new file mode 100644 index 000000000..eb8723241 --- /dev/null +++ b/src/api/modules/token_balance/token_balance.background.ts @@ -0,0 +1,27 @@ +import type { ModuleFunction } from "~api/background"; +import { ExtensionStorage } from "~utils/storage"; +import { getAoTokenBalance, getNativeTokenBalance } from "~tokens/aoTokens/ao"; +import { AO_NATIVE_TOKEN } from "~utils/ao_import"; +import { isAddress } from "~utils/assertions"; + +const background: ModuleFunction = async (_, id?: string) => { + // validate input + isAddress(id); + const address = await ExtensionStorage.get("active_address"); + + let balance: string | null = null; + try { + if (id === AO_NATIVE_TOKEN) { + balance = await getNativeTokenBalance(address); + } else { + const balanceResult = await getAoTokenBalance(address, id); + balance = balanceResult.toString(); + } + } catch (error) { + console.error(`Error fetching balance for token ${id}:`, error); + } + + return balance; +}; + +export default background; diff --git a/src/api/modules/token_balance/token_balance.foreground.ts b/src/api/modules/token_balance/token_balance.foreground.ts new file mode 100644 index 000000000..220afbdbe --- /dev/null +++ b/src/api/modules/token_balance/token_balance.foreground.ts @@ -0,0 +1,6 @@ +import type { ModuleFunction } from "~api/module"; + +// no need to transform anything in the foreground +const foreground: ModuleFunction = () => {}; + +export default foreground; From d8af6eb73396df127295ed40fa4d8edf4aed6d69 Mon Sep 17 00:00:00 2001 From: Pawan Paudel Date: Tue, 15 Oct 2024 23:44:49 +0545 Subject: [PATCH 2/4] refactor: improve error handling for AO token balance retrieval --- .../token_balance/token_balance.background.ts | 15 ++++----------- src/tokens/aoTokens/ao.ts | 7 +++++++ 2 files changed, 11 insertions(+), 11 deletions(-) diff --git a/src/api/modules/token_balance/token_balance.background.ts b/src/api/modules/token_balance/token_balance.background.ts index eb8723241..6abd5eeaa 100644 --- a/src/api/modules/token_balance/token_balance.background.ts +++ b/src/api/modules/token_balance/token_balance.background.ts @@ -9,17 +9,10 @@ const background: ModuleFunction = async (_, id?: string) => { isAddress(id); const address = await ExtensionStorage.get("active_address"); - let balance: string | null = null; - try { - if (id === AO_NATIVE_TOKEN) { - balance = await getNativeTokenBalance(address); - } else { - const balanceResult = await getAoTokenBalance(address, id); - balance = balanceResult.toString(); - } - } catch (error) { - console.error(`Error fetching balance for token ${id}:`, error); - } + const balance = + id === AO_NATIVE_TOKEN + ? await getNativeTokenBalance(address) + : (await getAoTokenBalance(address, id)).toString(); return balance; }; diff --git a/src/tokens/aoTokens/ao.ts b/src/tokens/aoTokens/ao.ts index 648a9b70d..e686d58ec 100644 --- a/src/tokens/aoTokens/ao.ts +++ b/src/tokens/aoTokens/ao.ts @@ -225,6 +225,10 @@ export async function getAoTokenBalance( const aoToken = aoTokens.find((token) => token.processId === process); + if (!aoToken) { + throw new Error(`Token not found for address '${address}'.`); + } + const res = await dryrun({ Id, Owner: address, @@ -239,6 +243,9 @@ export async function getAoTokenBalance( return new Quantity(BigInt(balance), BigInt(aoToken.Denomination)); } } + + // default return + return new Quantity(0, BigInt(aoToken.Denomination)); } export async function getNativeTokenBalance(address: string): Promise { From 971fa76c60bd0b64f48aaea7e6b2ac6150bf5754 Mon Sep 17 00:00:00 2001 From: Pawan Paudel Date: Wed, 16 Oct 2024 23:16:30 +0545 Subject: [PATCH 3/4] refactor: Update getAoTokenBalance with error handling --- src/tokens/aoTokens/ao.ts | 27 ++++++++++++++++++++------- 1 file changed, 20 insertions(+), 7 deletions(-) diff --git a/src/tokens/aoTokens/ao.ts b/src/tokens/aoTokens/ao.ts index e686d58ec..449002567 100644 --- a/src/tokens/aoTokens/ao.ts +++ b/src/tokens/aoTokens/ao.ts @@ -16,6 +16,8 @@ import { import type { Alarms } from "webextension-polyfill"; import type { KeystoneSigner } from "~wallets/hardware/keystone"; import browser from "webextension-polyfill"; +import { fetchTokenByProcessId } from "~lib/transactions"; +import { tokenTypeRegistry } from "~tokens/token"; export type AoInstance = ReturnType; @@ -223,11 +225,7 @@ export async function getAoTokenBalance( ): Promise { const aoTokens = (await ExtensionStorage.get("ao_tokens")) || []; - const aoToken = aoTokens.find((token) => token.processId === process); - - if (!aoToken) { - throw new Error(`Token not found for address '${address}'.`); - } + let aoToken = aoTokens.find((token) => token.processId === process); const res = await dryrun({ Id, @@ -236,16 +234,31 @@ export async function getAoTokenBalance( tags: [{ name: "Action", value: "Balance" }] }); + if (!res?.Messages) { + throw new Error("Balance handler missing"); + } + + if ((res as any)?.error) { + throw new Error((res as any)?.error); + } + for (const msg of res.Messages as Message[]) { const balance = getTagValue("Balance", msg.Tags); - if (balance && aoToken) { + if (balance && +balance) { + if (!aoToken) { + aoToken = await fetchTokenByProcessId(process); + if (!aoToken) { + throw new Error("Could not load token info."); + } + } + return new Quantity(BigInt(balance), BigInt(aoToken.Denomination)); } } // default return - return new Quantity(0, BigInt(aoToken.Denomination)); + return new Quantity(0n, 12n); } export async function getNativeTokenBalance(address: string): Promise { From b5528b5ffe77b57528b358b3f584416d8f3692b6 Mon Sep 17 00:00:00 2001 From: Pawan Paudel Date: Thu, 17 Oct 2024 10:55:19 +0545 Subject: [PATCH 4/4] refactor: throw error if result Messages is empty --- src/tokens/aoTokens/ao.ts | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/src/tokens/aoTokens/ao.ts b/src/tokens/aoTokens/ao.ts index 449002567..7be33ac23 100644 --- a/src/tokens/aoTokens/ao.ts +++ b/src/tokens/aoTokens/ao.ts @@ -234,12 +234,16 @@ export async function getAoTokenBalance( tags: [{ name: "Action", value: "Balance" }] }); - if (!res?.Messages) { - throw new Error("Balance handler missing"); + const errorMessage = (res as any)?.error || res?.Error; + + if (errorMessage) { + throw new Error(errorMessage); } - if ((res as any)?.error) { - throw new Error((res as any)?.error); + if (res.Messages.length === 0) { + throw new Error( + "Invalid token process: Balance action handler missing or unsupported." + ); } for (const msg of res.Messages as Message[]) {