feat: use the Connect signaturePrivateKey as an additional owner in the smart account

Author: pcarranzav

apps/connect/src/components/logout.tsx

@@ -1,3 +1,4 @@
+import { Connect } from '@graphprotocol/hypergraph';
 import { usePrivy } from '@privy-io/react-auth';
 import { useRouter } from '@tanstack/react-router';
 import { Loader2 } from 'lucide-react';
@@ -10,6 +11,7 @@ export function Logout() {
   const [isLoading, setIsLoading] = useState(false);
   const disconnectWallet = async () => {
     setIsLoading(true);
+    Connect.wipeAllAuthData(localStorage);
     await privyLogout();
     router.navigate({
       to: '/login',

apps/connect/src/routes/authenticate.tsx

@@ -3,16 +3,17 @@ import { Button } from '@/components/ui/button';
 import { usePrivateSpaces } from '@/hooks/use-private-spaces';
 import { type PublicSpaceData, usePublicSpaces } from '@/hooks/use-public-spaces';
 import { Connect, Identity, Key, type Messages, StoreConnect, Utils } from '@graphprotocol/hypergraph';
-import { GEOGENESIS, GEO_TESTNET, getSmartAccountWalletClient } from '@graphprotocol/hypergraph/connect/smart-account';
+
 import { useIdentityToken, usePrivy, useWallets } from '@privy-io/react-auth';
 import { createFileRoute } from '@tanstack/react-router';
 import { createStore } from '@xstate/store';
 import { useSelector } from '@xstate/store/react';
 import { Effect, Schema } from 'effect';
 import { useEffect } from 'react';
 import { createWalletClient, custom } from 'viem';
+import { privateKeyToAccount } from 'viem/accounts';
 
-const CHAIN = import.meta.env.VITE_HYPERGRAPH_CHAIN === 'geogenesis' ? GEOGENESIS : GEO_TESTNET;
+const CHAIN = import.meta.env.VITE_HYPERGRAPH_CHAIN === 'geogenesis' ? Connect.GEOGENESIS : Connect.GEO_TESTNET;
 
 type AuthenticateSearch = {
   data: unknown;
@@ -380,9 +381,11 @@ function AuthenticateComponent() {
           })) ?? [];
       console.log('spaces', spaces);
 
+      const localAccount = privateKeyToAccount(keys.signaturePrivateKey as `0x${string}`);
+      console.log('local account', localAccount.address);
       // TODO: add additional actions (must be passed from the app)
       const permissionId = await Connect.createSmartSession(
-        walletClient,
+        localAccount,
         accountAddress,
         newAppIdentity.addressPrivateKey,
         CHAIN,
@@ -394,22 +397,23 @@ function AuthenticateComponent() {
         },
       );
       console.log('smart session created');
-      const smartAccountClient = await getSmartAccountWalletClient({
-        owner: walletClient,
+      const smartAccountClient = await Connect.getSmartAccountWalletClient({
+        owner: localAccount,
         address: accountAddress,
         chain: CHAIN,
         rpcUrl: import.meta.env.VITE_HYPERGRAPH_RPC_URL,
       });
 
+      console.log('encrypting app identity');
       const { ciphertext, nonce } = await Connect.encryptAppIdentity(
         signer,
         newAppIdentity.address,
         newAppIdentity.addressPrivateKey,
         permissionId,
         keys,
       );
+      console.log('proving ownership');
       const { accountProof, keyProof } = await Identity.proveIdentityOwnership(
-        walletClient,
         smartAccountClient,
         accountAddress,
         keys,

packages/hypergraph/src/connect/abis.ts

@@ -52,6 +52,25 @@ export const safeOwnerManagerAbi = [
     stateMutability: 'nonpayable',
     type: 'function',
   },
+  {
+    inputs: [
+      {
+        internalType: 'address',
+        name: 'owner',
+        type: 'address',
+      },
+    ],
+    name: 'isOwner',
+    outputs: [
+      {
+        internalType: 'bool',
+        name: '',
+        type: 'bool',
+      },
+    ],
+    stateMutability: 'view',
+    type: 'function',
+  },
 ];
 
 // We only use this for revokeEnableSignature to use as a noop when creating a smart session

packages/hypergraph/src/connect/auth-storage.ts

@@ -65,3 +65,12 @@ export const storeAccountAddress = (storage: Storage, accountId: string) => {
 export const wipeAccountAddress = (storage: Storage) => {
   storage.removeItem(buildAccountAddressStorageKey());
 };
+
+export const wipeAllAuthData = (storage: Storage) => {
+  const accountAddress = loadAccountAddress(storage);
+  wipeAccountAddress(storage);
+  if (accountAddress) {
+    wipeKeys(storage, accountAddress);
+    wipeSyncServerSessionToken(storage, accountAddress);
+  }
+};

packages/hypergraph/src/connect/login.ts

@@ -1,14 +1,16 @@
 import * as Schema from 'effect/Schema';
 import type { SmartAccountClient } from 'permissionless';
 import type { Address, Chain, Hex, WalletClient } from 'viem';
+import { privateKeyToAccount } from 'viem/accounts';
 import { proveIdentityOwnership } from '../identity/prove-ownership.js';
 import * as Messages from '../messages/index.js';
 import { store } from '../store-connect.js';
-import { loadAccountAddress, storeAccountAddress, storeKeys, wipeAccountAddress } from './auth-storage.js';
+import { loadAccountAddress, storeAccountAddress, storeKeys } from './auth-storage.js';
 import { createIdentityKeys } from './create-identity-keys.js';
 import { decryptIdentity, encryptIdentity } from './identity-encryption.js';
 import {
   type SmartAccountParams,
+  addSmartAccountOwner,
   getSmartAccountWalletClient,
   isSmartAccountDeployed,
   smartAccountNeedsUpdate,
@@ -31,15 +33,23 @@ export async function signup(
   syncServerUri: string,
   storage: Storage,
   identityToken: string,
+  chain: Chain,
+  rpcUrl: string,
 ) {
   const keys = createIdentityKeys();
   const { ciphertext, nonce } = await encryptIdentity(signer, keys);
-  const { accountProof, keyProof } = await proveIdentityOwnership(
-    walletClient,
-    smartAccountClient,
-    accountAddress,
-    keys,
-  );
+
+  const localAccount = privateKeyToAccount(keys.signaturePrivateKey as `0x${string}`);
+  // This will deploy the smart account if it's not deployed
+  await addSmartAccountOwner(smartAccountClient, localAccount.address, chain, rpcUrl);
+  const localSmartAccountClient = await getSmartAccountWalletClient({
+    owner: localAccount,
+    address: accountAddress,
+    rpcUrl,
+    chain,
+  });
+
+  const { accountProof, keyProof } = await proveIdentityOwnership(localSmartAccountClient, accountAddress, keys);
 
   const req: Messages.RequestConnectCreateIdentity = {
     keyBox: { signer: await signer.getAddress(), accountAddress, ciphertext, nonce },
@@ -103,7 +113,7 @@ export async function restoreKeys(
   throw new Error(`Error fetching identity ${res.status}`);
 }
 
-const getAndDeploySmartAccount = async (walletClient: WalletClient, rpcUrl: string, chain: Chain, storage: Storage) => {
+const getAndUpdateSmartAccount = async (walletClient: WalletClient, rpcUrl: string, chain: Chain, storage: Storage) => {
   const accountAddressFromStorage = loadAccountAddress(storage) as Hex;
   const smartAccountParams: SmartAccountParams = {
     owner: walletClient,
@@ -128,21 +138,6 @@ const getAndDeploySmartAccount = async (walletClient: WalletClient, rpcUrl: stri
     // Create the client again to ensure we have the 7579 config now
     return getSmartAccountWalletClient(smartAccountParams);
   }
-  if (!(await isSmartAccountDeployed(smartAccountWalletClient))) {
-    // TODO: remove this once we manage to get counterfactual signatures working
-    console.log('sending dummy userOp to deploy smart account');
-    if (!walletClient.account) {
-      throw new Error('Wallet client account not found');
-    }
-    const tx = await smartAccountWalletClient.sendUserOperation({
-      calls: [{ to: walletClient.account.address, data: '0x' }],
-      account: smartAccountWalletClient.account,
-    });
-
-    console.log('tx', tx);
-    const receipt = await smartAccountWalletClient.waitForUserOperationReceipt({ hash: tx });
-    console.log('receipt', receipt);
-  }
   return smartAccountWalletClient;
 };
 
@@ -163,13 +158,7 @@ export async function login({
   rpcUrl: string;
   chain: Chain;
 }) {
-  let smartAccountWalletClient: SmartAccountClient;
-  try {
-    smartAccountWalletClient = await getAndDeploySmartAccount(walletClient, rpcUrl, chain, storage);
-  } catch (error) {
-    wipeAccountAddress(storage);
-    smartAccountWalletClient = await getAndDeploySmartAccount(walletClient, rpcUrl, chain, storage);
-  }
+  const smartAccountWalletClient = await getAndUpdateSmartAccount(walletClient, rpcUrl, chain, storage);
   if (!smartAccountWalletClient.account) {
     throw new Error('Smart account wallet client account not found');
   }
@@ -189,6 +178,8 @@ export async function login({
       syncServerUri,
       storage,
       identityToken,
+      chain,
+      rpcUrl,
     );
   } else {
     authData = await restoreKeys(signer, accountAddress, syncServerUri, storage, identityToken);

packages/hypergraph/src/connect/smart-account.ts

@@ -36,6 +36,7 @@ import {
   ContractFunctionExecutionError,
   type Hex,
   type Narrow,
+  type PrivateKeyAccount,
   type SignableMessage,
   type WalletClient,
   createPublicClient,
@@ -582,13 +583,61 @@ const getSpaceActions = (space: { address: Hex; type: 'personal' | 'public' }) =
   return actions;
 };
 
+export const addSmartAccountOwner = async (
+  smartAccountClient: SmartAccountClient,
+  newOwner: Address,
+  chain: Chain,
+  rpcUrl: string,
+) => {
+  if (!smartAccountClient.account) {
+    throw new Error('Invalid smart account');
+  }
+  const publicClient = createPublicClient({
+    transport: http(rpcUrl),
+    chain,
+  });
+  if (await isSmartAccountDeployed(smartAccountClient)) {
+    const isOwner = await publicClient.readContract({
+      abi: safeOwnerManagerAbi,
+      address: smartAccountClient.account.address,
+      functionName: 'isOwner',
+      args: [newOwner],
+    });
+    if (isOwner) {
+      return;
+    }
+  }
+
+  const tx = await smartAccountClient.sendUserOperation({
+    calls: [
+      {
+        to: smartAccountClient.account.address,
+        data: encodeFunctionData({
+          abi: safeOwnerManagerAbi,
+          functionName: 'addOwnerWithThreshold',
+          args: [newOwner, BigInt(1)],
+        }),
+        value: BigInt(0),
+      },
+    ],
+    account: smartAccountClient.account,
+  });
+  const receipt = await smartAccountClient.waitForUserOperationReceipt({
+    hash: tx,
+  });
+  if (!receipt.success) {
+    throw new Error('Transaction to add smart account owner failed');
+  }
+  return receipt;
+};
+
 // This is the function that the Connect app uses to create a smart session and
 // enable it on the smart account.
 // It will prompt the user to sign the message to enable the session, and then
 // execute the transaction to enable the session.
 // It will return the permissionId that can be used to create a smart session client.
 export const createSmartSession = async (
-  owner: WalletClient,
+  owner: WalletClient | PrivateKeyAccount,
   accountAddress: Hex,
   sessionPrivateKey: Hex,
   chain: Chain,
@@ -624,9 +673,6 @@ export const createSmartSession = async (
   if (!smartAccountClient.chain) {
     throw new Error('Invalid smart account chain');
   }
-  if (!owner.account) {
-    throw new Error('Invalid wallet client');
-  }
 
   const sessionKeyAccount = privateKeyToAccount(sessionPrivateKey);
   const transport = http(rpcUrl);
@@ -751,9 +797,9 @@ export const createSmartSession = async (
 
   console.log('signing session details');
   // This will prompt the user to sign the message to enable the session
-  sessionDetails.enableSessionData.enableSession.permissionEnableSig = await owner.signMessage({
+  sessionDetails.enableSessionData.enableSession.permissionEnableSig = await (owner as WalletClient).signMessage({
     message: { raw: sessionDetails.permissionEnableHash },
-    account: owner.account.address,
+    account: (owner as WalletClient).account?.address ?? (owner as PrivateKeyAccount).address,
   });
   console.log('session details signed');
   const smartSessions = getSmartSessionsValidator({});

packages/hypergraph/src/identity/prove-ownership.ts

@@ -20,7 +20,6 @@ export const accountProofDomain = {
 };
 
 export const proveIdentityOwnership = async (
-  walletClient: WalletClient,
   smartAccountClient: SmartAccountClient,
   accountAddress: string,
   keys: IdentityKeys,