[WIP] Keystone USB integration

package.json

@@ -39,8 +39,8 @@
     "@avalabs/core-wallets-sdk": "3.1.0-alpha.32",
     "@avalabs/evm-module": "1.2.0",
     "@avalabs/glacier-sdk": "3.1.0-alpha.32",
-    "@avalabs/hw-app-avalanche": "0.14.1",
     "@avalabs/hvm-module": "1.2.0",
+    "@avalabs/hw-app-avalanche": "0.14.1",
     "@avalabs/types": "3.1.0-alpha.32",
     "@avalabs/vm-module-types": "1.2.0",
     "@blockaid/client": "0.10.0",
@@ -50,8 +50,12 @@
     "@ethereumjs/common": "2.6.5",
     "@ethereumjs/tx": "3.4.0",
     "@hpke/core": "1.2.5",
+    "@keystonehq/alias-sampling": "^0.1.2",
     "@keystonehq/animated-qr": "0.5.2",
     "@keystonehq/bc-ur-registry-eth": "0.15.2",
+    "@keystonehq/hw-app-avalanche": "^0.0.2",
+    "@keystonehq/hw-app-eth": "^0.4.4",
+    "@keystonehq/hw-transport-webusb": "^0.5.1",
     "@keystonehq/ur-decoder": "0.9.2",
     "@ledgerhq/hw-app-btc": "10.2.4",
     "@ledgerhq/hw-app-eth": "6.36.1",

src/background/connections/extensionConnection/models.ts

