move loginWithWallet function to core and add tests

packages/hypergraph-react/src/HypergraphAppContext.tsx

@@ -4,13 +4,7 @@ import * as automerge from '@automerge/automerge';
 import { uuid } from '@automerge/automerge';
 import { RepoContext } from '@automerge/automerge-repo-react-hooks';
 import { Identity, Key, Messages, SpaceEvents, type SpaceStorageEntry, Utils, store } from '@graphprotocol/hypergraph';
-import {
-  getSessionNonce,
-  identityExists,
-  prepareSiweMessage,
-  restoreKeys,
-  signup,
-} from '@graphprotocol/hypergraph/identity/login';
+import { getSessionNonce, identityExists, prepareSiweMessage } from '@graphprotocol/hypergraph/identity/login';
 import { useSelector as useSelectorStore } from '@xstate/store/react';
 import { Effect, Exit } from 'effect';
 import * as Schema from 'effect/Schema';
@@ -133,58 +127,6 @@ export function HypergraphAppProvider({
   const sessionToken = useSelectorStore(store, (state) => state.context.sessionToken);
   const keys = useSelectorStore(store, (state) => state.context.keys);
 
-  async function loginWithWallet(signer: Identity.Signer, accountId: Address, retryCount = 0) {
-    const sessionToken = Identity.loadSyncServerSessionToken(storage, accountId);
-    if (!sessionToken) {
-      const sessionNonce = await getSessionNonce(accountId, syncServerUri);
-      // Use SIWE to login with the server and get a token
-      const message = prepareSiweMessage(
-        accountId,
-        sessionNonce,
-        { host: window.location.host, origin: window.location.origin },
-        chainId,
-      );
-      const signature = await signer.signMessage(message);
-      const loginReq = { accountId, message, signature } as const satisfies Messages.RequestLogin;
-      const res = await fetch(new URL('/login', syncServerUri), {
-        method: 'POST',
-        headers: {
-          'Content-Type': 'application/json',
-        },
-        body: JSON.stringify(loginReq),
-      });
-      const decoded = Schema.decodeUnknownSync(Messages.ResponseLogin)(await res.json());
-      Identity.storeAccountId(storage, accountId);
-      Identity.storeSyncServerSessionToken(storage, accountId, decoded.sessionToken);
-      const keys = await restoreKeys(signer, accountId, decoded.sessionToken, syncServerUri, storage);
-      return {
-        accountId,
-        sessionToken: decoded.sessionToken,
-        keys,
-      };
-    }
-    // use whoami to check if the session token is still valid
-    const res = await fetch(new URL('/whoami', syncServerUri), {
-      headers: {
-        Authorization: `Bearer ${sessionToken}`,
-      },
-    });
-    if (res.status !== 200 || (await res.text()) !== accountId) {
-      console.warn('Session token is invalid, wiping state and retrying login with wallet');
-      Identity.wipeSyncServerSessionToken(storage, accountId);
-      if (retryCount > 3) {
-        throw new Error('Could not login with wallet after several attempts');
-      }
-      return await loginWithWallet(signer, accountId, retryCount + 1);
-    }
-    const keys = await restoreKeys(signer, accountId, sessionToken, syncServerUri, storage);
-    return {
-      accountId,
-      sessionToken,
-      keys,
-    };
-  }
-
   async function loginWithKeys(keys: Identity.IdentityKeys, accountId: Address, retryCount = 0) {
     const sessionToken = Identity.loadSyncServerSessionToken(storage, accountId);
     if (sessionToken) {
@@ -255,15 +197,16 @@ export function HypergraphAppProvider({
       sessionToken: string;
       keys: Identity.IdentityKeys;
     };
+    const location = {
+      host: window.location.host,
+      origin: window.location.origin,
+    };
     if (!keys && !(await identityExists(accountId, syncServerUri))) {
-      authData = await signup(signer, accountId, syncServerUri, chainId, storage, {
-        host: window.location.host,
-        origin: window.location.origin,
-      });
+      authData = await Identity.signup(signer, accountId, syncServerUri, chainId, storage, location);
     } else if (keys) {
       authData = await loginWithKeys(keys, accountId);
     } else {
-      authData = await loginWithWallet(signer, accountId);
+      authData = await Identity.loginWithWallet(signer, accountId, syncServerUri, chainId, storage, location);
     }
     console.log('Identity initialized');
     store.send({

packages/hypergraph/src/identity/login.ts

@@ -3,7 +3,14 @@ import { SiweMessage } from 'siwe';
 import type { Address, Hex } from 'viem';
 import { privateKeyToAccount } from 'viem/accounts';
 import * as Messages from '../messages/index.js';
-import { loadKeys, storeAccountId, storeKeys, storeSyncServerSessionToken } from './auth-storage.js';
+import {
+  loadKeys,
+  loadSyncServerSessionToken,
+  storeAccountId,
+  storeKeys,
+  storeSyncServerSessionToken,
+  wipeSyncServerSessionToken,
+} from './auth-storage.js';
 import { createIdentityKeys } from './create-identity-keys.js';
 import { decryptIdentity, encryptIdentity } from './identity-encryption.js';
 import { proveIdentityOwnership } from './prove-ownership.js';
@@ -65,7 +72,6 @@ export async function restoreKeys(
     },
   });
   if (res.status === 200) {
-    console.log('Identity found');
     const decoded = Schema.decodeUnknownSync(Messages.ResponseIdentityEncrypted)(await res.json());
     const { keyBox } = decoded;
     const { ciphertext, nonce } = keyBox;
@@ -123,3 +129,58 @@ export async function signup(
     keys,
   };
 }
+
+export async function loginWithWallet(
+  signer: Signer,
+  accountId: Address,
+  syncServerUri: string,
+  chainId: number,
+  storage: Storage,
+  location: { host: string; origin: string },
+  retryCount = 0,
+) {
+  const sessionToken = loadSyncServerSessionToken(storage, accountId);
+  if (!sessionToken) {
+    const sessionNonce = await getSessionNonce(accountId, syncServerUri);
+    // Use SIWE to login with the server and get a token
+    const message = prepareSiweMessage(accountId, sessionNonce, location, chainId);
+    const signature = await signer.signMessage(message);
+    const loginReq = { accountId, message, signature } as const satisfies Messages.RequestLogin;
+    const res = await fetch(new URL('/login', syncServerUri), {
+      method: 'POST',
+      headers: {
+        'Content-Type': 'application/json',
+      },
+      body: JSON.stringify(loginReq),
+    });
+    const decoded = Schema.decodeUnknownSync(Messages.ResponseLogin)(await res.json());
+    storeAccountId(storage, accountId);
+    storeSyncServerSessionToken(storage, accountId, decoded.sessionToken);
+    const keys = await restoreKeys(signer, accountId, decoded.sessionToken, syncServerUri, storage);
+    return {
+      accountId,
+      sessionToken: decoded.sessionToken,
+      keys,
+    };
+  }
+  // use whoami to check if the session token is still valid
+  const res = await fetch(new URL('/whoami', syncServerUri), {
+    headers: {
+      Authorization: `Bearer ${sessionToken}`,
+    },
+  });
+  if (res.status !== 200 || (await res.text()) !== accountId) {
+    console.warn('Session token is invalid, wiping state and retrying login with wallet');
+    wipeSyncServerSessionToken(storage, accountId);
+    if (retryCount > 3) {
+      throw new Error('Could not login with wallet after several attempts');
+    }
+    return await loginWithWallet(signer, accountId, syncServerUri, chainId, storage, location, retryCount + 1);
+  }
+  const keys = await restoreKeys(signer, accountId, sessionToken, syncServerUri, storage);
+  return {
+    accountId,
+    sessionToken,
+    keys,
+  };
+}

packages/hypergraph/test/identity/login.test.ts

@@ -1,7 +1,7 @@
 import { privateKeyToAccount } from 'viem/accounts';
 import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
 import { encryptIdentity } from '../../src/identity/identity-encryption';
-import { getSessionNonce, prepareSiweMessage, restoreKeys, signup } from '../../src/identity/login';
+import { getSessionNonce, loginWithWallet, prepareSiweMessage, restoreKeys, signup } from '../../src/identity/login';
 import type { IdentityKeys, Signer, Storage } from '../../src/identity/types';
 import type * as Messages from '../../src/messages';
 
@@ -124,14 +124,22 @@ describe('restoreKeys', () => {
     removeItem: vi.fn(),
   };
 
+  const encryptedKeys: Messages.ResponseIdentityEncrypted = {
+    keyBox: {
+      accountId: account.address,
+      ciphertext: 'foo',
+      nonce: 'bar',
+    },
+  };
+
   // mock fetch to http://localhost:3000/login/nonce to return a nonce
   beforeEach(() => {
     vi.clearAllMocks();
     vi.spyOn(global, 'fetch').mockImplementation((url) => {
       if (url.toString() === 'http://localhost:3000/identity/encrypted') {
         return Promise.resolve({
           status: 200,
-          json: () => Promise.resolve({ sessionNonce: 'Sv2dJppgx9SKDGCIb' }),
+          json: () => Promise.resolve(encryptedKeys),
         } as Response);
       }
       return vi.fn() as never;
@@ -166,7 +174,7 @@ describe('restoreKeys', () => {
 
     const encryptedIdentity = await encryptIdentity(signer, accountId, identityKeys);
 
-    const mockEncryptedKeys: Messages.ResponseIdentityEncrypted = {
+    const encryptedKeys: Messages.ResponseIdentityEncrypted = {
       keyBox: {
         accountId,
         ciphertext: encryptedIdentity.ciphertext,
@@ -180,7 +188,7 @@ describe('restoreKeys', () => {
       if (url.toString() === 'http://localhost:3000/identity/encrypted') {
         return Promise.resolve({
           status: 200,
-          json: () => Promise.resolve(mockEncryptedKeys),
+          json: () => Promise.resolve(encryptedKeys),
         } as Response);
       }
       return vi.fn() as never;
@@ -191,14 +199,6 @@ describe('restoreKeys', () => {
     expect(result).toEqual(identityKeys);
     expect(mockStorage.getItem).toHaveBeenCalledWith('hypergraph:dev:keys:0x4aAE31951Dfd101d95c2b90e6a8a44b49867346E');
     expect(JSON.parse(mockStorage.setItem.mock.calls[0][1])).toEqual(identityKeys);
-    expect(fetch).toHaveBeenCalledWith(
-      expect.any(URL),
-      expect.objectContaining({
-        headers: {
-          Authorization: 'Bearer session-token',
-        },
-      }),
-    );
   });
 
   it('should throw error if fetch fails', async () => {
@@ -303,3 +303,146 @@ describe('signup', () => {
     );
   });
 });
+
+describe('loginWithWallet', () => {
+  let mockStorage: Storage;
+  const accountPrivateKey = '0xda75e4ea10de7b3ba2d4212cc16bfbb6cac6aed04ac59a28c3231994d8027a9f';
+  const account = privateKeyToAccount(accountPrivateKey);
+
+  const signer: Signer = {
+    signMessage: (message) => {
+      return account.signMessage({ message });
+    },
+    getAddress: () => account.address,
+  };
+
+  const location = { host: 'localhost', origin: 'http://localhost' };
+
+  beforeEach(() => {
+    mockStorage = {
+      getItem: vi.fn(),
+      setItem: vi.fn(),
+      removeItem: vi.fn(),
+    };
+
+    vi.useFakeTimers();
+    vi.setSystemTime(new Date('2024-01-01'));
+  });
+
+  afterEach(() => {
+    vi.useRealTimers();
+    vi.restoreAllMocks();
+  });
+
+  it('should successfully login and restore keys', async () => {
+    const sessionToken = 'test-session-token';
+    const accountId = account.address;
+    const identityKeys: IdentityKeys = {
+      encryptionPrivateKey: 'encryptionPrivateKey',
+      encryptionPublicKey: 'encryptionPublicKey',
+      signaturePrivateKey: 'signaturePrivateKey',
+      signaturePublicKey: 'signaturePublicKey',
+    };
+
+    const encryptedIdentity = await encryptIdentity(signer, accountId, identityKeys);
+
+    const encryptedKeys: Messages.ResponseIdentityEncrypted = {
+      keyBox: {
+        accountId,
+        ciphertext: encryptedIdentity.ciphertext,
+        nonce: encryptedIdentity.nonce,
+      },
+    };
+
+    // @ts-expect-error
+    mockStorage.getItem.mockReturnValue(null);
+
+    // Mock the nonce endpoint
+    vi.spyOn(global, 'fetch').mockImplementation((url) => {
+      if (url.toString() === 'http://localhost:3000/identity/encrypted') {
+        return Promise.resolve({
+          status: 200,
+          json: () => Promise.resolve(encryptedKeys),
+        } as Response);
+      }
+      if (url.toString() === 'http://localhost:3000/login/nonce') {
+        return Promise.resolve({
+          status: 200,
+          json: () => Promise.resolve({ sessionNonce: 'Sv2dJppgx9SKDGCIb' }),
+        } as Response);
+      }
+      if (url.toString() === 'http://localhost:3000/login') {
+        return Promise.resolve({
+          status: 200,
+          json: () => Promise.resolve({ sessionToken }),
+        } as Response);
+      }
+      if (url.toString() === 'http://localhost:3000/whoami') {
+        return Promise.resolve({
+          status: 200,
+          text: () => Promise.resolve(accountId),
+        } as Response);
+      }
+      return vi.fn() as never;
+    });
+
+    const result = await loginWithWallet(signer, accountId, 'http://localhost:3000', 1, mockStorage, location);
+
+    expect(result.accountId).toBe(accountId);
+    expect(result.sessionToken).toBe(sessionToken);
+    expect(result.keys).toEqual(identityKeys);
+  });
+
+  it('should successfully restore keys in case of session token exists', async () => {
+    const sessionToken = 'test-session-token';
+    const accountId = account.address;
+    const identityKeys: IdentityKeys = {
+      encryptionPrivateKey: 'encryptionPrivateKey',
+      encryptionPublicKey: 'encryptionPublicKey',
+      signaturePrivateKey: 'signaturePrivateKey',
+      signaturePublicKey: 'signaturePublicKey',
+    };
+
+    const encryptedIdentity = await encryptIdentity(signer, accountId, identityKeys);
+
+    const encryptedKeys: Messages.ResponseIdentityEncrypted = {
+      keyBox: {
+        accountId,
+        ciphertext: encryptedIdentity.ciphertext,
+        nonce: encryptedIdentity.nonce,
+      },
+    };
+
+    // @ts-expect-error
+    mockStorage.getItem.mockImplementation((key) => {
+      if (key === `hypergraph:dev:session-token:${accountId}`) {
+        return sessionToken;
+      }
+
+      return null;
+    });
+
+    // Mock the nonce endpoint
+    vi.spyOn(global, 'fetch').mockImplementation((url) => {
+      if (url.toString() === 'http://localhost:3000/identity/encrypted') {
+        return Promise.resolve({
+          status: 200,
+          json: () => Promise.resolve(encryptedKeys),
+        } as Response);
+      }
+      if (url.toString() === 'http://localhost:3000/whoami') {
+        return Promise.resolve({
+          status: 200,
+          text: () => Promise.resolve(accountId),
+        } as Response);
+      }
+      return vi.fn() as never;
+    });
+
+    const result = await loginWithWallet(signer, accountId, 'http://localhost:3000', 1, mockStorage, location);
+
+    expect(result.accountId).toBe(accountId);
+    expect(result.sessionToken).toBe(sessionToken);
+    expect(result.keys).toEqual(identityKeys);
+  });
+});