perf(sdk-core): reuse Argon2 session for account password rotation

modules/sdk-api/src/bitgoAPI.ts

@@ -20,6 +20,7 @@ import {
   GetSharingKeyOptions,
   GetSigningKeyApi,
   GlobalCoinFactory,
+  IEncryptionSession,
   IRequestTracer,
   makeRandomKey,
   sanitizeLegacyPath,
@@ -859,7 +860,7 @@ export class BitGoAPI implements BitGoBase {
    * v1: returns a shim that satisfies the same interface but runs SJCL PBKDF2 per call. Lets
    * callers that must produce v1 envelopes use the same factory as v2 callers.
    */
-  async createEncryptionSession(password: string, encryptionVersion?: EncryptionVersion) {
+  async createEncryptionSession(password: string, encryptionVersion?: EncryptionVersion): Promise<IEncryptionSession> {
     return createEncryptionSession(password, { encryptionVersion });
   }
 
@@ -2032,54 +2033,67 @@ export class BitGoAPI implements BitGoBase {
     // we just need to choose a coin that exists in the current environment
     const coin = common.Environments[this.getEnv()].network === 'bitcoin' ? 'btc' : 'tbtc';
 
-    const updateKeychainPasswordParams = { oldPassword, newPassword, encryptionVersion };
-    const v1KeychainUpdatePWResult = await this.keychains().updatePassword(updateKeychainPasswordParams);
-    const v2Keychains = await this.coin(coin).keychains().updatePassword(updateKeychainPasswordParams);
-
-    const [hmacOldPassword, hmacNewPassword] = await Promise.all([
-      this._hmacAuthStrategy.calculateHMAC(user.username, oldPassword),
-      this._hmacAuthStrategy.calculateHMAC(user.username, newPassword),
-    ]);
-
-    const updatePasswordParams = {
-      keychains: v1KeychainUpdatePWResult.keychains,
-      v2_keychains: v2Keychains,
-      version: v1KeychainUpdatePWResult.version,
-      oldPassword: hmacOldPassword,
-      password: hmacNewPassword,
-    };
-
-    // Calculate payload size in KB
-    const payloadSizeBytes = JSON.stringify(updatePasswordParams).length;
-    const payloadSizeKB = Math.ceil(payloadSizeBytes / 1024);
-
-    // Check if batching flow is enabled
+    // Argon2 is expensive, so one v2 session covers every matching keychain. V1 uses
+    // direct SJCL calls because it has no Argon2 derivation to cache.
+    const encryptionSession =
+      encryptionVersion === 2 ? await this.createEncryptionSession(newPassword, encryptionVersion) : undefined;
     try {
-      const batchingFlowCheck = await this.get(this.url('/user/checkBatchingPasswordFlow', 2))
-        .query({ payloadSize: payloadSizeKB.toString() })
-        .result();
-
-      if (batchingFlowCheck.isBatchingFlowEnabled) {
-        await this.processKeychainPasswordUpdatesInBatches(
-          updatePasswordParams.keychains,
-          updatePasswordParams.v2_keychains,
-          batchingFlowCheck.maxBatchSizeKB,
-          3
-        );
-        // Call changepassword API without keychains for batching flow
-        return this.post(this.url('/user/changepassword'))
-          .send({
-            version: updatePasswordParams.version,
-            oldPassword: updatePasswordParams.oldPassword,
-            password: updatePasswordParams.password,
-          })
+      const updateKeychainPasswordParams = {
+        oldPassword,
+        newPassword,
+        encryptionVersion,
+        encryptionSession,
+      };
+      const v1KeychainUpdatePWResult = await this.keychains().updatePassword(updateKeychainPasswordParams);
+      const v2Keychains = await this.coin(coin).keychains().updatePassword(updateKeychainPasswordParams);
+
+      const [hmacOldPassword, hmacNewPassword] = await Promise.all([
+        this._hmacAuthStrategy.calculateHMAC(user.username, oldPassword),
+        this._hmacAuthStrategy.calculateHMAC(user.username, newPassword),
+      ]);
+
+      const updatePasswordParams = {
+        keychains: v1KeychainUpdatePWResult.keychains,
+        v2_keychains: v2Keychains,
+        version: v1KeychainUpdatePWResult.version,
+        oldPassword: hmacOldPassword,
+        password: hmacNewPassword,
+      };
+
+      // Calculate payload size in KB
+      const payloadSizeBytes = JSON.stringify(updatePasswordParams).length;
+      const payloadSizeKB = Math.ceil(payloadSizeBytes / 1024);
+
+      // Check if batching flow is enabled
+      try {
+        const batchingFlowCheck = await this.get(this.url('/user/checkBatchingPasswordFlow', 2))
+          .query({ payloadSize: payloadSizeKB.toString() })
           .result();
+
+        if (batchingFlowCheck.isBatchingFlowEnabled) {
+          await this.processKeychainPasswordUpdatesInBatches(
+            updatePasswordParams.keychains,
+            updatePasswordParams.v2_keychains,
+            batchingFlowCheck.maxBatchSizeKB,
+            3
+          );
+          // Call changepassword API without keychains for batching flow
+          return this.post(this.url('/user/changepassword'))
+            .send({
+              version: updatePasswordParams.version,
+              oldPassword: updatePasswordParams.oldPassword,
+              password: updatePasswordParams.password,
+            })
+            .result();
+        }
+      } catch (error) {
+        // batching flow check failed
       }
-    } catch (error) {
-      // batching flow check failed
-    }
 
-    return this.post(this.url('/user/changepassword')).send(updatePasswordParams).result();
+      return this.post(this.url('/user/changepassword')).send(updatePasswordParams).result();
+    } finally {
+      encryptionSession?.destroy();
+    }
   }
 
   /**

modules/sdk-api/test/unit/bitgoAPI.ts

@@ -5,6 +5,7 @@ import { ProxyAgent } from 'proxy-agent';
 import * as sinon from 'sinon';
 import nock from 'nock';
 import type { IHmacAuthStrategy } from '@bitgo/sdk-hmac';
+import type { IEncryptionSession } from '@bitgo/sdk-core';
 
 describe('Constructor', function () {
   describe('cookiesPropagationEnabled argument', function () {
@@ -1131,6 +1132,39 @@ describe('Constructor', function () {
       sinon.assert.calledWithMatch(v1UpdatePasswordStub, { encryptionVersion: 2 });
       sinon.assert.calledWithMatch(v2UpdatePasswordStub, { encryptionVersion: 2 });
     });
+    it('shares one encryption session across keychain password updates', async function () {
+      nock(ROOT).get('/api/v2/user/checkBatchingPasswordFlow').query(true).reply(200, { isBatchingFlowEnabled: false });
+      nock(ROOT)
+        .post('/api/v1/user/changepassword', (body: unknown) => {
+          if (!body || typeof body !== 'object') {
+            return false;
+          }
+          return 'keychains' in body && 'v2_keychains' in body;
+        })
+        .reply(200, {});
+
+      const destroy = sandbox.stub();
+      const session: IEncryptionSession = {
+        encrypt: sandbox.stub().resolves('session-encrypted'),
+        decrypt: sandbox.stub().resolves('session-decrypted'),
+        destroy,
+      };
+      const createSession = sandbox.stub(bitgo, 'createEncryptionSession').resolves(session);
+
+      await bitgo.changePassword({ oldPassword: 'oldpw', newPassword: 'newpw', encryptionVersion: 2 });
+
+      sinon.assert.calledOnce(createSession);
+      sinon.assert.calledWithExactly(createSession, 'newpw', 2);
+      sinon.assert.calledWithMatch(v1UpdatePasswordStub, {
+        encryptionVersion: 2,
+        encryptionSession: session,
+      });
+      sinon.assert.calledWithMatch(v2UpdatePasswordStub, {
+        encryptionVersion: 2,
+        encryptionSession: session,
+      });
+      sinon.assert.calledOnce(destroy);
+    });
   });
 
   describe('createUserEcdhKeychain - encryptionVersion threading', function () {

modules/sdk-core/src/bitgo/keychain/iKeychains.ts

@@ -104,6 +104,8 @@ export interface UpdatePasswordOptions {
    * Sept 15 breaking-change window closes).
    */
   encryptionVersion?: EncryptionVersion;
+  /** Reuse one password-derived session across all matching keychains. */
+  encryptionSession?: IEncryptionSession;
 }
 
 export interface UpdateSingleKeychainPasswordOptions {
@@ -115,6 +117,8 @@ export interface UpdateSingleKeychainPasswordOptions {
    * Pass `2` to opt in to the Argon2id upgrade.
    */
   encryptionVersion?: EncryptionVersion;
+  /** Reuse one password-derived session for this keychain's new envelope. */
+  encryptionSession?: IEncryptionSession;
 }
 
 /**

modules/sdk-core/src/bitgo/keychain/keychains.ts

@@ -125,6 +125,7 @@ export class Keychains implements IKeychains {
             oldPassword: params.oldPassword,
             newPassword: params.newPassword,
             encryptionVersion: params.encryptionVersion,
+            encryptionSession: params.encryptionSession,
           });
           if (updatedKeychain.encryptedPrv) {
             // Both TSS and multi-user-ofc keys have multiple public keys in their key document and thus need to use objectID
@@ -208,11 +209,13 @@ export class Keychains implements IKeychains {
     const oldEncryptedPrv = params.keychain.encryptedPrv;
     try {
       const decryptedPrv = await this.bitgo.decrypt({ input: oldEncryptedPrv, password: params.oldPassword });
-      const newEncryptedPrv = await this.bitgo.encrypt({
-        input: decryptedPrv,
-        password: params.newPassword,
-        encryptionVersion: params.encryptionVersion ?? this.getEncryptionVersion(oldEncryptedPrv),
-      });
+      const newEncryptedPrv = params.encryptionSession
+        ? await params.encryptionSession.encrypt(decryptedPrv)
+        : await this.bitgo.encrypt({
+            input: decryptedPrv,
+            password: params.newPassword,
+            encryptionVersion: params.encryptionVersion ?? this.getEncryptionVersion(oldEncryptedPrv),
+          });
       return _.assign({}, params.keychain, { encryptedPrv: newEncryptedPrv });
     } catch (e) {
       // catching an error here means that the password was incorrect or, less likely, the input to decrypt is corrupted

modules/sdk-core/test/unit/bitgo/keychain/keychains.ts

@@ -1,6 +1,7 @@
 import * as sinon from 'sinon';
 import 'should';
 import { Keychains, decodeDerivableEd25519Pub } from '../../../../src';
+import type { IEncryptionSession } from '../../../../src/api';
 
 /**
  * Slot-④ root keys are generated as XLM keychains, so `create()` yields a 56-char StrKey pub.
@@ -147,4 +148,31 @@ describe('Keychains.createBackup', function () {
       body.provider!.should.equal('krs-provider');
     });
   });
+  describe('password rotation encryption session', function () {
+    it('uses the supplied session for the new encrypted private key', async function () {
+      mockBitGo.decrypt = sinon.stub().resolves('decrypted-prv');
+      const sessionEncrypt = sinon.stub().resolves('session-encrypted');
+      const session: IEncryptionSession = {
+        encrypt: sessionEncrypt,
+        decrypt: sinon.stub().resolves('decrypted-prv'),
+        destroy: sinon.stub(),
+      };
+
+      const updatedKeychain = await keychains.updateSingleKeychainPassword({
+        keychain: {
+          id: 'key-id',
+          encryptedPrv: 'legacy-encrypted-prv',
+          type: 'independent',
+        },
+        oldPassword: 'old-password',
+        newPassword: 'new-password',
+        encryptionVersion: 2,
+        encryptionSession: session,
+      });
+
+      sessionEncrypt.calledOnceWithExactly('decrypted-prv').should.equal(true);
+      mockBitGo.encrypt.called.should.equal(false);
+      updatedKeychain.encryptedPrv!.should.equal('session-encrypted');
+    });
+  });
 });