feat: generateSignature helper for registerKeysOnBehalf

package.json

@@ -1,6 +1,6 @@
 {
   "name": "@scopelift/stealth-address-sdk",
-  "version": "0.2.0",
+  "version": "0.2.2",
   "main": "dist/index.js",
   "types": "dist/index.d.ts",
   "module": "dist/index.js",

src/lib/actions/prepareRegisterKeysOnBehalf/prepareRegisterKeysOnBehalf.test.ts

@@ -1,8 +1,8 @@
 import { beforeAll, describe, expect, test } from 'bun:test';
 import type { Address, TransactionReceipt } from 'viem';
 import {
-  ERC6538RegistryAbi,
   VALID_SCHEME_ID,
+  generateSignatureForRegisterKeysOnBehalf,
   parseStealthMetaAddressURI
 } from '../../..';
 import setupTestEnv from '../../helpers/test/setupTestEnv';
@@ -39,57 +39,20 @@ describe('prepareRegisterKeysOnBehalf', () => {
     const chain = walletClient.chain;
     if (!chain) throw new Error('No chain found');
 
-    const generateSignature = async (account: Address) => {
-      // Get the registrant's current nonce for the signature
-      const nonce = await walletClient.readContract({
-        address: ERC6538Address,
-        abi: ERC6538RegistryAbi,
-        functionName: 'nonceOf',
-        args: [account]
-      });
-
-      // Prepare the signature domain
-      const domain = {
-        name: 'ERC6538Registry',
-        version: '1.0',
-        chainId,
-        verifyingContract: ERC6538Address
-      } as const;
-
-      // Taken from the ERC6538Registry contract
-      const primaryType = 'Erc6538RegistryEntry';
-
-      // Prepare the signature types
-      const types = {
-        [primaryType]: [
-          { name: 'schemeId', type: 'uint256' },
-          { name: 'stealthMetaAddress', type: 'bytes' },
-          { name: 'nonce', type: 'uint256' }
-        ]
-      } as const;
-
-      const message = {
-        schemeId: BigInt(schemeId),
-        stealthMetaAddress: stealthMetaAddressToRegister,
-        nonce
-      };
-
-      const signature = await walletClient.signTypedData({
-        account,
-        primaryType,
-        domain,
-        types,
-        message
-      });
-
-      return signature;
-    };
+    const signature = await generateSignatureForRegisterKeysOnBehalf({
+      walletClient,
+      account,
+      ERC6538Address,
+      chainId,
+      schemeId,
+      stealthMetaAddressToRegister
+    });
 
     args = {
       registrant: account,
       schemeId,
       stealthMetaAddress: stealthMetaAddressToRegister,
-      signature: await generateSignature(account)
+      signature
     } satisfies RegisterKeysOnBehalfArgs;
 
     const prepared = await stealthClient.prepareRegisterKeysOnBehalf({

src/lib/helpers/test/setupTestEnv.test.ts

@@ -1,8 +1,21 @@
-import { beforeEach, describe, expect, mock, test } from 'bun:test';
+import {
+  afterAll,
+  afterEach,
+  beforeEach,
+  describe,
+  expect,
+  mock,
+  test
+} from 'bun:test';
 import { VALID_CHAINS } from '../types';
 import { LOCAL_ENDPOINT } from './setupTestEnv';
 
 describe('setupTestEnv with different environment configurations', () => {
+  afterEach(() => {
+    process.env.USE_FORK = undefined;
+    process.env.RPC_URL = undefined;
+  });
+
   test('should use local node endpoint url when USE_FORK is true and RPC_URL is defined', async () => {
     const exampleRpcUrl = 'http://example-rpc-url.com';
     process.env.USE_FORK = 'true';
@@ -61,6 +74,11 @@ describe('fetchChainId', async () => {
     process.env.RPC_URL = 'http://example-rpc-url.com';
   });
 
+  afterAll(() => {
+    process.env.USE_FORK = undefined;
+    process.env.RPC_URL = undefined;
+  });
+
   test('successful fetch returns chain ID', async () => {
     mock.module('./setupTestEnv', () => ({
       fetchJson: () =>

src/lib/helpers/test/setupTestEnv.ts

@@ -73,7 +73,7 @@ const getChainInfo = async () => {
 
 export const fetchChainId = async (): Promise<number> => {
   // If not running fork test script, use the foundry chain ID
-  if (!process.env.USE_FORK) return foundry.id;
+  if (!isUsingFork()) return foundry.id;
 
   if (!process.env.RPC_URL) {
     throw new Error('RPC_URL not defined in env');
@@ -115,5 +115,10 @@ const fetchJson = async <T>(url: string, options: FetchRequestInit) => {
   return response.json() as T;
 };
 
+function isUsingFork(): boolean {
+  const useFork = process.env.USE_FORK;
+  return useFork === 'true';
+}
+
 export { getValidChainId, getRpcUrl, getChainInfo, fetchJson };
 export default setupTestEnv;

src/utils/helpers/generateSignatureForRegisterKeysOnBehalf.ts

@@ -0,0 +1,91 @@
+import { readContract } from 'viem/actions';
+import { ERC6538RegistryAbi } from '../../lib';
+import {
+  GenerateSignatureForRegisterKeysError,
+  type GenerateSignatureForRegisterKeysParams
+} from './types';
+
+const DOMAIN_NAME = 'ERC6538Registry';
+const DOMAIN_VERSION = '1.0';
+const PRIMARY_TYPE = 'Erc6538RegistryEntry';
+
+const SIGNATURE_TYPES = {
+  [PRIMARY_TYPE]: [
+    { name: 'schemeId', type: 'uint256' },
+    { name: 'stealthMetaAddress', type: 'bytes' },
+    { name: 'nonce', type: 'uint256' }
+  ]
+} as const;
+
+/**
+ * Generates a typed signature for registering keys on behalf of a user (account) in the ERC6538 Registry.
+ *
+ * This function creates an EIP-712 compliant signature for the `registerKeysOnBehalf` function
+ * in the ERC6538 Registry contract. It retrieves the current nonce for the account, prepares
+ * the domain separator and message, and signs the data using the provided Viem wallet client.
+ *
+ * @param {GenerateSignatureForRegisterKeysParams} params - The parameters for generating the signature.
+ * @returns {Promise<`0x${string}`>} A promise that resolves to the generated signature as a hexadecimal string.
+ *
+ * @throws {GenerateSignatureForRegisterKeysError} If the contract read fails or if the signing process encounters an issue.
+ *
+ * @example
+ * const signature = await generateSignatureForRegisterKeysOnBehalf({
+ *   walletClient,
+ *   account: '0x1234...5678',
+ *   ERC6538Address: '0xabcd...ef01',
+ *   chainId: 1,
+ *   schemeId: 1,
+ *   stealthMetaAddressToRegister: '0x9876...5432'
+ * });
+ */
+async function generateSignatureForRegisterKeysOnBehalf({
+  walletClient,
+  account,
+  ERC6538Address,
+  chainId,
+  schemeId,
+  stealthMetaAddressToRegister
+}: GenerateSignatureForRegisterKeysParams): Promise<`0x${string}`> {
+  try {
+    // Get the registrant's current nonce for the signature
+    const nonce = await readContract(walletClient, {
+      address: ERC6538Address,
+      abi: ERC6538RegistryAbi,
+      functionName: 'nonceOf',
+      args: [account]
+    });
+
+    // Prepare the signature domain
+    const domain = {
+      name: DOMAIN_NAME,
+      version: DOMAIN_VERSION,
+      chainId,
+      verifyingContract: ERC6538Address
+    } as const;
+
+    const message = {
+      schemeId: BigInt(schemeId),
+      stealthMetaAddress: stealthMetaAddressToRegister,
+      nonce
+    };
+
+    const signature = await walletClient.signTypedData({
+      account,
+      primaryType: PRIMARY_TYPE,
+      domain,
+      types: SIGNATURE_TYPES,
+      message
+    });
+
+    return signature;
+  } catch (error) {
+    console.error('Error generating signature:', error);
+    throw new GenerateSignatureForRegisterKeysError(
+      'Failed to generate signature for registerKeysOnBehalf',
+      error
+    );
+  }
+}
+
+export default generateSignatureForRegisterKeysOnBehalf;

src/utils/helpers/index.ts

@@ -1,5 +1,11 @@
+export type {
+  GenerateSignatureForRegisterKeysError,
+  GenerateSignatureForRegisterKeysParams
+} from './types';
+
 export { default as generateKeysFromSignature } from './generateKeysFromSignature';
 export { default as generateRandomStealthMetaAddress } from './generateRandomStealthMetaAddress';
+export { default as generateSignatureForRegisterKeysOnBehalf } from './generateSignatureForRegisterKeysOnBehalf';
 export { default as generateStealthMetaAddressFromKeys } from './generateStealthMetaAddressFromKeys';
 export { default as generateStealthMetaAddressFromSignature } from './generateStealthMetaAddressFromSignature';
 export { default as getViewTagFromMetadata } from './getViewTagFromMetadata';

src/utils/helpers/test/generateKeysFromSignature.test.ts

@@ -1,9 +1,9 @@
-import { afterAll, beforeAll, describe, expect, mock, test } from 'bun:test';
+import { beforeAll, describe, expect, test } from 'bun:test';
 import { signMessage } from 'viem/actions';
 import setupTestWallet from '../../../lib/helpers/test/setupTestWallet';
 import type { SuperWalletClient } from '../../../lib/helpers/types';
 import type { HexString } from '../../crypto/types';
-import generateKeysFromSignature from '../generateKeysFromSignature';
+import generateKeysFromSignature, {} from '../generateKeysFromSignature';
 import isValidPublicKey from '../isValidPublicKey';
 
 describe('generateKeysFromSignature', () => {
@@ -13,7 +13,6 @@ describe('generateKeysFromSignature', () => {
   beforeAll(async () => {
     walletClient = await setupTestWallet();
     if (!walletClient.account) throw new Error('No account found');
-
     // Generate a signature to use in the tests
     signature = await signMessage(walletClient, {
       account: walletClient.account,
@@ -22,35 +21,18 @@ describe('generateKeysFromSignature', () => {
     });
   });
 
-  afterAll(() => {
-    mock.restore();
-  });
-
   test('should generate valid public keys from a correct signature', () => {
     const result = generateKeysFromSignature(signature);
-
     expect(isValidPublicKey(result.spendingPublicKey)).toBe(true);
     expect(isValidPublicKey(result.viewingPublicKey)).toBe(true);
+    expect(result.spendingPrivateKey).toBeDefined();
+    expect(result.viewingPrivateKey).toBeDefined();
   });
 
   test('should throw an error for an invalid signature', () => {
-    const invalidSignature = '0x123';
-
-    expect(() => {
-      generateKeysFromSignature(invalidSignature);
-    }).toThrow('Invalid signature');
-  });
-
-  test('should throw an error for incorrectly parsed signatures', () => {
-    const notMatchingSignature = '0x123';
-
-    // Mock the output from extractPortions to return an signature that doesn't match the one passed in
-    mock.module('../generateKeysFromSignature', () => ({
-      extractPortions: () => notMatchingSignature
-    }));
-
+    const invalid = '0x123';
     expect(() => {
-      generateKeysFromSignature(signature);
-    }).toThrow('Signature incorrectly generated or parsed');
+      generateKeysFromSignature(invalid);
+    }).toThrow(`Invalid signature: ${invalid}`);
   });
 });

src/utils/helpers/test/generateSignatureForRegisterKeysOnBehalf.test.ts

@@ -0,0 +1,79 @@
+import { beforeAll, describe, expect, test } from 'bun:test';
+import type { WalletClient } from 'viem';
+import { signMessage } from 'viem/actions';
+import setupTestEnv from '../../../lib/helpers/test/setupTestEnv';
+import setupTestWallet from '../../../lib/helpers/test/setupTestWallet';
+import type { VALID_CHAIN_IDS } from '../../../lib/helpers/types';
+import { VALID_SCHEME_ID } from '../../crypto';
+import generateSignatureForRegisterKeysOnBehalf from '../generateSignatureForRegisterKeysOnBehalf';
+import generateStealthMetaAddressFromSignature from '../generateStealthMetaAddressFromSignature';
+import { GenerateSignatureForRegisterKeysError } from '../types';
+
+describe('generateSignatureForRegisterKeysOnBehalf', () => {
+  let params: Parameters<typeof generateSignatureForRegisterKeysOnBehalf>[0];
+  let walletClient: WalletClient;
+
+  beforeAll(async () => {
+    walletClient = await setupTestWallet();
+    const { ERC6538Address } = await setupTestEnv();
+    const account = walletClient.account;
+    const chainId = walletClient.chain?.id as VALID_CHAIN_IDS | undefined;
+    if (!account) throw new Error('No account found');
+    if (!chainId) throw new Error('No chain found');
+
+    const signatureForStealthMetaAddress = await signMessage(walletClient, {
+      account,
+      message:
+        'Sign this message to generate your stealth address keys.\nChain ID: 31337'
+    });
+
+    const stealthMetaAddressToRegister =
+      generateStealthMetaAddressFromSignature(signatureForStealthMetaAddress);
+
+    params = {
+      walletClient,
+      account: account.address,
+      ERC6538Address,
+      chainId,
+      schemeId: VALID_SCHEME_ID.SCHEME_ID_1,
+      stealthMetaAddressToRegister
+    };
+  });
+
+  test('should generate signature successfully', async () => {
+    const result = await generateSignatureForRegisterKeysOnBehalf(params);
+    expect(result).toBeTypeOf('string');
+    expect(result.startsWith('0x')).toBe(true);
+  });
+
+  test('should throw GenerateSignatureForRegisterKeysError when contract read fails', async () => {
+    const invalidParams = {
+      ...params,
+      ERC6538Address:
+        '0x1234567890123456789012345678901234567890' as `0x${string}`
+    };
+
+    expect(
+      generateSignatureForRegisterKeysOnBehalf(invalidParams)
+    ).rejects.toBeInstanceOf(GenerateSignatureForRegisterKeysError);
+  });
+
+  test('should throw GenerateSignatureForRegisterKeysError when signing fails', async () => {
+    // Create a wallet client that will fail on signTypedData
+    const failingWalletClient = {
+      ...walletClient,
+      signTypedData: async () => {
+        throw new Error('Signing failed');
+      }
+    } as WalletClient;
+
+    const failingParams = {
+      ...params,
+      walletClient: failingWalletClient
+    };
+
+    expect(
+      generateSignatureForRegisterKeysOnBehalf(failingParams)
+    ).rejects.toBeInstanceOf(GenerateSignatureForRegisterKeysError);
+  });
+});

src/utils/helpers/test/generateStealthMetaAddressFromSignature.test.ts

@@ -6,7 +6,7 @@ import { VALID_SCHEME_ID, parseKeysFromStealthMetaAddress } from '../../crypto';
 import type { HexString } from '../../crypto/types';
 import generateStealthMetaAddressFromSignature from '../generateStealthMetaAddressFromSignature';
 
-describe('getStealthMetaAddressFromSignature', () => {
+describe('generateStealthMetaAddressFromSignature', () => {
   let walletClient: SuperWalletClient;
   let signature: HexString;