package.json

 {
   "name": "@hiveio/metamask-snap",
-  "version": "1.2.1",
+  "version": "1.3.0",
   "description": "Hive wallet extension allowing you to sign transactions using keys derived from your Metamask wallet",
   "main": "./dist/bundle.js",
   "files": [

snap.manifest.json

 {
-  "version": "1.2.1",
+  "version": "1.3.0",
   "description": "Hive wallet extension allowing you to sign transactions using keys derived from your Metamask wallet",
   "proposedName": "Hive Wallet",
   "repository": {
@@ -7,7 +7,7 @@
     "url": "git+https://gitlab.syncad.com/hive/metamask-snap.git"
   },
   "source": {
-    "shasum": "MXajRhEb2TjGR6uz1Q/BWyrqfzvNX0WPxcdKx+2ZlaQ=",
+    "shasum": "Qn3MhblQGIyqmEaAKxChWSC9ALvWQW/iOxk8Xxp74Js=",
     "location": {
       "npm": {
         "filePath": "dist/bundle.js",

src/hive/wax.ts

-import { createWaxFoundation, type IWaxBaseInterface } from "@hiveio/wax";
+import { createWaxFoundation, DEFAULT_WAX_OPTIONS, type IWaxBaseInterface } from "@hiveio/wax";
 
-let _wax: undefined | IWaxBaseInterface;
-export const getWax = async (): Promise<IWaxBaseInterface> => {
-  if (!_wax)
-    _wax = await createWaxFoundation();
+const waxInstances: Record<string, IWaxBaseInterface> = {};
+export const getWax = async (chainId: string = DEFAULT_WAX_OPTIONS.chainId): Promise<IWaxBaseInterface> => {
+  if (!waxInstances[chainId])
+    waxInstances[chainId] = await createWaxFoundation(chainId !== undefined ? { chainId } : undefined);
 
-  return _wax;
+  return waxInstances[chainId];
 };

src/index.ts

@@ -27,7 +27,7 @@ export const onRpcRequest = async ({
 
     case 'hive_signTransaction':
       return {
-        signatures: await signTransaction(origin, request.params.transaction, request.params.keys)
+        signatures: await signTransaction(origin, request.params.transaction, request.params.keys, request.params.chainId)
       };
 
     case 'hive_decrypt':

src/rpc.ts

@@ -20,6 +20,7 @@ export type SignTransactionRequest = {
   method: 'hive_signTransaction';
   params: {
     transaction: string;
+    chainId?: string;
     keys: KeyIndex[];
   };
 }

src/snap/decodeBuffer.ts

 import type { KeyIndex } from "../rpc";
 import { getWax } from "../hive/wax";
-import { remove0x } from "@metamask/utils";
-import { keyIndexToPath } from "../utils/key-management";
+import { importPrivateKeyToWallet } from "../utils/key-management";
 import { getTempWallet } from "../hive/beekeeper";
 import { ConfirmBufferDecode } from "./dialogs/ConfirmBufferDecode";
-import { SLIP10Node } from "@metamask/key-tree";
 import type { THexString } from "@hiveio/wax";
 
 export const decodeBuffer = async (origin: string, buffer: THexString, decodeKey: KeyIndex): Promise<string> => {
@@ -18,22 +16,7 @@ export const decodeBuffer = async (origin: string, buffer: THexString, decodeKey
   const wallet = await getTempWallet();
 
   try {
-    const snapResponse = await snap.request({
-      method: 'snap_getBip32Entropy',
-      params: {
-        curve: "secp256k1",
-        path: keyIndexToPath(decodeKey)
-      }
-    });
-
-    const node = await SLIP10Node.fromJSON(snapResponse);
-
-    if (!node.privateKey)
-      throw new Error('No private key found');
-
-    const wif = wax.convertRawPrivateKeyToWif(remove0x(node.privateKey));
-
-    await wallet.importKey(wif);
+    await importPrivateKeyToWallet(wallet, decodeKey);
 
     const response = wax.decrypt(wallet, buffer);
 

src/snap/dialogs/ConfirmBufferDecode.tsx

-import { Bold, Copyable, Text, Box, Italic } from "@metamask/snaps-sdk/jsx";
+import { Bold, Copyable, Text, Box } from "@metamask/snaps-sdk/jsx";
 import type { KeyIndex } from "../../rpc";
+import { KeyTypeNotice } from "./components/KeyTypeNotice";
 
 export const ConfirmBufferDecode = (origin: string, buffer: string, decodeKey: KeyIndex) => snap.request({
   method: 'snap_dialog',
@@ -14,9 +15,7 @@ export const ConfirmBufferDecode = (origin: string, buffer: string, decodeKey: K
         <Text>
           Confirm if you want to sign it using your:
         </Text>
-        <Text>
-          - <Bold>{decodeKey.role}</Bold> key (account index: <Italic>#{ String(decodeKey.accountIndex ?? 0) }</Italic>)
-        </Text>
+        {KeyTypeNotice(decodeKey)}
       </Box>
     )
   }

src/snap/dialogs/ConfirmBufferSign.tsx

-import { Bold, Copyable, Text, Box, Italic } from "@metamask/snaps-sdk/jsx";
+import { Bold, Copyable, Text, Box } from "@metamask/snaps-sdk/jsx";
 import { KeyIndex } from "../../rpc";
+import { KeyTypeNotice } from "./components/KeyTypeNotice";
 
 export const ConfirmBufferSign = (origin: string, buffer: string, firstKey: KeyIndex, secondKey?: KeyIndex | string) => snap.request({
   method: 'snap_dialog',
@@ -14,14 +15,11 @@ export const ConfirmBufferSign = (origin: string, buffer: string, firstKey: KeyI
         <Text>
           Confirm if you want to sign it using your:
         </Text>
-        <Text>
-          - <Bold>{firstKey.role}</Bold> key (account index: <Italic>#{ String(firstKey.accountIndex ?? 0) }</Italic>)
-        </Text>
+        {KeyTypeNotice(firstKey)}
         <Text>This message will be encrypted for: </Text>
-        { secondKey ? (typeof secondKey === "string" ?
-          <Copyable value={ secondKey }/> :
-          <Text>- <Bold>{secondKey.role}</Bold> key (account index: <Italic>#{ String(secondKey.accountIndex ?? 0) }</Italic>)</Text>
-        ) : <Text>- <Bold>yourself</Bold></Text> }
+        { secondKey && typeof secondKey === "string" && <Copyable value={ secondKey }/> }
+        { secondKey && typeof secondKey === "object" && KeyTypeNotice(secondKey) }
+        { !secondKey && <Text>- <Bold>yourself</Bold></Text> }
       </Box>
     )
   }

src/snap/dialogs/ConfirmTransactionSign.tsx

-import { Bold, Copyable, Text, Box, Italic } from "@metamask/snaps-sdk/jsx";
+import { Bold, Copyable, Text, Box, Banner, Italic } from "@metamask/snaps-sdk/jsx";
 import { KeyIndex } from "../../rpc";
+import { KeyTypeNotice } from "./components/KeyTypeNotice";
+import { DEFAULT_WAX_OPTIONS } from "@hiveio/wax";
 
-export const ConfirmTransactionSign = (origin: string, transaction: string, keys: KeyIndex[]) => snap.request({
+export const ConfirmTransactionSign = (origin: string, transaction: string, keys: KeyIndex[], chainId?: string) => snap.request({
   method: 'snap_dialog',
   params: {
     type: 'confirmation',
     content: (
       <Box>
         <Text>
-          <Bold>{ origin }</Bold> asked to sign a transaction:
+          <Bold>{origin}</Bold> asked to sign a transaction:
         </Text>
-        <Copyable value={ transaction } />
+        <Copyable value={transaction} />
+        {chainId && chainId !== DEFAULT_WAX_OPTIONS.chainId && <Banner title="Custom chain signing" severity="warning">
+          <Text>
+            <Bold>Warning:</Bold> You are signing this transaction for a custom chain. Make sure you trust this chain: <Italic>{chainId}</Italic>
+          </Text>
+        </Banner>}
         <Text>
           Confirm if you want to sign it using your:
         </Text>
-        { keys.map(key => (<Text>
-          - <Bold>{key.role}</Bold> key (account index: <Italic>#{ String(key.accountIndex ?? 0) }</Italic>)
-        </Text>))}
+        {KeyTypeNotice(...keys)}
       </Box>
     )
   }

src/snap/dialogs/components/KeyTypeNotice.tsx

+import { Bold, Text, Italic } from "@metamask/snaps-sdk/jsx";
+import type { KeyIndex } from "../../../rpc";
+
+export const KeyTypeNotice = (...keys: KeyIndex[]) => keys.map(key => (
+  <Text>
+    - <Bold>{key.role}</Bold> key (account index: <Italic>#{ String(key.accountIndex ?? 0) }</Italic>)
+  </Text>
+));

src/snap/encodeBuffer.ts

 import type { KeyIndex } from "../rpc";
 import { getWax } from "../hive/wax";
-import { remove0x } from "@metamask/utils";
-import { keyIndexToPath } from "../utils/key-management";
+import { importPrivateKeyToWallet } from "../utils/key-management";
 import { getTempWallet } from "../hive/beekeeper";
 import { ConfirmBufferSign } from "./dialogs/ConfirmBufferSign";
-import { SLIP10Node } from "@metamask/key-tree";
 
 export const encodeBuffer = async (origin: string, buffer: string, firstKey: KeyIndex, secondKey?: KeyIndex | string): Promise<string> => {
   const confirmDecode = await ConfirmBufferSign(origin, buffer, firstKey, secondKey);
@@ -17,38 +15,13 @@ export const encodeBuffer = async (origin: string, buffer: string, firstKey: Key
   const wallet = await getTempWallet();
 
   try {
-    const firstKeyBip32 = await snap.request({
-      method: 'snap_getBip32Entropy',
-      params: {
-        curve: "secp256k1",
-        path: keyIndexToPath(firstKey)
-      }
-    });
-    const nodeFirstKey = await SLIP10Node.fromJSON(firstKeyBip32);
-    if (!nodeFirstKey.privateKey)
-      throw new Error('No private key found');
-
-    const wifFirstKey = wax.convertRawPrivateKeyToWif(remove0x(nodeFirstKey.privateKey));
-    const publicKeyFirstKey = await wallet.importKey(wifFirstKey);
+    const publicKeyFirstKey = await importPrivateKeyToWallet(wallet, firstKey);
 
     let publicKeySecondKey: string | undefined;
     if(typeof secondKey === "string")
       publicKeySecondKey = secondKey;
-    else if (secondKey) {
-      const secondKeyBip32 = await snap.request({
-        method: 'snap_getBip32Entropy',
-        params: {
-          curve: "secp256k1",
-          path: keyIndexToPath(secondKey)
-        }
-      });
-      const nodeSecondKey = await SLIP10Node.fromJSON(secondKeyBip32);
-      if (!nodeSecondKey.privateKey)
-        throw new Error('No private key found');
-
-      const wifSecondKey = wax.convertRawPrivateKeyToWif(remove0x(nodeSecondKey.privateKey));
-      publicKeySecondKey = await wallet.importKey(wifSecondKey);
-    }
+    else if (secondKey)
+      publicKeySecondKey = await importPrivateKeyToWallet(wallet, secondKey);
 
     const response = wax.encrypt(wallet, buffer, publicKeyFirstKey, publicKeySecondKey);
 

src/snap/getPublicKeys.ts

 import type { KeyIndex, PublicKeyData } from "../rpc";
-import { getWax } from "../hive/wax";
-import { remove0x } from "@metamask/utils";
-import { keyIndexToPath } from "../utils/key-management";
+import { getPublicKeyWifFromKeyIndex } from "../utils/key-management";
 
 export const getPublicKeys = async (keys: KeyIndex[]): Promise<PublicKeyData[]> => {
-  const wax = await getWax();
-
   const publicKeys: PublicKeyData[] = [];
 
   for(const key of keys) {
-    const bip32 = await snap.request({
-      method: 'snap_getBip32PublicKey',
-      params: {
-        curve: "secp256k1",
-        path: keyIndexToPath(key),
-        compressed: true
-      }
-    });
-
-    const publicKey = wax.convertRawPublicKeyToWif(remove0x(bip32));
+    const publicKey = await getPublicKeyWifFromKeyIndex(key);
 
     publicKeys.push({
       accountIndex: key.accountIndex ?? 0,

src/snap/signTransaction.ts

 import type { KeyIndex } from "../rpc";
 import { getWax } from "../hive/wax";
-import { remove0x } from "@metamask/utils";
-import { keyIndexToPath } from "../utils/key-management";
+import { importPrivateKeyToWallet } from "../utils/key-management";
 import { getTempWallet } from "../hive/beekeeper";
 import { ConfirmTransactionSign } from "./dialogs/ConfirmTransactionSign";
 import type { THexString } from "@hiveio/wax";
-import { SLIP10Node } from "@metamask/key-tree";
 
-export const signTransaction = async (origin: string, transaction: string, keys: KeyIndex[]): Promise<THexString[]> => {
+export const signTransaction = async (origin: string, transaction: string, keys: KeyIndex[], chainId?: string): Promise<THexString[]> => {
   if (keys.length < 1)
     throw new Error('No keys provided');
 
-  const confirmSign = await ConfirmTransactionSign(origin, transaction, keys);
+  const confirmSign = await ConfirmTransactionSign(origin, transaction, keys, chainId);
 
   if(!confirmSign)
     throw new Error('User denied the transaction');
 
   // The order is important: First create wax, then transaction and if all success then create wallet
-  const wax = await getWax();
+  const wax = await getWax(chainId);
   const tx = wax.createTransactionFromJson(transaction);
   const wallet = await getTempWallet();
 
@@ -25,22 +23,7 @@ export const signTransaction = async (origin: string, transaction: string, keys:
     const signatures: THexString[] = [];
 
     for(const key of keys) {
-      const snapResponse = await snap.request({
-        method: 'snap_getBip32Entropy',
-        params: {
-          curve: "secp256k1",
-          path: keyIndexToPath(key)
-        }
-      });
-
-      const node = await SLIP10Node.fromJSON(snapResponse);
-
-      if (!node.privateKey)
-        throw new Error('No private key found');
-
-      const wif = wax.convertRawPrivateKeyToWif(remove0x(node.privateKey));
-
-      const publicKey = await wallet.importKey(wif);
+      const publicKey = await importPrivateKeyToWallet(wallet, key);
 
       const signature = tx.sign(wallet, publicKey);
 

src/utils/key-management.ts

-import type { TRole } from "@hiveio/wax";
+import type { TPublicKey, TRole } from "@hiveio/wax";
 import type { KeyIndex } from "../rpc";
+import { getWax } from "../hive/wax";
+import { remove0x } from "@metamask/utils";
+import { SLIP10Node } from "@metamask/key-tree";
+import type { IBeekeeperUnlockedWallet } from "@hiveio/beekeeper";
 
 const KeyIndexToPathMap = {
   owner: 0,
@@ -15,3 +19,49 @@ export const keyIndexToPath = (keyIndex: KeyIndex): string[] => {
 
   return ["m", "48'", "13'", `${keyIndexType}'`, `${keyIndex.accountIndex ?? 0}'`, "0'"];
 };
+
+export const getPublicKeyWifFromKeyIndex = async (keyIndex: KeyIndex): Promise<TPublicKey> => {
+  const wax = await getWax();
+
+  const bip32 = await snap.request({
+    method: 'snap_getBip32PublicKey',
+    params: {
+      curve: "secp256k1",
+      path: keyIndexToPath(keyIndex),
+      compressed: true
+    }
+  });
+
+  const publicKey = wax.convertRawPublicKeyToWif(remove0x(bip32));
+
+  return publicKey;
+};
+
+export const getPrivateKeyWifFromKeyIndex = async (keyIndex: KeyIndex): Promise<string> => {
+  const wax = await getWax();
+
+  const snapResponse = await snap.request({
+    method: 'snap_getBip32Entropy',
+    params: {
+      curve: "secp256k1",
+      path: keyIndexToPath(keyIndex)
+    }
+  });
+
+  const node = await SLIP10Node.fromJSON(snapResponse);
+
+  if (!node.privateKey)
+    throw new Error('No private key found');
+
+  const wif = wax.convertRawPrivateKeyToWif(remove0x(node.privateKey));
+
+  return wif;
+};
+
+export const importPrivateKeyToWallet = async (wallet: IBeekeeperUnlockedWallet, keyIndex: KeyIndex): Promise<TPublicKey> => {
+  const privateKeyWif = await getPrivateKeyWifFromKeyIndex(keyIndex);
+
+  const publicKey = await wallet.importKey(privateKeyWif);
+
+  return publicKey;
+};