Merge pull request #18 from Layr-Labs/jb/custom-derivation-paths-2

[1.1.2] Support custom derivation paths

Author: jbrower95

CHANGELOG.md

@@ -1,5 +1,8 @@
 
 **[Current]** 
+1.1.2:
+- Bump to support custom derivation paths.
+
 1.1.1:
 - Minor patches to the zeus state machine to make errors more legible.
 

package.json

@@ -1,6 +1,6 @@
 {
   "name": "@layr-labs/zeus",
-  "version": "1.1.1",
+  "version": "1.1.2",
   "description": "web3 deployer / metadata manager",
   "main": "src/index.ts",
   "scripts": {

src/commands/prompts.ts

@@ -1,5 +1,5 @@
 import { select } from './utils';
-import { privateKeyToAccount } from 'viem/accounts';
+import { mnemonicToAccount, privateKeyToAccount } from 'viem/accounts';
 import { search, input, password as inquirerPassword } from '@inquirer/prompts';
 import chalk from 'chalk';
 import * as AllChains from "viem/chains";
@@ -125,8 +125,32 @@ export const privateKey: (chainId: number, overridePrompt?: string) => Promise<`
     return res as `0x${string}`;
 }
 
-export const accountIndex = async () => {
-    const cont = await wouldYouLikeToContinue("Would you like to use a custom bip39 account index? (NOTE: the default 'm/44'/60'/0'/0/[0]s' will be used otherwise)");
+export const bip32Path = async () => {
+    const cont = await wouldYouLikeToContinue("Would you like to use a custom bip-32 derivation path? (NOTE: the default 'm/44'/60'/0'/0/0' will be used otherwise)");
+    if (!cont) {
+        return `m/44'/60'/0'/0/0`
+    }
+
+    const val = await envVarOrPrompt({
+        title: `Enter the derivation path (default m/44'/60'/0'/0/0)`,
+        directEntryInputType: 'text',
+        reuseKey: `derivationPathSuffix`,
+        isValid: (val: string) => {
+            try {
+                const account = mnemonicToAccount('legal winner thank year wave sausage worth useful legal winner thank yellow') 
+                account.getHdKey().derive(val);
+                return true;
+            } catch {
+                return false;
+            }
+        }
+    })
+
+    return val;
+}
+
+export const addressIndex = async () => {
+    const cont = await wouldYouLikeToContinue("Would you like to use a custom bip39 address index? (NOTE: the default 'm/44'/60'/0'/0/[0]s' will be used otherwise)");
     if (!cont) {
         return 0;
     }

src/signing/strategies/eoa/ledger.ts

@@ -4,32 +4,31 @@ import { ICachedArg, TStrategyOptions } from "../../strategy";
 import { SavebleDocument, Transaction } from "../../../metadata/metadataStore";
 import { TDeploy } from "../../../metadata/schema";
 import * as prompts from '../../../commands/prompts';
-import { DEFAULT_BASE_DERIVATION_PATH } from "../ledgerTransport";
 
 export class LedgerSigningStrategy extends EOABaseSigningStrategy {
     id = "ledger";
     description = "Signing w/ ledger";
 
-    public accountIndex: ICachedArg<number> 
+    public bip32Path: ICachedArg<string> 
 
     constructor(deploy: SavebleDocument<TDeploy>, transaction: Transaction, options?: TStrategyOptions) {
         super(deploy, transaction, options);
-        this.accountIndex = this.arg(async () => {
-            return await prompts.accountIndex();
-        }, 'accountIndex')
+        this.bip32Path = this.arg(async () => {
+            return await prompts.bip32Path();
+        }, 'bip32path')
     }
 
     async getSignerAddress(): Promise<`0x${string}`> {
         console.warn(`If your ledger is not working, you may need to open LedgerLive, navigate to: Accounts -> <Signer> -> Receive and follow the prompts on device. Once your Ledger says "Application is Ready", you can force quit LedgerLive and retry Zeus.`)
-        const signer = await getLedgerAccount(await this.accountIndex.get());
+        const signer = await getLedgerAccount(await this.bip32Path.get());
         return await signer.address as `0x${string}`;
     }
 
     async subclassForgeArgs(): Promise<string[]> {
         const derivationPathArgs = await (async () => {
             try {
-                const accountIndex = await this.accountIndex.getImmediately();
-                return [`--mnemonic-derivation-paths`, `${DEFAULT_BASE_DERIVATION_PATH}/${accountIndex}`]
+                const path = await this.bip32Path.getImmediately();
+                return [`--mnemonic-derivation-paths`, path]
             } catch {
                 return [];
             }

src/signing/strategies/gnosis/api/gnosisLedger.ts

@@ -15,17 +15,17 @@ export class GnosisLedgerStrategy extends GnosisApiStrategy {
     id = "gnosis.api.ledger";
     description = "Gnosis API - Ledger (Not For Private Hotfixes)";
 
-    public accountIndex: ICachedArg<number> 
+    public bip32Path: ICachedArg<string> 
 
     constructor(deploy: SavebleDocument<TDeploy>, transaction: Transaction, options?: TStrategyOptions) {
         super(deploy, transaction, options);
-        this.accountIndex = this.arg(async () => {
-            return await prompts.accountIndex();
-        }, 'accountIndex')
+        this.bip32Path = this.arg(async () => {
+            return await prompts.bip32Path();
+        }, 'bip32path')
     }
 
     async getSignature(version: string, txn: SafeTransaction, safeAddress: `0x${string}`): Promise<`0x${string}`> {
-        const signer = await getLedgerAccount(await this.accountIndex.get());
+        const signer = await getLedgerAccount(await this.bip32Path.get());
         if (!signer.signTypedData) {
             throw new Error(`This ledger does not support signing typed data, and cannot be used with zeus.`);
         }
@@ -81,7 +81,7 @@ export class GnosisLedgerStrategy extends GnosisApiStrategy {
         try {
             while (true) {
                 try {
-                    const accountIndex = await this.accountIndex.get();
+                    const accountIndex = await this.bip32Path.get();
                     const signer = await getLedgerAccount(accountIndex);
                     console.log(`Detected ledger address: ${signer.address}`);    
 

src/signing/strategies/ledgerTransport.ts

@@ -1,17 +1,17 @@
-import {ledgerToAccount, BASE_DERIVATION_PATH} from './ledger/account'
+import {ledgerToAccount} from './ledger/account'
 import { Account } from "viem";
 
 export const DEFAULT_BASE_DERIVATION_PATH = `m/44'/60'/0'/0`
 
-const ledgerAccounts: Record<number, Account> = {};
+const ledgerAccounts: Record<string, Account> = {};
 
-export const getLedgerAccount = async (accountIndex = 0) => {
-    if (ledgerAccounts[accountIndex] === undefined) {
-        ledgerAccounts[accountIndex] = await ledgerToAccount({
-            derivationPath: `${BASE_DERIVATION_PATH}/${accountIndex}`
+export const getLedgerAccount = async (derivationPath = DEFAULT_BASE_DERIVATION_PATH) => {
+    if (ledgerAccounts[derivationPath] === undefined) {
+        ledgerAccounts[derivationPath] = await ledgerToAccount({
+            derivationPath: derivationPath.slice(2) as `44'/60'/0'/0${string}`
         })
     }
 
-    return ledgerAccounts[accountIndex];
+    return ledgerAccounts[derivationPath];
 }
 

src/signing/strategy.ts

@@ -72,7 +72,7 @@ export interface TExecuteOptions {
     nonInteractive?: boolean,
     fork?: TForkType,
     etherscanApiKey?: string | boolean,
-    accountIndex?: number,
+    bip32path?: string,
     anvil?: AnvilService,
     testClient?: TestClient,
     rpcUrl?: string,