fix: added cache

package.json

@@ -57,7 +57,8 @@
     "rimraf": "^3.0.2",
     "rxjs": "^7.2.0",
     "winston": "^3.3.3",
-    "ws": "^8.10.0"
+    "ws": "^8.10.0",
+    "lru-cache": "^9.1.1"
   },
   "devDependencies": {
     "@nestjs/cli": "^8.2.5",

src/guardian/keys-validation/constants.ts

@@ -0,0 +1,2 @@
+// TODO: put in config
+export const KEYS_LRU_CACHE_SIZE = 32000;

src/guardian/keys-validation/keys-validation.service.ts

@@ -3,56 +3,121 @@ import { WINSTON_MODULE_NEST_PROVIDER } from 'nest-winston';
 import {
   KeyValidatorInterface,
   bufferFromHexString,
+  Pubkey,
+  WithdrawalCredentialsBuffer,
+  Key,
 } from '@lido-nestjs/key-validation';
 import { RegistryKey } from 'keys-api/interfaces/RegistryKey';
-import { ProviderService } from 'provider';
 import { GENESIS_FORK_VERSION_BY_CHAIN_ID } from 'bls/bls.constants';
+import { LRUCache } from 'lru-cache';
+import { KEYS_LRU_CACHE_SIZE } from './constants';
+import { ProviderService } from 'provider';
+
+type DepositData = {
+  key: Pubkey;
+  depositSignature: string;
+  withdrawalCredentials: WithdrawalCredentialsBuffer;
+  genesisForkVersion: Buffer;
+};
 
 @Injectable()
 export class KeysValidationService {
+  private keysCache: LRUCache<string, { signature: string; isValid: boolean }>;
+
   constructor(
     @Inject(WINSTON_MODULE_NEST_PROVIDER)
     protected readonly logger: LoggerService,
     private readonly keyValidator: KeyValidatorInterface,
     private readonly provider: ProviderService,
-  ) {}
+  ) {
+    this.keysCache = new LRUCache({ max: KEYS_LRU_CACHE_SIZE });
+  }
 
   /**
    *
    * Return list of invalid keys
    */
-  async validateKeys(
+  public async findInvalidKeys(
     vettedKeys: RegistryKey[],
     withdrawalCredentials: string,
   ): Promise<{ key: string; depositSignature: string }[]> {
     const forkVersion: Uint8Array = await this.forkVersion();
 
-    const validatedKeys: [
-      {
-        key: string;
-        depositSignature: string;
-        used: boolean;
-      },
-      boolean,
-    ][] = await this.keyValidator.validateKeys(
-      vettedKeys.map((key) => ({
-        key: key.key,
-        depositSignature: key.depositSignature,
-        used: false,
-        withdrawalCredentials: bufferFromHexString(withdrawalCredentials),
-        genesisForkVersion: Buffer.from(forkVersion.buffer),
-      })),
+    const { keysNeedingValidation, unchangedAndInvalidKeys } =
+      this.divideKeys(vettedKeys);
+
+    const keysForValidation = keysNeedingValidation.map((key) =>
+      this.toDepositData(key, withdrawalCredentials, forkVersion),
     );
 
-    return validatedKeys
-      .filter(([, result]) => !result)
+    const validatedKeys: [Key & DepositData, boolean][] =
+      await this.keyValidator.validateKeys(keysForValidation);
+
+    this.updateCache(validatedKeys);
+
+    // this list will not include invalid keys from cache
+    const invalidKeysFromCurrentValidation = validatedKeys
+      .filter(([, isValid]) => !isValid)
       .map(([key]) => ({
         key: key.key,
         depositSignature: key.depositSignature,
       }));
+
+    this.logger.log('Validation keys information', {
+      vettedKeysCount: vettedKeys.length,
+      currentCacheSize: this.keysCache.size,
+      cacheInvalidKeysCount: unchangedAndInvalidKeys.length,
+      newInvalidKeys: invalidKeysFromCurrentValidation.length,
+    });
+
+    const unchangedAndInvalidKeysValues = unchangedAndInvalidKeys.map(
+      (key) => ({
+        key: key.key,
+        depositSignature: key.depositSignature,
+      }),
+    );
+
+    // merge just checked invalid keys and invalid keys from cache but only from vettedKeys
+    return [
+      ...invalidKeysFromCurrentValidation,
+      ...unchangedAndInvalidKeysValues,
+    ];
+  }
+
+  private divideKeys(vettedKeys: RegistryKey[]): {
+    keysNeedingValidation: RegistryKey[];
+    unchangedAndInvalidKeys: RegistryKey[];
+  } {
+    const keysNeedingValidation: RegistryKey[] = [];
+    const unchangedAndInvalidKeys: RegistryKey[] = [];
+
+    vettedKeys.forEach((key) => {
+      const cachedEntry = this.keysCache.get(key.key);
+
+      if (!cachedEntry || cachedEntry.signature !== key.depositSignature) {
+        keysNeedingValidation.push(key);
+      } else if (!cachedEntry.isValid) {
+        unchangedAndInvalidKeys.push(key);
+      }
+    });
+
+    return { keysNeedingValidation, unchangedAndInvalidKeys };
   }
 
-  async forkVersion(): Promise<Uint8Array> {
+  private toDepositData(
+    key: RegistryKey,
+    withdrawalCredentials: string,
+    forkVersion: Uint8Array,
+  ): DepositData {
+    return {
+      key: key.key,
+      depositSignature: key.depositSignature,
+      withdrawalCredentials: bufferFromHexString(withdrawalCredentials),
+      genesisForkVersion: Buffer.from(forkVersion.buffer),
+    };
+  }
+
+  private async forkVersion(): Promise<Uint8Array> {
     const chainId = await this.provider.getChainId();
     const forkVersion = GENESIS_FORK_VERSION_BY_CHAIN_ID[chainId];
 
@@ -62,4 +127,10 @@ export class KeysValidationService {
 
     return forkVersion;
   }
+
+  private async updateCache(validatedKeys: [Key & DepositData, boolean][]) {
+    validatedKeys.forEach(([key, isValid]) =>
+      this.keysCache.set(key.key, { signature: key.depositSignature, isValid }),
+    );
+  }
 }

src/guardian/keys-validation/keys-validation.spec.ts

@@ -0,0 +1,189 @@
+import { Test } from '@nestjs/testing';
+import { KeysValidationModule } from './keys-validation.module';
+import {
+  KeyValidatorInterface,
+  KeyValidatorModule,
+  bufferFromHexString,
+} from '@lido-nestjs/key-validation';
+import { KeysValidationService } from './keys-validation.service';
+import { LoggerModule } from 'common/logger';
+import { ConfigModule } from 'common/config';
+import { MockProviderModule } from 'provider';
+import {
+  invalidKey,
+  invalidKey2,
+  invalidKey2GoodSign,
+  validKeys,
+} from './keys.fixtures';
+import { GENESIS_FORK_VERSION_BY_CHAIN_ID } from 'bls/bls.constants';
+
+describe('KeysValidationService', () => {
+  let keysValidationService: KeysValidationService;
+  let keysValidator: KeyValidatorInterface;
+
+  let validateKeysFun: jest.SpyInstance;
+
+  const wc =
+    '0x010000000000000000000000dc62f9e8c34be08501cdef4ebde0a280f576d762';
+
+  beforeEach(async () => {
+    const moduleRef = await Test.createTestingModule({
+      imports: [
+        ConfigModule.forRoot(),
+        MockProviderModule.forRoot(),
+        LoggerModule,
+        KeyValidatorModule.forFeature({ multithreaded: true }),
+        KeysValidationModule,
+      ],
+    }).compile();
+
+    keysValidationService = moduleRef.get(KeysValidationService);
+    keysValidator = moduleRef.get(KeyValidatorInterface);
+
+    validateKeysFun = jest.spyOn(keysValidator, 'validateKeys');
+  });
+
+  it('add new key in empty cache', async () => {
+    // add a new keys
+    const result = await keysValidationService.findInvalidKeys(
+      [...validKeys, invalidKey],
+      wc,
+    );
+
+    const expected = [invalidKey].map((key) => ({
+      key: key.key,
+      depositSignature: key.depositSignature,
+    }));
+
+    const fork = GENESIS_FORK_VERSION_BY_CHAIN_ID[5];
+
+    const depositData = [...validKeys, invalidKey].map((key) => ({
+      key: key.key,
+      depositSignature: key.depositSignature,
+      withdrawalCredentials: bufferFromHexString(wc),
+      genesisForkVersion: Buffer.from(fork.buffer),
+    }));
+
+    expect(validateKeysFun).toBeCalledTimes(1);
+    expect(validateKeysFun).toBeCalledWith(depositData);
+    expect(result).toEqual(expect.arrayContaining(expected));
+    expect(result.length).toEqual(1);
+
+    validateKeysFun.mockClear();
+
+    const newResult = await keysValidationService.findInvalidKeys(
+      [...validKeys, invalidKey],
+      wc,
+    );
+
+    expect(validateKeysFun).toBeCalledTimes(1);
+    expect(validateKeysFun).toBeCalledWith([]);
+    // will be read from cache
+    expect(newResult).toEqual(expect.arrayContaining(expected));
+    expect(result.length).toEqual(1);
+  });
+
+  it('add new key in non empty cache', async () => {
+    // will include in result only invalid keys from actual list
+    // add a new keys
+    const result = await keysValidationService.findInvalidKeys(
+      [...validKeys, invalidKey, invalidKey2],
+      wc,
+    );
+
+    const expected = [invalidKey, invalidKey2].map((key) => ({
+      key: key.key,
+      depositSignature: key.depositSignature,
+    }));
+
+    const fork = GENESIS_FORK_VERSION_BY_CHAIN_ID[5];
+
+    const depositData = [...validKeys, invalidKey, invalidKey2].map((key) => ({
+      key: key.key,
+      depositSignature: key.depositSignature,
+      withdrawalCredentials: bufferFromHexString(wc),
+      genesisForkVersion: Buffer.from(fork.buffer),
+    }));
+
+    expect(validateKeysFun).toBeCalledTimes(1);
+    expect(validateKeysFun).toBeCalledWith(depositData);
+    expect(result).toEqual(expect.arrayContaining(expected));
+    expect(result.length).toEqual(expected.length);
+
+    // one invalid key was deleted
+    validateKeysFun.mockClear();
+    const newResult = await keysValidationService.findInvalidKeys(
+      [...validKeys, invalidKey],
+      wc,
+    );
+
+    const newExpected = [invalidKey].map((key) => ({
+      key: key.key,
+      depositSignature: key.depositSignature,
+    }));
+
+    expect(validateKeysFun).toBeCalledTimes(1);
+    expect(validateKeysFun).toBeCalledWith([]);
+    expect(newResult).toEqual(expect.arrayContaining(newExpected));
+    expect(newResult.length).toEqual(newExpected.length);
+  });
+
+  it('validate key again if signature was changed', async () => {
+    // if signature was changed we need to repeat validation
+    // invalid key could become valid and visa versa
+    // add a new keys
+    const result = await keysValidationService.findInvalidKeys(
+      [...validKeys, invalidKey, invalidKey2],
+      wc,
+    );
+
+    const expected = [invalidKey, invalidKey2].map((key) => ({
+      key: key.key,
+      depositSignature: key.depositSignature,
+    }));
+
+    const fork = GENESIS_FORK_VERSION_BY_CHAIN_ID[5];
+
+    const depositData = [...validKeys, invalidKey, invalidKey2].map((key) => ({
+      key: key.key,
+      depositSignature: key.depositSignature,
+      withdrawalCredentials: bufferFromHexString(wc),
+      genesisForkVersion: Buffer.from(fork.buffer),
+    }));
+
+    expect(validateKeysFun).toBeCalledTimes(1);
+    expect(validateKeysFun).toBeCalledWith(depositData);
+    expect(result).toEqual(expect.arrayContaining(expected));
+    expect(result.length).toEqual(expected.length);
+
+    // repeat and check that cache will be used
+    validateKeysFun.mockClear();
+    const newResult = await keysValidationService.findInvalidKeys(
+      [
+        ...validKeys,
+        invalidKey,
+        { ...invalidKey2, depositSignature: invalidKey2GoodSign },
+      ],
+      wc,
+    );
+
+    const newDepositData = [
+      { ...invalidKey2, depositSignature: invalidKey2GoodSign },
+    ].map((key) => ({
+      key: key.key,
+      depositSignature: key.depositSignature,
+      withdrawalCredentials: bufferFromHexString(wc),
+      genesisForkVersion: Buffer.from(fork.buffer),
+    }));
+
+    const newExpected = [invalidKey].map((key) => ({
+      key: key.key,
+      depositSignature: key.depositSignature,
+    }));
+
+    expect(validateKeysFun).toBeCalledTimes(1);
+    expect(validateKeysFun).toBeCalledWith(newDepositData);
+    expect(newResult).toEqual(expect.arrayContaining(newExpected));
+    expect(newResult.length).toEqual(newExpected.length);
+  });
+});

src/guardian/keys-validation/keys.fixtures.ts

@@ -0,0 +1,42 @@
+export const validKeys = [
+  {
+    key: '0xa9bfaa8207ee6c78644c079ffc91b6e5abcc5eede1b7a06abb8fb40e490a75ea269c178dd524b65185299d2bbd2eb7b2',
+    depositSignature:
+      '0xaa5f2a1053ba7d197495df44d4a32b7ae10265cf9e38560a16b782978c0a24271a113c9538453b7e45f35cb64c7adb460d7a9fe8c8ce6b8c80ca42fd5c48e180c73fc08f7d35ba32e39f32c902fd333faf47611827f0b7813f11c4c518dd2e59',
+    operatorIndex: 1,
+    used: false,
+    moduleAddress: '0x9D4AF1Ee19Dad8857db3a45B0374c81c8A1C6320',
+    index: 51,
+  },
+  {
+    key: '0xb3c90525010a5710d43acbea46047fc37ed55306d032527fa15dd7e8cd8a9a5fa490347cc5fce59936fb8300683cd9f3',
+    depositSignature:
+      '0x8a77d9411781360cc107344a99f6660b206d2c708ae7fa35565b76ec661a0b86b6c78f5b5691d2cf469c27d0655dfc6311451a9e0501f3c19c6f7e35a770d1a908bfec7cba2e07339dc633b8b6626216ce76ec0fa48ee56aaaf2f9dc7ccb2fe2',
+    operatorIndex: 1,
+    used: false,
+    moduleAddress: '0x9D4AF1Ee19Dad8857db3a45B0374c81c8A1C6320',
+    index: 52,
+  },
+];
+export const invalidKey = {
+  key: '0x84e85db03bee714dbecf01914460d9576b7f7226030bdbeae9ee923bf5f8e01eec4f7dfe54aa7eca6f4bccce59a0bf42',
+  depositSignature:
+    '0xb45b15f6e043d91eabbda838eae32f7dcb998578919bd813d8add67de9b14bc268a4fde41d08058a9dc2c40b881f47970c30fd3beee46517e4e5eebd4aba52060425e021302c987d365347d478681b2cabfd31208d0607f71f3766a53ca1ada0',
+  operatorIndex: 28,
+  used: false,
+  moduleAddress: '0x11a93807078f8BB880c1BD0ee4C387537de4b4b6',
+  index: 5,
+};
+
+export const invalidKey2 = {
+  key: '0x9100e67cfb22cb7f1c3924e91bc8f70111f0634fa87d3361f807585e7ab06f84a0f504b7390683ce01567e5de3ad7445',
+  depositSignature:
+    '0x8d4ed47875fab45e9cfec65bf67c956be0b00d4d4cde2b6b898b09d07eed10457b4e2a8f496077e4a145e523d5b18749035b87c2412360d4fbbc850051b307f704a758f4ef35ca4af6c5f8f4e4a95603dc688bb3773b5a22c6c21b5440c71e13',
+  operatorIndex: 1,
+  used: false,
+  moduleAddress: '0x9D4AF1Ee19Dad8857db3a45B0374c81c8A1C6320',
+  index: 54,
+};
+
+export const invalidKey2GoodSign =
+  '0x889531faa742982deab20afd3c76e4c0e4af784aed814c15ccb25fe2b77cbaaddda39dc78f364b06990972690958bae7077efa352e51c57283129598612d2ce4f3f4a4df06695d42d804ebc923a1811c80b60503b8c87e19ceee8c0bc1bb9650';