Refactored some auth utils (32)

src/controllers/UserLoginController.ts

@@ -9,23 +9,17 @@ import {
   REQUEST_DHE_USER_CREDENTIALS_CMD,
 } from '../common';
 import {
+  authWithPrivateKey,
   createAuthenticationMethodQuickPick,
-  EC_SENTINEL,
   generateBase64KeyPair,
   Logger,
   promptForOperateAs,
   promptForPassword,
   promptForUsername,
   runUserLoginWorkflow,
-  signWithPrivateKey,
+  uploadPublicKey,
 } from '../util';
-import type {
-  Base64Nonce,
-  IAsyncCacheService,
-  Lazy,
-  ServerState,
-  Username,
-} from '../types';
+import type { IAsyncCacheService, Lazy, ServerState, Username } from '../types';
 
 const logger = new Logger('UserLoginController');
 
@@ -96,57 +90,32 @@ export class UserLoginController extends ControllerBase {
       type: 'password',
     } as const satisfies DheLoginCredentials;
 
-    const [publicKey, privateKey] = generateBase64KeyPair();
-
-    const dheClient = await this.dheClientCache.get(serverUrl);
-    await dheClient.login(dheCredentials);
-
-    const { dbAclWriterHost, dbAclWriterPort } =
-      await dheClient.getServerConfigValues();
-
-    const publicKeyWithSentinel = `${EC_SENTINEL}${publicKey}`;
-
-    const body = {
-      user: dheCredentials.username,
-      encodedStr: publicKeyWithSentinel,
-      algorithm: 'EC',
-      comment: `Generated by vscode extension ${new Date().valueOf()}`,
-    };
-
-    const uploadKeyResult = await fetch(
-      `https://${dbAclWriterHost}:${dbAclWriterPort}/acl/publickey`,
-      {
-        method: 'POST',
-        headers: {
-          /* eslint-disable @typescript-eslint/naming-convention */
-          Authorization: await dheClient.createAuthToken('DbAclWriteServer'),
-          'Content-Type': 'application/json',
-          /* eslint-enable @typescript-eslint/naming-convention */
-        },
-        body: JSON.stringify(body),
-      }
+    const keyPair = generateBase64KeyPair();
+    const { type, publicKey } = keyPair;
+
+    let dheClient = await this.dheClientCache.get(serverUrl);
+
+    const uploadKeyResult = await uploadPublicKey(
+      dheClient,
+      dheCredentials,
+      publicKey,
+      type
     );
     logger.debug('uploadKeyResult:', uploadKeyResult.status);
 
-    // TODO: This needs to be moved to the DHE credentials cache as lazy loaded
-    // credentials.
-    try {
-      const { nonce }: { nonce: Base64Nonce } = await (
-        dheClient as any
-      ).getChallengeNonce();
+    // TODO: Need to move the login logic to lazy credentials call
 
-      const signedNonce = signWithPrivateKey(nonce, privateKey);
+    // Have to use a new client to login with the private key
+    this.dheClientCache.invalidate(serverUrl);
+    dheClient = await this.dheClientCache.get(serverUrl);
+
+    await authWithPrivateKey({
+      dheClient,
+      keyPair,
+      username,
+      operateAs: username,
+    });
 
-      const authResult = await (dheClient as any).challengeResponse(
-        signedNonce,
-        publicKeyWithSentinel,
-        dheCredentials.username,
-        dheCredentials.username
-      );
-      console.log('authResult:', authResult);
-    } catch (e) {
-      console.error(e);
-    }
     return;
 
     // TODO: Need to store public key + algorithm in the server keys
@@ -157,7 +126,7 @@ export class UserLoginController extends ControllerBase {
     // Store the new private key for the user
     await this.secretService.storeServerKeys(serverUrl, {
       ...serverKeys,
-      [dheCredentials.username]: privateKey,
+      [dheCredentials.username]: keyPair,
     });
 
     // Remove credentials from cache since presumably a valid key pair was

src/types/commonTypes.d.ts

@@ -45,9 +45,13 @@ export type Base64PrivateKey = Brand<'Base64PrivateKey', string>;
 export type Base64PublicKey = Brand<'Base64PublicKey', string>;
 export type Base64Nonce = Brand<'Base64Nonce', string>;
 export type Base64Signature = Brand<'Base64Signature', string>;
-export type DHPrivateKey = Brand<'DHPrivateKey', string>;
-export type DHPublicKey = Brand<'DHPublicKey', string>;
-export type ServerSecretKeys = Record<string, Base64PrivateKey>;
+export type KeyPairType = 'ec';
+export type Base64KeyPair = {
+  type: KeyPairType;
+  publicKey: Base64PublicKey;
+  privateKey: Base64PrivateKey;
+};
+export type ServerSecretKeys = Record<string, Base64KeyPair>;
 export type UserLoginPreferences = {
   lastLogin?: Username;
   operateAsUser: Record<Username, OperateAsUsername>;

src/util/authUtils.ts

@@ -1,18 +1,19 @@
 import { generateKeyPairSync, sign } from 'node:crypto';
 import type {
+  EnterpriseClient,
+  LoginCredentials as DheLoginCredentials,
+} from '@deephaven-enterprise/jsapi-types';
+import type {
+  Base64KeyPair,
   Base64Nonce,
   Base64PrivateKey,
   Base64PublicKey,
   Base64Signature,
-  DHPrivateKey,
-  DHPublicKey,
+  KeyPairType,
 } from '../types';
+import { Logger } from './Logger';
 
-/*
- * Base64 encoded value of 'EC:'. Used to identify that a key is an EC key when
- * passing to DH server.
- */
-export const EC_SENTINEL = 'RUM6' as const;
+const logger = new Logger('UserLoginController');
 
 /*
  * Named curve to use for generating key pairs.
@@ -22,11 +23,13 @@ const NAMED_CURVE = 'prime256v1' as const;
 
 /**
  * Generate a base64 encoded asymmetric key pair using eliptic curve.
- * @returns A tuple containing the base64 encoded public and private keys.
+ * @returns The base64 encoded public and private keys.
  */
-export function generateBase64KeyPair(): [Base64PublicKey, Base64PrivateKey] {
+export function generateBase64KeyPair(): Base64KeyPair {
+  const type: KeyPairType = 'ec';
+
   const { publicKey: publicKeyBuffer, privateKey: privateKeyBuffer } =
-    generateKeyPairSync('ec', {
+    generateKeyPairSync(type, {
       namedCurve: NAMED_CURVE,
       publicKeyEncoding: { type: 'spki', format: 'der' },
       privateKeyEncoding: { type: 'pkcs8', format: 'der' },
@@ -35,28 +38,21 @@ export function generateBase64KeyPair(): [Base64PublicKey, Base64PrivateKey] {
   const publicKey = publicKeyBuffer.toString('base64') as Base64PublicKey;
   const privateKey = privateKeyBuffer.toString('base64') as Base64PrivateKey;
 
-  return [publicKey, privateKey];
+  return { type, publicKey, privateKey };
 }
 
-export function formatDHPublicKey(
-  userName: string,
-  base64PublicKey: Base64PublicKey
-): string {
-  return `${userName} ${base64PublicKey}` as DHPublicKey;
-}
+/**
+ * Prepend a sentinal value to a private key based on the given type where
+ * sentinel is the uppercase type followed by a colon.
+ * @param type Keypair type.
+ * @param key
+ * @returns
+ */
+export function keyWithSentinel(type: 'ec', key: Base64PublicKey): string {
+  const sentinelBytes = Buffer.from(`${type.toUpperCase()}:`);
+  const keyBytes = Buffer.from(key, 'base64');
 
-export function formatDHPrivateKey(
-  userName: string,
-  operateAs: string,
-  base64PublicKey: Base64PublicKey,
-  base64PrivateKey: Base64PrivateKey
-): string {
-  return [
-    `user ${userName}`,
-    `operateas ${operateAs}`,
-    `public ${base64PublicKey}`,
-    `private ${base64PrivateKey}`,
-  ].join('\n') as DHPrivateKey;
+  return Buffer.concat([sentinelBytes, keyBytes]).toString('base64');
 }
 
 /**
@@ -78,3 +74,95 @@ export function signWithPrivateKey(
     type: 'pkcs8',
   }).toString('base64') as Base64Signature;
 }
+
+// Temporary until `jaspi-types` is updated on DHE servers
+declare module '@deephaven-enterprise/jsapi-types' {
+  // eslint-disable-next-line no-unused-vars
+  interface EnterpriseClient {
+    challengeResponse: (
+      signedNonce: Base64Signature,
+      publicKeyWithSentinel: string,
+      username: string,
+      operateAs: string
+    ) => Promise<unknown>;
+    getChallengeNonce(): Promise<{
+      algorithm: 'SHA256withDSA';
+      nonce: Base64Nonce;
+    }>;
+  }
+}
+
+/**
+ * Upload a public key to a DHE server.
+ * @param dheClient
+ * @param dheCredentials
+ * @param publicKey
+ * @param type
+ * @returns The response from the server.
+ */
+export async function uploadPublicKey(
+  dheClient: EnterpriseClient,
+  dheCredentials: DheLoginCredentials,
+  publicKey: Base64PublicKey,
+  type: KeyPairType
+): Promise<Response> {
+  await dheClient.login(dheCredentials);
+
+  const { dbAclWriterHost, dbAclWriterPort } =
+    await dheClient.getServerConfigValues();
+
+  const body = {
+    user: dheCredentials.username,
+    encodedStr: keyWithSentinel(type, publicKey),
+    algorithm: type.toUpperCase(),
+    comment: `Generated by vscode extension ${new Date().valueOf()}`,
+  };
+
+  return fetch(`https://${dbAclWriterHost}:${dbAclWriterPort}/acl/publickey`, {
+    method: 'POST',
+    headers: {
+      /* eslint-disable @typescript-eslint/naming-convention */
+      Authorization: await dheClient.createAuthToken('DbAclWriteServer'),
+      'Content-Type': 'application/json',
+      /* eslint-enable @typescript-eslint/naming-convention */
+    },
+    body: JSON.stringify(body),
+  });
+}
+
+/**
+ * Authenticate using public / private key.
+ * @param dheClient The DHE client to use.
+ * @param keyPair The base64 encoded key pair + type.
+ * @param username The username to authenticate as.
+ * @param operateAs The optional username to operate as. Defaults to `username`.
+ */
+export async function authWithPrivateKey({
+  dheClient,
+  keyPair: { type, publicKey, privateKey },
+  username,
+  operateAs = username,
+}: {
+  dheClient: EnterpriseClient;
+  keyPair: Base64KeyPair;
+  username: string;
+  operateAs?: string;
+}): Promise<void> {
+  try {
+    const { nonce } = await dheClient.getChallengeNonce();
+    const signedNonce = signWithPrivateKey(nonce, privateKey);
+
+    await dheClient.challengeResponse(
+      signedNonce,
+      keyWithSentinel(type, publicKey),
+      username,
+      operateAs
+    );
+  } catch (e) {
+    logger.error(
+      'An error occurred when signing in with public / private key',
+      e
+    );
+    throw e;
+  }
+}