@@ -102,6 +102,8 @@ export enum ExtensionRequest {
   LEDGER_VERSION_WARNING_CLOSED = 'ledger_version_warning_closed',
   LEDGER_MIGRATE_MISSING_PUBKEYS = 'ledger_migrate_missing_pubkeys',
 
+  KEYSTONE_INIT_TRANSPORT = 'keystone_init_transport',
+  KEYSTONE_CLOSE_TRANSPORT = 'keystone_close_transport',
   KEYSTONE_SUBMIT_SIGNATURE = 'keystone_submit_signature',
 
   NAVIGATION_HISTORY_GET = 'navigation_history_get',

src/background/services/featureFlags/models.ts

@@ -13,6 +13,7 @@ export enum FeatureGates {
   BUY_MOONPAY = 'buy-feature-moonpay',
   BUY_COINBASE = 'buy-feature-coinbase',
   KEYSTONE = 'keystone',
+  KEYSTONE_USB = 'keystone-usb',
   NFT_MARKETPLACE = 'nft-marketplace',
   BOTTOM_NAVIGATION = 'bottom-navigation',
   DEFI = 'defi-feature',
@@ -65,6 +66,7 @@ export const DISABLED_FLAG_VALUES: FeatureFlags = {
   [FeatureGates.BUY_MOONPAY]: false,
   [FeatureGates.BUY_COINBASE]: false,
   [FeatureGates.KEYSTONE]: false,
+  [FeatureGates.KEYSTONE_USB]: false,
   [FeatureGates.NFT_MARKETPLACE]: false,
   [FeatureGates.BOTTOM_NAVIGATION]: false,
   [FeatureGates.DEFI]: false,
@@ -116,6 +118,7 @@ export const DEFAULT_FLAGS: FeatureFlags = {
   [FeatureGates.BUY_MOONPAY]: true,
   [FeatureGates.BUY_COINBASE]: true,
   [FeatureGates.KEYSTONE]: true,
+  [FeatureGates.KEYSTONE_USB]: true,
   [FeatureGates.NFT_MARKETPLACE]: true,
   [FeatureGates.BOTTOM_NAVIGATION]: true,
   [FeatureGates.DEFI]: true,

src/background/services/keystone/BitcoinKeystoneWallet.ts

@@ -1,7 +1,11 @@
 import { BitcoinProviderAbstract, createPsbt } from '@avalabs/core-wallets-sdk';
 import { CryptoPSBT } from '@keystonehq/bc-ur-registry-eth';
+import KeystoneUSBEthSDK from '@keystonehq/hw-app-eth';
 import { Psbt, Transaction } from 'bitcoinjs-lib';
+import { createKeystoneTransport } from '@keystonehq/hw-transport-webusb';
+import { UREncoder } from '@ngraveio/bc-ur';
 
+import { parseResponoseUR } from './KeystoneWallet';
 import { BtcTransactionRequest } from '../wallet/models';
 import { KeystoneTransport } from './models';
 
@@ -13,6 +17,7 @@ export class BitcoinKeystoneWallet {
     private keystoneTransport: KeystoneTransport,
     private provider: BitcoinProviderAbstract,
     private tabId?: number,
+    private viaUSB?: boolean,
   ) {}
 
   async signTx(
@@ -34,12 +39,26 @@ export class BitcoinKeystoneWallet {
     });
 
     const cryptoPSBT = new CryptoPSBT(psbt.toBuffer());
-    const { type, cbor } = cryptoPSBT.toUR();
+    const ur = cryptoPSBT.toUR();
+
+    if (this.viaUSB) {
+      const app = new KeystoneUSBEthSDK(
+        (await createKeystoneTransport()) as any,
+      );
+      const encodedUR = new UREncoder(ur!, Infinity).nextPart().toUpperCase();
+      const signedCborBuffer = parseResponoseUR(
+        (await app.signTransactionFromUr(encodedUR)).payload,
+      ).cbor;
+
+      const signedTx = CryptoPSBT.fromCBOR(signedCborBuffer).getPSBT();
+
+      return Psbt.fromBuffer(signedTx).extractTransaction();
+    }
 
     const signedCborBuffer = await this.keystoneTransport.requestSignature(
       {
-        type,
-        cbor: cbor.toString('hex'),
+        type: ur.type,
+        cbor: ur.cbor.toString('hex'),
       },
       this.tabId,
     );

src/background/services/keystone/KeystoneWallet.ts

@@ -12,22 +12,66 @@ import {
 } from '@keystonehq/bc-ur-registry-eth';
 import { TransactionRequest, hexlify } from 'ethers';
 import { BufferLike, rlp } from 'ethereumjs-util';
-
+import { Avalanche } from '@avalabs/core-wallets-sdk';
+import KeystoneUSBAvalancheSDK from '@keystonehq/hw-app-avalanche';
+import KeystoneUSBEthSDK from '@keystonehq/hw-app-eth';
+import { createKeystoneTransport } from '@keystonehq/hw-transport-webusb';
+import { UREncoder, URDecoder, UR } from '@ngraveio/bc-ur';
 import { makeBNLike } from '@src/utils/makeBNLike';
 
 import { CBOR, KeystoneTransport } from './models';
 import { convertTxData } from './utils';
 
+export const parseResponoseUR = (urPlayload: string): UR => {
+  const decoder = new URDecoder();
+  decoder.receivePart(urPlayload);
+  if (!decoder.isComplete()) {
+    throw new Error('UR incomplete');
+  }
+  const resultUR = decoder.resultUR();
+  return resultUR;
+};
+
 export class KeystoneWallet {
   constructor(
     private fingerprint: string,
     private activeAccountIndex: number,
     private keystoneTransport: KeystoneTransport,
     private chainId?: number,
     private tabId?: number,
+    private xpub?: string,
+    private xpubXP?: string,
   ) {}
 
   async signTransaction(txRequest: TransactionRequest): Promise<string> {
+    const isKeystone3 = !!this.xpubXP;
+    if (isKeystone3) {
+      const app = new KeystoneUSBEthSDK(
+        (await createKeystoneTransport()) as any,
+      );
+      const ur = await this.buildSignatureUR(
+        txRequest,
+        this.fingerprint,
+        this.activeAccountIndex,
+      );
+      const encodedUR = new UREncoder(ur!, Infinity).nextPart().toUpperCase();
+      const payload = (await app.signTransactionFromUr(encodedUR)).payload;
+      const signature = ETHSignature.fromCBOR(
+        parseResponoseUR(payload).cbor,
+      ).getSignature();
+
+      const r = hexlify(new Uint8Array(signature.slice(0, 32)));
+      const s = hexlify(new Uint8Array(signature.slice(32, 64)));
+      const v = new BN(signature.slice(64)).toNumber();
+
+      const signedTx = await this.getTxFromTransactionRequest(txRequest, {
+        r,
+        s,
+        v,
+      });
+
+      return '0x' + signedTx.serialize().toString('hex');
+    }
     const cborBuffer = await this.keystoneTransport.requestSignature(
       await this.buildSignatureRequest(
         txRequest,
@@ -51,11 +95,11 @@ export class KeystoneWallet {
     return '0x' + signedTx.serialize().toString('hex');
   }
 
-  private async buildSignatureRequest(
+  private async buildSignatureUR(
     txRequest: TransactionRequest,
     fingerprint: string,
     activeAccountIndex: number,
-  ): Promise<CBOR> {
+  ): Promise<UR> {
     const chainId = txRequest.chainId ?? this.chainId;
     const isLegacyTx = typeof txRequest.gasPrice !== 'undefined';
 
@@ -85,6 +129,19 @@ export class KeystoneWallet {
 
     const ur = ethSignRequest.toUR();
 
+    return ur;
+  }
+
+  private async buildSignatureRequest(
+    txRequest: TransactionRequest,
+    fingerprint: string,
+    activeAccountIndex: number,
+  ): Promise<CBOR> {
+    const ur = await this.buildSignatureUR(
+      txRequest,
+      fingerprint,
+      activeAccountIndex,
+    );
     return {
       cbor: ur.cbor.toString('hex'),
       type: ur.type,
@@ -148,4 +205,24 @@ export class KeystoneWallet {
       type: type || undefined,
     };
   }
+
+  /**
+   * Avalanche P/X chain transaction signing
+   */
+  public async signTx(
+    txRequest: Avalanche.SignTxRequest,
+  ): Promise<Avalanche.SignTxRequest['tx']> {
+    const tx = txRequest.tx;
+    const isEvmChain = tx.getVM() === 'EVM';
+    const app = new KeystoneUSBAvalancheSDK(await createKeystoneTransport());
+    const sig = await app.signTx(
+      tx as any,
+      this.fingerprint,
+      isEvmChain ? this.xpub! : this.xpubXP!,
+      this.activeAccountIndex,
+    );
+
+    tx.addSignature(Buffer.from(sig, 'hex'));
+    return tx;
+  }
 }

src/background/services/onboarding/handlers/keystoneOnboardingHandler.test.ts

@@ -108,7 +108,6 @@ describe('src/background/services/onboarding/handlers/keystoneOnboardingHandler.
     const request = getRequest([
       {
         xpub: 'xpub',
-        xpubXP: '',
         password: 'password',
         analyticsConsent: false,
         masterFingerprint: 'masterFingerprint',

src/background/services/onboarding/handlers/keystoneOnboardingHandler.ts

@@ -21,6 +21,7 @@ type HandlerType = ExtensionRequestHandler<
     {
       masterFingerprint: string;
       xpub: string;
+      xpubXP: string;
       password: string;
       analyticsConsent: boolean;
       walletName?: string;
@@ -44,8 +45,14 @@ export class KeystoneOnboardingHandler implements HandlerType {
   ) {}
 
   handle: HandlerType['handle'] = async ({ request }) => {
-    const { masterFingerprint, xpub, password, analyticsConsent, walletName } =
-      (request.params ?? [])[0] ?? {};
+    const {
+      masterFingerprint,
+      xpub,
+      xpubXP,
+      password,
+      analyticsConsent,
+      walletName,
+    } = (request.params ?? [])[0] ?? {};
 
     await startOnboarding({
       settingsService: this.settingsService,
@@ -56,8 +63,9 @@ export class KeystoneOnboardingHandler implements HandlerType {
     });
 
     const walletId = await this.walletService.init({
-      secretType: SecretType.Keystone,
+      secretType: xpubXP ? SecretType.Keystone3Pro : SecretType.Keystone,
       xpub,
+      xpubXP,
       masterFingerprint,
       derivationPath: DerivationPath.BIP44,
       name: walletName,

src/background/services/onboarding/models.ts

@@ -15,6 +15,7 @@ export enum OnboardingPhase {
   LEDGER_TROUBLE = 'ledger_trouble',
   ANALYTICS_CONSENT = 'analytics_consent',
   KEYSTONE = 'keystone',
+  KEYSTONE_USB = 'keystoneUsb',
   KEYSTONE_TUTORIAL = 'keystone_tutorial',
   SEEDLESS_GOOGLE = 'seedless_google',
   SEEDLESS_APPLE = 'seedless_apple',
@@ -25,6 +26,7 @@ export enum OnboardingURLs {
   CREATE_WALLET = '/onboarding/create-wallet',
   SEED_PHRASE = '/onboarding/seed-phrase',
   KEYSTONE = '/onboarding/keystone',
+  KEYSTONE_USB = '/onboarding/Keystone-usb',
   LEDGER = '/onboarding/ledger',
   CREATE_PASSWORD = '/onboarding/create-password',
   ANALYTICS_CONSENT = '/onboarding/analytics-consent',
@@ -39,6 +41,7 @@ export const ONBOARDING_EVENT_NAMES = {
   [OnboardingPhase.IMPORT_WALLET]: 'OnboardingImportMnemonicSelected',
   [OnboardingPhase.LEDGER]: 'OnboardingImportLedgerSelected',
   [OnboardingPhase.KEYSTONE]: 'OnboardingKeystoneSelected',
+  [OnboardingPhase.KEYSTONE_USB]: 'OnboardingKeystoneUSBSelected',
   [OnboardingPhase.SEEDLESS_GOOGLE]: 'OnboardingSeedlessGoogleSelected',
   [OnboardingPhase.SEEDLESS_APPLE]: 'OnboardingSeedlessAppleSelected',
 };

src/background/services/secrets/SecretsService.ts

@@ -62,6 +62,7 @@ export class SecretsService implements OnUnlock {
       [SecretType.Ledger]: 'Ledger',
       [SecretType.LedgerLive]: 'Ledger Live',
       [SecretType.Keystone]: 'Keystone',
+      [SecretType.Keystone3Pro]: 'Keystone3 Pro',
       [SecretType.Seedless]: 'Seedless',
     };
     const storedSecrets = await this.#loadSecrets(false);
@@ -762,7 +763,8 @@ export class SecretsService implements OnUnlock {
     if (
       secrets.secretType === SecretType.Ledger ||
       secrets.secretType === SecretType.Mnemonic ||
-      secrets.secretType === SecretType.Keystone
+      secrets.secretType === SecretType.Keystone ||
+      secrets.secretType === SecretType.Keystone3Pro
     ) {
       // C-avax... this address uses the same public key as EVM
       const cPubkey = getAddressPublicKeyFromXPub(secrets.xpub, index);

src/background/services/secrets/models.ts

@@ -19,6 +19,7 @@ export enum SecretType {
   Ledger = 'ledger',
   LedgerLive = 'ledger-live',
   Keystone = 'keystone',
+  Keystone3Pro = 'keystone3-pro',
   Seedless = 'seedless',
   // Importable wallets types
   PrivateKey = 'private-key',
@@ -57,10 +58,10 @@ interface MnemonicSecrets extends PrimarySecretsBase {
 }
 
 interface KeystoneSecrets extends PrimarySecretsBase {
-  secretType: SecretType.Keystone;
+  secretType: SecretType.Keystone | SecretType.Keystone3Pro;
   masterFingerprint: string;
   xpub: string;
-  xpubXP?: never;
+  xpubXP?: string;
   derivationPath: DerivationPath;
 }
 

src/background/services/storage/schemaMigrations/migrations/wallet_v4/utils/getSecretsType.ts

@@ -46,6 +46,10 @@ export function getSecretsType(walletKeys: WalletKeys) {
     return SecretType.Keystone;
   }
 
+  if (masterFingerprint && xpub && xpubXP) {
+    return SecretType.Keystone3Pro;
+  }
+
   if (xpub) {
     return SecretType.Ledger;
   }

src/background/services/storage/schemaMigrations/migrations/wallet_v4/wallet_v4.ts

@@ -52,7 +52,8 @@ const up = async (walletStorage: PreviousSchema) => {
       ? 'Recovery Phrase'
       : secretType === SecretType.Seedless
         ? 'Seedless'
-        : secretType === SecretType.Keystone
+        : secretType === SecretType.Keystone ||
+            secretType === SecretType.Keystone3Pro
           ? 'Keystone'
           : 'Ledger';