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..6abd5eeaa --- /dev/null +++ b/src/api/modules/token_balance/token_balance.background.ts @@ -0,0 +1,20 @@ +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"); + + const balance = + id === AO_NATIVE_TOKEN + ? await getNativeTokenBalance(address) + : (await getAoTokenBalance(address, id)).toString(); + + 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; diff --git a/src/tokens/aoTokens/ao.ts b/src/tokens/aoTokens/ao.ts index 648a9b70d..7be33ac23 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,7 +225,7 @@ export async function getAoTokenBalance( ): Promise { const aoTokens = (await ExtensionStorage.get("ao_tokens")) || []; - const aoToken = aoTokens.find((token) => token.processId === process); + let aoToken = aoTokens.find((token) => token.processId === process); const res = await dryrun({ Id, @@ -232,13 +234,35 @@ export async function getAoTokenBalance( tags: [{ name: "Action", value: "Balance" }] }); + const errorMessage = (res as any)?.error || res?.Error; + + if (errorMessage) { + throw new Error(errorMessage); + } + + 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[]) { 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(0n, 12n); } export async function getNativeTokenBalance(address: string): Promise {