fixed loginWithKeys function and moved it to core + add tests

apps/server/src/handlers/getIdentity.ts

@@ -1,9 +1,14 @@
 import { prisma } from '../prisma.js';
 
-type Params = {
-  accountId?: string;
-  signaturePublicKey?: string;
-};
+type Params =
+  | {
+      accountId: string;
+      signaturePublicKey?: string;
+    }
+  | {
+      accountId?: string;
+      signaturePublicKey: string;
+    };
 
 export const getIdentity = async (params: Params) => {
   if (!params.accountId && !params.signaturePublicKey) {

apps/server/src/index.ts

@@ -1,8 +1,8 @@
-import { parse } from 'node:url';
 import { Identity, Messages, SpaceEvents, Utils } from '@graphprotocol/hypergraph';
 import cors from 'cors';
 import { Effect, Exit, Schema } from 'effect';
 import express, { type Request, type Response } from 'express';
+import { parse } from 'node:url';
 import { SiweMessage } from 'siwe';
 import type { Hex } from 'viem';
 import WebSocket, { WebSocketServer } from 'ws';
@@ -97,7 +97,7 @@ app.post('/login/with-signing-key', async (req, res) => {
   try {
     const message = Schema.decodeUnknownSync(Messages.RequestLoginWithSigningKey)(req.body);
     const accountId = message.accountId;
-    const nonce = await getSessionNonce({ accountId });
+
     // getIdentity will throw if it doesn't exist
     const identity = await getIdentity({ signaturePublicKey: message.publicKey });
     if (identity.accountId !== accountId) {
@@ -109,6 +109,7 @@ app.post('/login/with-signing-key', async (req, res) => {
       res.status(401).send('Unauthorized');
       return;
     }
+    const nonce = await getSessionNonce({ accountId });
     const { data: siweMessage } = await siweObject.verify({ signature: message.signature, nonce });
     if (!siweMessage.expirationTime) {
       res.status(400).send('Expiration time not set');

packages/hypergraph-react/src/HypergraphAppContext.tsx

@@ -5,7 +5,6 @@ import { uuid } from '@automerge/automerge';
 import type { DocHandle } from '@automerge/automerge-repo';
 import { RepoContext } from '@automerge/automerge-repo-react-hooks';
 import { Identity, Key, Messages, SpaceEvents, type SpaceStorageEntry, Utils, store } from '@graphprotocol/hypergraph';
-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';
@@ -19,9 +18,7 @@ import {
   useRef,
   useState,
 } from 'react';
-import type { Hex } from 'viem';
 import { type Address, getAddress } from 'viem';
-import { privateKeyToAccount } from 'viem/accounts';
 
 const decodeResponseMessage = Schema.decodeUnknownEither(Messages.ResponseMessage);
 
@@ -128,61 +125,6 @@ export function HypergraphAppProvider({
   const sessionToken = useSelectorStore(store, (state) => state.context.sessionToken);
   const keys = useSelectorStore(store, (state) => state.context.keys);
 
-  async function loginWithKeys(keys: Identity.IdentityKeys, accountId: Address, retryCount = 0) {
-    const sessionToken = Identity.loadSyncServerSessionToken(storage, accountId);
-    if (sessionToken) {
-      // 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 keys');
-        Identity.wipeSyncServerSessionToken(storage, accountId);
-        if (retryCount > 3) {
-          throw new Error('Could not login with keys after several attempts');
-        }
-        return await loginWithKeys(keys, accountId, retryCount + 1);
-      }
-      throw new Error('Could not login with keys');
-    }
-
-    const account = privateKeyToAccount(keys.signaturePrivateKey as Hex);
-    const sessionNonce = await getSessionNonce(account.address, syncServerUri);
-    const message = prepareSiweMessage(
-      account.address,
-      sessionNonce,
-      { host: window.location.host, origin: window.location.origin },
-      chainId,
-    );
-    const signature = await account.signMessage({ message });
-    const req = {
-      accountId,
-      message,
-      publicKey: keys.signaturePublicKey,
-      signature,
-    } as const satisfies Messages.RequestLoginWithSigningKey;
-    const res = await fetch(new URL('/login/with-signing-key', syncServerUri), {
-      method: 'POST',
-      headers: {
-        'Content-Type': 'application/json',
-      },
-      body: JSON.stringify(req),
-    });
-    if (res.status !== 200) {
-      throw new Error('Error logging in with signing key');
-    }
-    const decoded = Schema.decodeUnknownSync(Messages.ResponseLogin)(await res.json());
-    Identity.storeAccountId(storage, accountId);
-    Identity.storeSyncServerSessionToken(storage, accountId, decoded.sessionToken);
-    return {
-      accountId,
-      sessionToken: decoded.sessionToken,
-      keys,
-    };
-  }
-
   async function login(signer: Identity.Signer) {
     if (!signer) {
       return;
@@ -202,10 +144,10 @@ export function HypergraphAppProvider({
       host: window.location.host,
       origin: window.location.origin,
     };
-    if (!keys && !(await identityExists(accountId, syncServerUri))) {
+    if (!keys && !(await Identity.identityExists(accountId, syncServerUri))) {
       authData = await Identity.signup(signer, accountId, syncServerUri, chainId, storage, location);
     } else if (keys) {
-      authData = await loginWithKeys(keys, accountId);
+      authData = await Identity.loginWithKeys(keys, accountId, syncServerUri, chainId, storage, location);
     } else {
       authData = await Identity.loginWithWallet(signer, accountId, syncServerUri, chainId, storage, location);
     }

packages/hypergraph/src/identity/login.ts

@@ -14,7 +14,7 @@ import {
 import { createIdentityKeys } from './create-identity-keys.js';
 import { decryptIdentity, encryptIdentity } from './identity-encryption.js';
 import { proveIdentityOwnership } from './prove-ownership.js';
-import type { Signer, Storage } from './types.js';
+import type { IdentityKeys, Signer, Storage } from './types.js';
 
 export function prepareSiweMessage(
   address: Address,
@@ -184,3 +184,65 @@ export async function loginWithWallet(
     keys,
   };
 }
+
+export async function loginWithKeys(
+  keys: IdentityKeys,
+  accountId: Address,
+  syncServerUri: string,
+  chainId: number,
+  storage: Storage,
+  location: { host: string; origin: string },
+  retryCount = 0,
+) {
+  const sessionToken = loadSyncServerSessionToken(storage, accountId);
+  if (sessionToken) {
+    // 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 keys');
+      wipeSyncServerSessionToken(storage, accountId);
+      if (retryCount > 3) {
+        throw new Error('Could not login with keys after several attempts');
+      }
+      return await loginWithKeys(keys, accountId, syncServerUri, chainId, storage, location, retryCount + 1);
+    }
+    return {
+      accountId,
+      sessionToken,
+      keys,
+    };
+  }
+
+  const account = privateKeyToAccount(keys.signaturePrivateKey as Hex);
+  const sessionNonce = await getSessionNonce(accountId, syncServerUri);
+  const message = prepareSiweMessage(account.address, sessionNonce, location, chainId);
+  const signature = await account.signMessage({ message });
+  const req = {
+    accountId,
+    message,
+    publicKey: keys.signaturePublicKey,
+    signature,
+  } as const satisfies Messages.RequestLoginWithSigningKey;
+  const res = await fetch(new URL('/login/with-signing-key', syncServerUri), {
+    method: 'POST',
+    headers: {
+      'Content-Type': 'application/json',
+    },
+    body: JSON.stringify(req),
+  });
+  if (res.status !== 200) {
+    throw new Error('Error logging in with signing key');
+  }
+  const decoded = Schema.decodeUnknownSync(Messages.ResponseLogin)(await res.json());
+  storeAccountId(storage, accountId);
+  storeSyncServerSessionToken(storage, accountId, decoded.sessionToken);
+  return {
+    accountId,
+    sessionToken: decoded.sessionToken,
+    keys,
+  };
+}

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

@@ -1,7 +1,14 @@
 import { privateKeyToAccount } from 'viem/accounts';
 import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
 import { encryptIdentity } from '../../src/identity/identity-encryption';
-import { getSessionNonce, loginWithWallet, prepareSiweMessage, restoreKeys, signup } from '../../src/identity/login';
+import {
+  getSessionNonce,
+  loginWithKeys,
+  loginWithWallet,
+  prepareSiweMessage,
+  restoreKeys,
+  signup,
+} from '../../src/identity/login';
 import type { IdentityKeys, Signer, Storage } from '../../src/identity/types';
 import type * as Messages from '../../src/messages';
 
@@ -446,3 +453,140 @@ describe('loginWithWallet', () => {
     expect(result.keys).toEqual(identityKeys);
   });
 });
+
+describe('loginWithKeys', () => {
+  let mockStorage: Storage;
+  const accountId = '0x742d35Cc6634C0532925a3b844Bc454e4438f44e';
+  const identityKeys: IdentityKeys = {
+    signaturePublicKey: '0x0262701b2eb1b6b37ad03e24445dfcad1b91309199e43017b657ce2604417c12f5',
+    signaturePrivateKey: '0x88bb6f20de8dc1787c722dc847f4cf3d00285b8955445f23c483d1237fe85366',
+    encryptionPrivateKey: '0xbbf164a93b0f78a85346017fa2673cf367c64d81b1c3d6af7ad45e308107a812',
+    encryptionPublicKey: '0x595e1a6b0bb346d83bc382998943d2e6d9210fd341bc8b9f41a7229eede27240',
+  };
+  const location = { host: 'localhost', origin: 'http://localhost' };
+
+  beforeEach(() => {
+    mockStorage = {
+      getItem: vi.fn(),
+      setItem: vi.fn(),
+      removeItem: vi.fn(),
+    };
+  });
+
+  afterEach(() => {
+    vi.restoreAllMocks();
+  });
+
+  it('should successfully login with keys and validate the valid session token', async () => {
+    const sessionToken = 'test-session-token';
+
+    // @ts-expect-error
+    mockStorage.getItem.mockImplementation((key) => {
+      if (key === `hypergraph:dev:session-token:${accountId}`) {
+        return sessionToken;
+      }
+
+      return null;
+    });
+
+    vi.spyOn(global, 'fetch').mockImplementation((url) => {
+      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 loginWithKeys(identityKeys, 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 login with keys and retrieve a new session token', async () => {
+    const sessionToken = 'test-session-token';
+
+    vi.spyOn(global, 'fetch').mockImplementation((url) => {
+      if (url.toString() === 'http://localhost:3000/login/with-signing-key') {
+        return Promise.resolve({
+          status: 200,
+          json: () => Promise.resolve({ sessionToken }),
+        } as Response);
+      }
+      if (url.toString() === 'http://localhost:3000/login/nonce') {
+        return Promise.resolve({
+          status: 200,
+          json: () => Promise.resolve({ sessionNonce: 'Sv2dJppgx9SKDGCIb' }),
+        } as Response);
+      }
+
+      return vi.fn() as never;
+    });
+
+    const result = await loginWithKeys(identityKeys, accountId, 'http://localhost:3000', 1, mockStorage, location);
+
+    expect(result.accountId).toBe(accountId);
+    expect(result.sessionToken).toBe(sessionToken);
+    expect(result.keys).toEqual(identityKeys);
+  });
+
+  it('should get new session token if existing token is invalid', async () => {
+    const sessionToken = 'test-session-token';
+
+    const storageData = {};
+    const storage: Storage = {
+      getItem: vi.fn((key) => storageData[key] || null),
+      setItem: vi.fn((key, value) => {
+        storageData[key] = value;
+      }),
+      removeItem: vi.fn((key) => {
+        delete storageData[key];
+      }),
+    };
+
+    storage.setItem(`hypergraph:dev:session-token:${accountId}`, 'wrong-session-token');
+
+    let whoamiCounter = 0;
+
+    vi.spyOn(global, 'fetch').mockImplementation((url) => {
+      if (url.toString() === 'http://localhost:3000/whoami') {
+        return Promise.resolve({
+          status: 200,
+          text: () => {
+            whoamiCounter++;
+            if (whoamiCounter === 1) {
+              return Promise.resolve('wrong-account-id');
+            }
+            return Promise.resolve(accountId);
+          },
+        } as Response);
+      }
+      if (url.toString() === 'http://localhost:3000/login/with-signing-key') {
+        return Promise.resolve({
+          status: 200,
+          json: () => Promise.resolve({ sessionToken }),
+        } as Response);
+      }
+      if (url.toString() === 'http://localhost:3000/login/nonce') {
+        return Promise.resolve({
+          status: 200,
+          json: () => Promise.resolve({ sessionNonce: 'Sv2dJppgx9SKDGCIb' }),
+        } as Response);
+      }
+
+      return vi.fn() as never;
+    });
+
+    const result = await loginWithKeys(identityKeys, accountId, 'http://localhost:3000', 1, storage, location);
+
+    expect(result.accountId).toBe(accountId);
+    expect(result.sessionToken).toBe(sessionToken);
+    expect(result.keys).toEqual(identityKeys);
+    expect(storage.removeItem).toHaveBeenCalledWith(`hypergraph:dev:session-token:${accountId}`);
+    expect(storage.setItem).toHaveBeenCalledWith(`hypergraph:dev:session-token:${accountId}`, sessionToken);
+  });
+});