Refactored secret storage to support user maps (#32)

src/controllers/UserLoginController.ts

@@ -6,7 +6,7 @@ import {
   GENERATE_DHE_KEY_PAIR_CMD,
   REQUEST_DHE_USER_CREDENTIALS_CMD,
 } from '../common';
-import { generateKeyPairForUser, Logger } from '../util';
+import { formatDHPublicKey, generateBase64KeyPair, Logger } from '../util';
 import type { ServerState } from '../types';
 
 const logger = new Logger('UserLoginController');
@@ -52,19 +52,24 @@ export class UserLoginController extends ControllerBase {
       return;
     }
 
-    const { publicKey, privateKey } = generateKeyPairForUser(
-      credentials.username
-    );
+    const [publicKey, privateKey] = generateBase64KeyPair();
 
     // TODO: Challenge response send public key to server
-    logger.debug('Public key:', publicKey);
-
-    await this.secretService.storePrivateKey(
-      credentials.username,
-      serverUrl,
-      privateKey
+    logger.debug(
+      'Public key:',
+      formatDHPublicKey(credentials.username, publicKey)
     );
 
+    // Get existing server keys or create a new object
+    const serverKeys =
+      (await this.secretService.getServerKeys(serverUrl)) ?? {};
+
+    // Store the new private key for the user
+    await this.secretService.storeServerKeys(serverUrl, {
+      ...serverKeys,
+      [credentials.username]: privateKey,
+    });
+
     // Remove credentials from cache since presumably a valid key pair was
     // generated and we'll want the user to authenticate with the private key
     // instead.

src/services/SecretService.ts

@@ -1,9 +1,13 @@
 import type { SecretStorage } from 'vscode';
-import type { DHPrivateKey } from '../types';
+import type { OperateAsUserStored, ServerSecretKeysStored } from '../types';
 
 class Key {
-  static privateKey(userName: string, serverUrl: URL): string {
-    return `privateKey.${userName}.${serverUrl.toString()}`;
+  static operateAsUser(serverUrl: URL): string {
+    return `operateAsUser.${serverUrl.toString()}`;
+  }
+
+  static serverKeys(serverUrl: URL): string {
+    return `serverKeys.${serverUrl.toString()}`;
   }
 }
 
@@ -20,20 +24,79 @@ export class SecretService {
 
   private readonly _secrets: SecretStorage;
 
-  getPrivateKey = async (
-    userName: string,
+  /**
+   * Parse a stored JSON string value to an object.
+   * @param key Secret storage key
+   * @returns An object of type T or null if a value cannot be found or parsed.
+   */
+  _getJson = async <T>(key: string): Promise<T | null> => {
+    const raw = await this._secrets.get(key);
+    if (raw == null) {
+      return null;
+    }
+
+    try {
+      return JSON.parse(raw);
+    } catch {
+      await this._secrets.delete(key);
+      return null;
+    }
+  };
+
+  /**
+   * Store a JSON-serializable value.
+   * @param key Secret storage key
+   * @param value Value to store
+   */
+  _storeJson = async <T>(key: string, value: T): Promise<void> => {
+    return this._secrets.store(key, JSON.stringify(value));
+  };
+
+  /**
+   * Get a map of user -> operatas users for a given server.
+   * @param serverUrl The server URL to get the map for.
+   * @returns The map of user -> operate as user or null.
+   */
+  getOperateAsUser = async (
+    serverUrl: URL
+  ): Promise<OperateAsUserStored | null> => {
+    return this._getJson(Key.operateAsUser(serverUrl));
+  };
+
+  /**
+   * Store a map of user -> operate as user for a given server.
+   * @param serverUrl The server URL to store the map for.
+   * @param operateAsUser The map of user -> operate as user.
+   */
+  storeOperateAsUser = async (
+    serverUrl: URL,
+    operateAsUser: string
+  ): Promise<void> => {
+    const key = Key.operateAsUser(serverUrl);
+    await this._storeJson(key, operateAsUser);
+  };
+
+  /**
+   * Get a map of user -> private keys for a given server.
+   * @param serverUrl
+   * @returns The map of user -> private key or null.
+   */
+  getServerKeys = async (
     serverUrl: URL
-  ): Promise<DHPrivateKey | undefined> => {
-    return this._secrets.get(Key.privateKey(userName, serverUrl)) as Promise<
-      DHPrivateKey | undefined
-    >;
+  ): Promise<ServerSecretKeysStored | null> => {
+    return this._getJson(Key.serverKeys(serverUrl));
   };
 
-  storePrivateKey = async (
-    userName: string,
+  /**
+   * Store a map of user -> private keys for a given server.
+   * @param serverUrl The server URL to store the map for.
+   * @param serverKeys The map of user -> private key.
+   */
+  storeServerKeys = async (
     serverUrl: URL,
-    privateKey: DHPrivateKey
+    serverKeys: ServerSecretKeysStored
   ): Promise<void> => {
-    await this._secrets.store(Key.privateKey(userName, serverUrl), privateKey);
+    const key = Key.serverKeys(serverUrl);
+    await this._storeJson(key, serverKeys);
   };
 }

src/types/commonTypes.d.ts

@@ -37,12 +37,12 @@ export interface EnterpriseConnectionConfig {
   experimentalWorkerConfig?: WorkerConfig;
 }
 
+export type Base64PrivateKey = Brand<'Base64PrivateKey', string>;
+export type Base64PublicKey = Brand<'Base64PublicKey', string>;
 export type DHPrivateKey = Brand<'DHPrivateKey', string>;
 export type DHPublicKey = Brand<'DHPublicKey', string>;
-export type DHKeyPair = {
-  privateKey: DHPrivateKey;
-  publicKey: DHPublicKey;
-};
+export type OperateAsUserStored = Record<string, string>;
+export type ServerSecretKeysStored = Record<string, Base64PrivateKey>;
 
 export type ServerConnectionConfig =
   | CoreConnectionConfig

src/util/authUtils.ts

@@ -1,38 +1,47 @@
 import { generateKeyPairSync } from 'node:crypto';
-import type { DHKeyPair, DHPrivateKey, DHPublicKey } from '../types';
+import type {
+  Base64PrivateKey,
+  Base64PublicKey,
+  DHPrivateKey,
+  DHPublicKey,
+} from '../types';
 
 // synonymous with secp256r1
 const NAMED_CURVE = 'prime256v1' as const;
 
-/**
- * Generate a new keypair in Deephaven format.
- * @param userName The userName to generate the keypair for.
- */
-export function generateKeyPairForUser<TUser extends string>(
-  userName: TUser
-): DHKeyPair {
+export function generateBase64KeyPair(): [Base64PublicKey, Base64PrivateKey] {
   const { publicKey: publicKeyBuffer, privateKey: privateKeyBuffer } =
     generateKeyPairSync('ec', {
       namedCurve: NAMED_CURVE,
       publicKeyEncoding: { type: 'spki', format: 'der' },
       privateKeyEncoding: { type: 'pkcs8', format: 'der' },
     });
 
-  const publicBase64Str = toEcBase64String(publicKeyBuffer);
-  const privateBase64Str = toEcBase64String(privateKeyBuffer);
+  const publicKey = toEcBase64String(publicKeyBuffer) as Base64PublicKey;
+  const privateKey = toEcBase64String(privateKeyBuffer) as Base64PrivateKey;
+
+  return [publicKey, privateKey];
+}
 
-  const publicKey = `${userName} ${publicBase64Str}` as DHPublicKey;
-  const privateKey = [
+export function formatDHPublicKey(
+  userName: string,
+  base64PublicKey: Base64PublicKey
+): string {
+  return `${userName} ${base64PublicKey}` as DHPublicKey;
+}
+
+export function formatDHPrivateKey(
+  userName: string,
+  operateAs: string,
+  base64PublicKey: Base64PublicKey,
+  base64PrivateKey: Base64PrivateKey
+): string {
+  return [
     `user ${userName}`,
-    `operateas ${userName}`,
-    `public ${publicBase64Str}`,
-    `private ${privateBase64Str}`,
+    `operateas ${operateAs}`,
+    `public ${base64PublicKey}`,
+    `private ${base64PrivateKey}`,
   ].join('\n') as DHPrivateKey;
-
-  return {
-    publicKey,
-    privateKey,
-  };
 }
 
 /**