From 88565e6da4815f08dd6a8fb01af20f4558451059 Mon Sep 17 00:00:00 2001
From: Eugene Kabanov <eugene.kabanov@status.im>
Date: Thu, 21 Sep 2023 21:07:02 +0300
Subject: [PATCH] Make KeystoreCache working with KeystoreManagement API.
 (#5443)

---
 beacon_chain/nimbus_beacon_node.nim            |  4 +++-
 beacon_chain/nimbus_validator_client.nim       |  7 ++++---
 beacon_chain/rpc/rest_key_management_api.nim   |  3 ++-
 .../validators/keystore_management.nim         | 18 ++++++++----------
 4 files changed, 17 insertions(+), 15 deletions(-)

diff --git a/beacon_chain/nimbus_beacon_node.nim b/beacon_chain/nimbus_beacon_node.nim
index 4041dfa0f..d95ae4959 100644
--- a/beacon_chain/nimbus_beacon_node.nim
+++ b/beacon_chain/nimbus_beacon_node.nim
@@ -700,6 +700,7 @@ proc init*(T: type BeaconNode,
     getStateField(dag.headState, genesis_validators_root)
 
   let
+    keystoreCache = KeystoreCacheRef.init()
     slashingProtectionDB =
       SlashingProtectionDB.init(
           getStateField(dag.headState, genesis_validators_root),
@@ -711,6 +712,7 @@ proc init*(T: type BeaconNode,
     keymanagerHost = if keymanagerInitResult.server != nil:
       newClone KeymanagerHost.init(
         validatorPool,
+        keystoreCache,
         rng,
         keymanagerInitResult.token,
         config.validatorsDir,
@@ -749,7 +751,7 @@ proc init*(T: type BeaconNode,
     restServer: restServer,
     keymanagerHost: keymanagerHost,
     keymanagerServer: keymanagerInitResult.server,
-    keystoreCache: KeystoreCacheRef.init(),
+    keystoreCache: keystoreCache,
     eventBus: eventBus,
     gossipState: {},
     blocksGossipState: {},
diff --git a/beacon_chain/nimbus_validator_client.nim b/beacon_chain/nimbus_validator_client.nim
index 3149bf733..7e5da53d8 100644
--- a/beacon_chain/nimbus_validator_client.nim
+++ b/beacon_chain/nimbus_validator_client.nim
@@ -364,9 +364,10 @@ proc asyncInit(vc: ValidatorClientRef): Future[ValidatorClientRef] {.async.} =
     vc.blockService = await BlockServiceRef.init(vc)
     vc.syncCommitteeService = await SyncCommitteeServiceRef.init(vc)
     vc.keymanagerServer = keymanagerInitResult.server
-    if vc.keymanagerServer != nil:
+    if not(isNil(vc.keymanagerServer)):
       vc.keymanagerHost = newClone KeymanagerHost.init(
         validatorPool,
+        vc.keystoreCache,
         vc.rng,
         keymanagerInitResult.token,
         vc.config.validatorsDir,
@@ -458,8 +459,8 @@ proc asyncRun*(vc: ValidatorClientRef) {.async.} =
   vc.blockService.start()
   vc.syncCommitteeService.start()
 
-  if not isNil(vc.keymanagerServer):
-    doAssert vc.keymanagerHost != nil
+  if not(isNil(vc.keymanagerServer)):
+    doAssert not(isNil(vc.keymanagerHost))
     vc.keymanagerServer.router.installKeymanagerHandlers(vc.keymanagerHost[])
     vc.keymanagerServer.start()
 
diff --git a/beacon_chain/rpc/rest_key_management_api.nim b/beacon_chain/rpc/rest_key_management_api.nim
index b6f9f5b99..90ec94005 100644
--- a/beacon_chain/rpc/rest_key_management_api.nim
+++ b/beacon_chain/rpc/rest_key_management_api.nim
@@ -185,7 +185,8 @@ proc installKeymanagerHandlers*(router: var RestRouter, host: KeymanagerHost) =
     for index, item in request.keystores:
       let res = importKeystore(host.validatorPool[], host.rng[],
                                host.validatorsDir, host.secretsDir,
-                               item, request.passwords[index])
+                               item, request.passwords[index],
+                               host.keystoreCache)
       if res.isErr():
         let failure = res.error()
         case failure.status
diff --git a/beacon_chain/validators/keystore_management.nim b/beacon_chain/validators/keystore_management.nim
index 1696db3e6..4673c3d8a 100644
--- a/beacon_chain/validators/keystore_management.nim
+++ b/beacon_chain/validators/keystore_management.nim
@@ -80,6 +80,7 @@ type
 
   KeymanagerHost* = object
     validatorPool*: ref ValidatorPool
+    keystoreCache*: KeystoreCacheRef
     rng*: ref HmacDrbgContext
     keymanagerToken*: string
     validatorsDir*: string
@@ -111,6 +112,7 @@ func dispose*(decryptor: var MultipleKeystoresDecryptor) =
 
 func init*(T: type KeymanagerHost,
            validatorPool: ref ValidatorPool,
+           keystoreCache: KeystoreCacheRef,
            rng: ref HmacDrbgContext,
            keymanagerToken: string,
            validatorsDir: string,
@@ -123,6 +125,7 @@ func init*(T: type KeymanagerHost,
            getForkFn: GetForkFn,
            getGenesisFn: GetGenesisFn): T =
   T(validatorPool: validatorPool,
+    keystoreCache: keystoreCache,
     rng: rng,
     keymanagerToken: keymanagerToken,
     validatorsDir: validatorsDir,
@@ -1339,18 +1342,13 @@ proc importKeystore*(pool: var ValidatorPool,
                      rng: var HmacDrbgContext,
                      validatorsDir, secretsDir: string,
                      keystore: Keystore,
-                     password: string): ImportResult[KeystoreData] {.
+                     password: string,
+                     cache: KeystoreCacheRef): ImportResult[KeystoreData] {.
      raises: [].} =
-  let keypass = KeystorePass.init(password)
-  let privateKey =
-    block:
-      let res = decryptKeystore(keystore, keypass)
-      if res.isOk():
-        res.get()
-      else:
-        return err(
-          AddValidatorFailure.init(AddValidatorStatus.failed, res.error()))
   let
+    keypass = KeystorePass.init(password)
+    privateKey = decryptKeystore(keystore, keypass, cache).valueOr:
+      return err(AddValidatorFailure.init(AddValidatorStatus.failed, error))
     publicKey = privateKey.toPubKey()
     keyName = publicKey.fsName
     keystoreDir = validatorsDir / keyName