perf(hashicorp-vault-signing-manager): cache keys to avoid repeated l…

…ookups (#29)

* perf(hashicorp-vault-signing-manager): cache keys to avoid repeated lookups

packages/hashicorp-vault-signing-manager/src/lib/hashicorp-vault-signing-manager.spec.ts

@@ -151,6 +151,20 @@ describe('class VaultSigner', () => {
       expect(result.signature).toBe(expectedSignature);
     });
 
+    it('should clear the cache if Vault returns an error', async () => {
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      const exposedSigner = signer as any;
+      const clearAddressCacheSpy = jest.spyOn(exposedSigner, 'clearAddressCache');
+      const payload = { address: accounts[0].address } as SignerPayloadJSON;
+
+      const expectedError = new Error('fake error');
+      mockHashicorpVault.signData.mockRejectedValue(expectedError);
+
+      await expect(signer.signPayload(payload)).rejects.toThrow(expectedError);
+
+      return expect(clearAddressCacheSpy).toHaveBeenCalled();
+    });
+
     it('should throw an error if the payload address is not present in the Vault', () => {
       mockHashicorpVault.fetchAllKeys.mockResolvedValue([]);
       return expect(
@@ -182,6 +196,48 @@ describe('class VaultSigner', () => {
       expect(result.signature).toBe(expectedSignature);
     });
 
+    it('should check and set the key cache', async () => {
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      const exposedSigner = signer as any;
+      const expectedAccount = accounts[0];
+      const getSpy = jest.spyOn(exposedSigner, 'getCachedKey');
+      const setSpy = jest.spyOn(exposedSigner, 'setCachedKey');
+
+      const { address, ...expectedKey } = expectedAccount;
+
+      const data = u8aToHex(stringToU8a('Hello, my name is Alice'));
+      const raw = {
+        address,
+        data,
+        type: 'bytes' as const,
+      };
+
+      await signer.signRaw(raw);
+
+      expect(getSpy).toHaveBeenCalledWith(address);
+      expect(setSpy).toHaveBeenCalledWith(address, expectedKey);
+    });
+
+    it('should clear the cache if Vault returns an error', async () => {
+      const data = u8aToHex(stringToU8a('Hello, my name is Alice'));
+      const raw = {
+        address: accounts[0].address,
+        data,
+        type: 'bytes' as const,
+      };
+
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      const exposedSigner = signer as any;
+      const clearAddressCacheSpy = jest.spyOn(exposedSigner, 'clearAddressCache');
+
+      const expectedError = new Error('fake error');
+      mockHashicorpVault.signData.mockRejectedValue(expectedError);
+
+      await expect(signer.signRaw(raw)).rejects.toThrow(expectedError);
+
+      return expect(clearAddressCacheSpy).toHaveBeenCalled();
+    });
+
     it('should throw an error if the payload address is not present in the Vault', () => {
       mockHashicorpVault.fetchAllKeys.mockResolvedValue([]);
       return expect(

packages/hashicorp-vault-signing-manager/src/lib/hashicorp-vault-signing-manager.ts

@@ -9,6 +9,10 @@ import { HashicorpVault, VaultKey } from './hashicorp-vault';
 
 export class VaultSigner implements PolkadotSigner {
   private currentId = -1;
+  /**
+   * This can be invalid in the event of a key rename in Vault
+   */
+  private addressCache: Record<string, VaultKey> = {};
 
   /**
    * @hidden
@@ -30,7 +34,10 @@ export class VaultSigner implements PolkadotSigner {
       version,
     });
 
-    return this.signData(name, keyVersion, signablePayload.toU8a(true));
+    return this.signData(name, keyVersion, signablePayload.toU8a(true)).catch(error => {
+      this.clearAddressCache(address);
+      throw error;
+    });
   }
 
   /**
@@ -41,7 +48,10 @@ export class VaultSigner implements PolkadotSigner {
 
     const { name, version } = await this.getVaultKey(address);
 
-    return this.signData(name, version, hexToU8a(data));
+    return this.signData(name, version, hexToU8a(data)).catch(error => {
+      this.clearAddressCache(address);
+      throw error;
+    });
   }
 
   /**
@@ -78,6 +88,11 @@ export class VaultSigner implements PolkadotSigner {
    * @throws if there is no key with that address
    */
   private async getVaultKey(address: string): Promise<VaultKey> {
+    const cachedKey = this.getCachedKey(address);
+    if (cachedKey) {
+      return cachedKey;
+    }
+
     const payloadPublicKey = u8aToHex(decodeAddress(address));
 
     const allKeys = await this.vault.fetchAllKeys();
@@ -88,8 +103,22 @@ export class VaultSigner implements PolkadotSigner {
       throw new Error('The signer cannot sign transactions on behalf of the calling Account');
     }
 
+    this.setCachedKey(address, foundKey);
+
     return foundKey;
   }
+
+  private getCachedKey(address: string): VaultKey {
+    return this.addressCache[address];
+  }
+
+  private setCachedKey(address: string, key: VaultKey): void {
+    this.addressCache[address] = key;
+  }
+
+  private clearAddressCache(address: string): void {
+    delete this.addressCache[address];
+  }
 }
 export class HashicorpVaultSigningManager implements SigningManager {
   private externalSigner: VaultSigner;