feat: use the Connect signaturePrivateKey as an additional owner in the smart account (#254)

Co-authored-by: Nik Graf <[email protected]>

Author: pcarranzav

apps/connect/src/Boot.tsx

@@ -18,17 +18,18 @@ declare module '@tanstack/react-router' {
 
 const queryClient = new QueryClient();
 
-const storage = localStorage;
+const addressStorage = localStorage;
+const keysStorage = sessionStorage;
 
 export function Boot() {
   // check if the user is already authenticated on initial render
   const initialRenderAuthCheckRef = useRef(false);
   // using a layout effect to avoid a re-render
   useLayoutEffect(() => {
     if (!initialRenderAuthCheckRef.current) {
-      const accountAddress = Connect.loadAccountAddress(storage);
+      const accountAddress = Connect.loadAccountAddress(addressStorage);
       if (accountAddress) {
-        const keys = Connect.loadKeys(storage, accountAddress);
+        const keys = Connect.loadKeys(keysStorage, accountAddress);
         if (keys) {
           // user is already authenticated, set state
           StoreConnect.store.send({

apps/connect/src/components/LogoutButton.tsx

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

apps/connect/src/routes/authenticate.tsx

@@ -12,6 +12,7 @@ import { Effect, Schema } from 'effect';
 import { TriangleAlert } from 'lucide-react';
 import { useEffect } from 'react';
 import { createWalletClient, custom } from 'viem';
+import { privateKeyToAccount } from 'viem/accounts';
 
 const CHAIN = import.meta.env.VITE_HYPERGRAPH_CHAIN === 'geogenesis' ? Connect.GEOGENESIS : Connect.GEO_TESTNET;
 
@@ -381,9 +382,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,
@@ -396,21 +399,22 @@ function AuthenticateComponent() {
       );
       console.log('smart session created');
       const smartAccountClient = await Connect.getSmartAccountWalletClient({
-        owner: walletClient,
+        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,

apps/connect/src/routes/login.lazy.tsx

@@ -8,7 +8,8 @@ import { type WalletClient, createWalletClient, custom } from 'viem';
 
 const CHAIN = import.meta.env.VITE_HYPERGRAPH_CHAIN === 'geogenesis' ? Connect.GEOGENESIS : Connect.GEO_TESTNET;
 const syncServerUri = import.meta.env.VITE_HYPERGRAPH_SYNC_SERVER_ORIGIN;
-const storage = localStorage;
+const addressStorage = localStorage;
+const keysStorage = sessionStorage;
 
 export const Route = createLazyFileRoute('/login')({
   component: () => <Login />,
@@ -52,7 +53,16 @@ function Login() {
       console.log(walletClient);
       console.log(rpcUrl);
       console.log(CHAIN);
-      await Connect.login({ walletClient, signer, syncServerUri, storage, identityToken, rpcUrl, chain: CHAIN });
+      await Connect.login({
+        walletClient,
+        signer,
+        syncServerUri,
+        addressStorage,
+        keysStorage,
+        identityToken,
+        rpcUrl,
+        chain: CHAIN,
+      });
     },
     [identityToken, signMessage],
   );

apps/server/src/index.ts

@@ -1,9 +1,9 @@
+import { parse } from 'node:url';
 import { Connect, Identity, Inboxes, Messages, SpaceEvents, Utils } from '@graphprotocol/hypergraph';
 import { bytesToHex, randomBytes } from '@noble/hashes/utils.js';
 import cors from 'cors';
 import { Effect, Exit, Schema } from 'effect';
 import express, { type Request, type Response } from 'express';
-import { parse } from 'node:url';
 import WebSocket, { WebSocketServer } from 'ws';
 import { addAppIdentityToSpaces } from './handlers/add-app-identity-to-spaces.js';
 import { applySpaceEvent } from './handlers/applySpaceEvent.js';

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 = (addressStorage: Storage, keysAndTokenStorage: Storage) => {
+  const accountAddress = loadAccountAddress(addressStorage);
+  wipeAccountAddress(addressStorage);
+  if (accountAddress) {
+    wipeKeys(keysAndTokenStorage, accountAddress);
+    wipeSyncServerSessionToken(keysAndTokenStorage, 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,
@@ -29,17 +31,26 @@ export async function signup(
   smartAccountClient: SmartAccountClient,
   accountAddress: Address,
   syncServerUri: string,
-  storage: Storage,
+  addressStorage: Storage,
+  keysStorage: 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 },
@@ -64,8 +75,8 @@ export async function signup(
   if (!decoded.success) {
     throw new Error('Error creating identity');
   }
-  storeKeys(storage, accountAddress, keys);
-  storeAccountAddress(storage, accountAddress);
+  storeKeys(keysStorage, accountAddress, keys);
+  storeAccountAddress(addressStorage, accountAddress);
   return {
     accountAddress,
     keys,
@@ -76,7 +87,8 @@ export async function restoreKeys(
   signer: Signer,
   accountAddress: Address,
   syncServerUri: string,
-  storage: Storage,
+  addressStorage: Storage,
+  keysStorage: Storage,
   identityToken: string,
 ) {
   const res = await fetch(new URL('/connect/identity/encrypted', syncServerUri), {
@@ -93,8 +105,8 @@ export async function restoreKeys(
     const { keyBox } = decoded;
     const { ciphertext, nonce } = keyBox;
     const keys = await decryptIdentity(signer, ciphertext, nonce);
-    storeKeys(storage, accountAddress, keys);
-    storeAccountAddress(storage, accountAddress);
+    storeKeys(keysStorage, accountAddress, keys);
+    storeAccountAddress(addressStorage, accountAddress);
     return {
       accountAddress,
       keys,
@@ -103,8 +115,13 @@ export async function restoreKeys(
   throw new Error(`Error fetching identity ${res.status}`);
 }
 
-const getAndDeploySmartAccount = async (walletClient: WalletClient, rpcUrl: string, chain: Chain, storage: Storage) => {
-  const accountAddressFromStorage = loadAccountAddress(storage) as Hex;
+const getAndUpdateSmartAccount = async (
+  walletClient: WalletClient,
+  rpcUrl: string,
+  chain: Chain,
+  addressStorage: Storage,
+) => {
+  const accountAddressFromStorage = loadAccountAddress(addressStorage) as Hex;
   const smartAccountParams: SmartAccountParams = {
     owner: walletClient,
     rpcUrl,
@@ -128,53 +145,34 @@ 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;
 };
 
 export async function login({
   walletClient,
   signer,
   syncServerUri,
-  storage,
+  addressStorage,
+  keysStorage,
   identityToken,
   rpcUrl,
   chain,
 }: {
   walletClient: WalletClient;
   signer: Signer;
   syncServerUri: string;
-  storage: Storage;
+  addressStorage: Storage;
+  keysStorage: Storage;
   identityToken: string;
   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, addressStorage);
   if (!smartAccountWalletClient.account) {
     throw new Error('Smart account wallet client account not found');
   }
   const accountAddress = smartAccountWalletClient.account.address;
-  // const keys = loadKeys(storage, accountAddress);
+
   let authData: {
     accountAddress: Address;
     keys: IdentityKeys;
@@ -187,11 +185,14 @@ export async function login({
       smartAccountWalletClient,
       accountAddress,
       syncServerUri,
-      storage,
+      addressStorage,
+      keysStorage,
       identityToken,
+      chain,
+      rpcUrl,
     );
   } else {
-    authData = await restoreKeys(signer, accountAddress, syncServerUri, storage, identityToken);
+    authData = await restoreKeys(signer, accountAddress, syncServerUri, addressStorage, keysStorage, identityToken);
   }
   store.send({ type: 'reset' });
   store.send({