[CORDA-3130] Add a cache for looking up external UUIDs from public keys (#5357)

node-api/src/main/kotlin/net/corda/nodeapi/internal/KeyOwningIdentity.kt

@@ -0,0 +1,44 @@
+package net.corda.nodeapi.internal
+
+import java.util.*
+
+/**
+ * A [KeyOwningIdentity] represents an entity that owns a public key. In this case, the "owner" refers to either an identifier provided
+ * when the key was generated, or the node itself if no identifier was supplied.
+ */
+sealed class KeyOwningIdentity {
+
+    abstract val uuid: UUID?
+
+    /**
+     * [UnmappedIdentity] is used for keys that are not assigned a UUID. This is any key created on this node that did not have a UUID
+     * assigned to it on generation. These keys are the node identity key, or confidential identity keys.
+     */
+    object UnmappedIdentity : KeyOwningIdentity() {
+        override fun toString(): String {
+            return "UNMAPPED_IDENTITY"
+        }
+
+        override val uuid: UUID? = null
+    }
+
+    /**
+     * [MappedIdentity] is used for keys that have an assigned UUID. Keys belonging to a mapped identity were assigned a UUID at the point
+     * they were generated. This UUID may refer to something outside the core of the node, for example an account.
+     */
+    data class MappedIdentity(override val uuid: UUID) : KeyOwningIdentity() {
+        override fun toString(): String {
+            return uuid.toString()
+        }
+    }
+
+    companion object {
+        fun fromUUID(uuid: UUID?): KeyOwningIdentity {
+            return if (uuid != null) {
+                MappedIdentity(uuid)
+            } else {
+                UnmappedIdentity
+            }
+        }
+    }
+}

node-api/src/main/kotlin/net/corda/nodeapi/internal/PublicKeyToOwningIdentityCache.kt

@@ -0,0 +1,17 @@
+package net.corda.nodeapi.internal
+
+import java.security.PublicKey
+
+/**
+ * A [PublicKeyToOwningIdentityCache] maps public keys to their owners. In this case, an owner could be the node identity, or it could be
+ * an external identity.
+ */
+interface PublicKeyToOwningIdentityCache {
+
+    /**
+     * Obtain the owning identity for a public key.
+     *
+     * If the key is unknown to the node, then this will return null.
+     */
+    operator fun get(key: PublicKey): KeyOwningIdentity?
+}

node/src/main/kotlin/net/corda/node/internal/AbstractNode.kt

@@ -183,6 +183,7 @@ abstract class AbstractNode<S>(val configuration: NodeConfiguration,
     @Suppress("LeakingThis")
     val networkParametersStorage = makeNetworkParametersStorage()
     val cordappProvider = CordappProviderImpl(cordappLoader, CordappConfigFileProvider(configuration.cordappDirectories), attachments).tokenize()
+    val pkToIdCache = PublicKeyToOwningIdentityCacheImpl(database, cacheFactory)
     @Suppress("LeakingThis")
     val keyManagementService = makeKeyManagementService(identityService).tokenize()
     val servicesForResolution = ServicesForResolutionImpl(identityService, attachments, cordappProvider, networkParametersStorage, transactionStorage).also {
@@ -818,7 +819,7 @@ abstract class AbstractNode<S>(val configuration: NodeConfiguration,
         // Place the long term identity key in the KMS. Eventually, this is likely going to be separated again because
         // the KMS is meant for derived temporary keys used in transactions, and we're not supposed to sign things with
         // the identity key. But the infrastructure to make that easy isn't here yet.
-        return BasicHSMKeyManagementService(cacheFactory, identityService, database, cryptoService)
+        return BasicHSMKeyManagementService(cacheFactory, identityService, database, cryptoService, pkToIdCache)
     }
 
     open fun stop() {

node/src/main/kotlin/net/corda/node/services/keys/BasicHSMKeyManagementService.kt

@@ -2,13 +2,14 @@ package net.corda.node.services.keys
 
 import net.corda.core.crypto.*
 import net.corda.core.crypto.internal.AliasPrivateKey
-import net.corda.core.identity.PartyAndCertificate
 import net.corda.core.internal.NamedCacheFactory
 import net.corda.core.serialization.SingletonSerializeAsToken
 import net.corda.core.serialization.serialize
 import net.corda.core.utilities.MAX_HASH_HEX_SIZE
 import net.corda.node.services.identity.PersistentIdentityService
+import net.corda.node.services.persistence.WritablePublicKeyToOwningIdentityCache
 import net.corda.node.utilities.AppendOnlyPersistentMap
+import net.corda.nodeapi.internal.KeyOwningIdentity
 import net.corda.nodeapi.internal.cryptoservice.SignOnlyCryptoService
 import net.corda.nodeapi.internal.persistence.CordaPersistence
 import net.corda.nodeapi.internal.persistence.NODE_DATABASE_PREFIX
@@ -29,8 +30,11 @@ import kotlin.collections.LinkedHashSet
  *
  * This class needs database transactions to be in-flight during method calls and init.
  */
-class BasicHSMKeyManagementService(cacheFactory: NamedCacheFactory, val identityService: PersistentIdentityService,
+class BasicHSMKeyManagementService(cacheFactory: NamedCacheFactory,
-                                   private val database: CordaPersistence, private val cryptoService: SignOnlyCryptoService) : SingletonSerializeAsToken(), KeyManagementServiceInternal {
+                                   override val identityService: PersistentIdentityService,
+                                   private val database: CordaPersistence,
+                                   private val cryptoService: SignOnlyCryptoService,
+                                   private val pkToIdCache: WritablePublicKeyToOwningIdentityCache) : SingletonSerializeAsToken(), KeyManagementServiceInternal {
     @Entity
     @Table(name = "${NODE_DATABASE_PREFIX}our_key_pairs")
     class PersistentKey(
@@ -93,33 +97,16 @@ class BasicHSMKeyManagementService(cacheFactory: NamedCacheFactory, val identity
         identityService.stripNotOurKeys(candidateKeys)
     }
 
-    // Unlike initial keys, freshkey() is related confidential keys and it utilises platform's software key generation
+    override fun freshKeyInternal(externalId: UUID?): PublicKey {
-    // thus, without using [cryptoService]).
-    override fun freshKey(): PublicKey {
         val keyPair = generateKeyPair()
         database.transaction {
             keysMap[keyPair.public] = keyPair.private
+            pkToIdCache[keyPair.public] = KeyOwningIdentity.fromUUID(externalId)
         }
         return keyPair.public
     }
 
-    override fun freshKey(externalId: UUID): PublicKey {
+    override fun getSigner(publicKey: PublicKey): ContentSigner {
-        val newKey = freshKey()
-        database.transaction { session.persist(PublicKeyHashToExternalId(externalId, newKey)) }
-        return newKey
-    }
-
-    override fun freshKeyAndCert(identity: PartyAndCertificate, revocationEnabled: Boolean): PartyAndCertificate {
-        return freshCertificate(identityService, freshKey(), identity, getSigner(identity.owningKey))
-    }
-
-    override fun freshKeyAndCert(identity: PartyAndCertificate, revocationEnabled: Boolean, externalId: UUID): PartyAndCertificate {
-        val newKeyWithCert = freshKeyAndCert(identity, revocationEnabled)
-        database.transaction { session.persist(PublicKeyHashToExternalId(externalId, newKeyWithCert.owningKey)) }
-        return newKeyWithCert
-    }
-
-    private fun getSigner(publicKey: PublicKey): ContentSigner {
         val signingPublicKey = getSigningPublicKey(publicKey)
         return if (signingPublicKey in originalKeysMap) {
             cryptoService.getSigner(originalKeysMap[signingPublicKey]!!)

node/src/main/kotlin/net/corda/node/services/keys/E2ETestKeyManagementService.kt

@@ -1,13 +1,12 @@
 package net.corda.node.services.keys
 
 import net.corda.core.crypto.*
-import net.corda.core.identity.PartyAndCertificate
+import net.corda.core.crypto.internal.AliasPrivateKey
 import net.corda.core.internal.ThreadBox
 import net.corda.core.node.services.IdentityService
 import net.corda.core.serialization.SingletonSerializeAsToken
-import net.corda.core.crypto.internal.AliasPrivateKey
-import net.corda.nodeapi.internal.cryptoservice.bouncycastle.BCCryptoService
 import net.corda.nodeapi.internal.cryptoservice.CryptoService
+import net.corda.nodeapi.internal.cryptoservice.bouncycastle.BCCryptoService
 import org.bouncycastle.operator.ContentSigner
 import java.security.KeyPair
 import java.security.PrivateKey
@@ -27,7 +26,7 @@ import javax.annotation.concurrent.ThreadSafe
  * etc.
  */
 @ThreadSafe
-class E2ETestKeyManagementService(val identityService: IdentityService, private val cryptoService: CryptoService? = null) : SingletonSerializeAsToken(), KeyManagementServiceInternal {
+class E2ETestKeyManagementService(override val identityService: IdentityService, private val cryptoService: CryptoService? = null) : SingletonSerializeAsToken(), KeyManagementServiceInternal {
     private class InnerState {
         val keys = HashMap<PublicKey, PrivateKey>()
     }
@@ -53,7 +52,10 @@ class E2ETestKeyManagementService(val identityService: IdentityService, private
         }
     }
 
-    override fun freshKey(): PublicKey {
+    override fun freshKeyInternal(externalId: UUID?): PublicKey {
+        if (externalId != null) {
+            throw UnsupportedOperationException("This operation is only supported by persistent key management service variants.")
+        }
         val keyPair = generateKeyPair()
         mutex.locked {
             keys[keyPair.public] = keyPair.private
@@ -61,19 +63,7 @@ class E2ETestKeyManagementService(val identityService: IdentityService, private
         return keyPair.public
     }
 
-    override fun freshKey(externalId: UUID): PublicKey {
+    override fun getSigner(publicKey: PublicKey): ContentSigner = getSigner(getSigningKeyPair(publicKey))
-        throw UnsupportedOperationException("This operation is only supported by persistent key management service variants.")
-    }
-
-    override fun freshKeyAndCert(identity: PartyAndCertificate, revocationEnabled: Boolean): PartyAndCertificate {
-        return freshCertificate(identityService, freshKey(), identity, getSigner(identity.owningKey))
-    }
-
-    override fun freshKeyAndCert(identity: PartyAndCertificate, revocationEnabled: Boolean, externalId: UUID): PartyAndCertificate {
-        throw UnsupportedOperationException("This operation is only supported by persistent key management service variants.")
-    }
-
-    private fun getSigner(publicKey: PublicKey): ContentSigner = getSigner(getSigningKeyPair(publicKey))
 
     private fun getSigningKeyPair(publicKey: PublicKey): KeyPair {
         return mutex.locked {

node/src/main/kotlin/net/corda/node/services/keys/KeyManagementServiceInternal.kt

@@ -1,32 +1,39 @@
 package net.corda.node.services.keys
 
-import net.corda.core.crypto.toStringShort
+import net.corda.core.identity.PartyAndCertificate
+import net.corda.core.node.services.IdentityService
 import net.corda.core.node.services.KeyManagementService
-import org.hibernate.annotations.Type
+import org.bouncycastle.operator.ContentSigner
 import java.security.KeyPair
 import java.security.PublicKey
 import java.util.*
-import javax.persistence.*
 
 interface KeyManagementServiceInternal : KeyManagementService {
+
+    val identityService: IdentityService
+
     fun start(initialKeyPairs: Set<KeyPair>)
+
+    fun freshKeyInternal(externalId: UUID?): PublicKey
+
+    fun getSigner(publicKey: PublicKey): ContentSigner
+
+    // Unlike initial keys, freshkey() is related confidential keys and it utilises platform's software key generation
+    // thus, without using [cryptoService]).
+    override fun freshKey(): PublicKey {
+        return freshKeyInternal(null)
+    }
+
+    override fun freshKey(externalId: UUID): PublicKey {
+        return freshKeyInternal(externalId)
+    }
+
+    override fun freshKeyAndCert(identity: PartyAndCertificate, revocationEnabled: Boolean): PartyAndCertificate {
+        return freshCertificate(identityService, freshKeyInternal(null), identity, getSigner(identity.owningKey))
+    }
+
+    override fun freshKeyAndCert(identity: PartyAndCertificate, revocationEnabled: Boolean, externalId: UUID): PartyAndCertificate {
+        return freshCertificate(identityService, freshKeyInternal(externalId), identity, getSigner(identity.owningKey))
+    }
 }
 
-@Entity
-@Table(name = "pk_hash_to_ext_id_map", indexes = [Index(name = "pk_hash_to_xid_idx", columnList = "public_key_hash")])
-class PublicKeyHashToExternalId(
-        @Id
-        @GeneratedValue
-        @Column(name = "id", unique = true, nullable = false)
-        val key: Long?,
-
-        @Column(name = "external_id", nullable = false)
-        @Type(type = "uuid-char")
-        val externalId: UUID,
-
-        @Column(name = "public_key_hash", nullable = false)
-        val publicKeyHash: String
-) {
-    constructor(accountId: UUID, publicKey: PublicKey)
-            : this(null, accountId, publicKey.toStringShort())
-}

node/src/main/kotlin/net/corda/node/services/keys/PersistentKeyManagementService.kt

@@ -1,13 +1,14 @@
 package net.corda.node.services.keys
 
 import net.corda.core.crypto.*
-import net.corda.core.identity.PartyAndCertificate
 import net.corda.core.internal.NamedCacheFactory
 import net.corda.core.internal.toSet
 import net.corda.core.serialization.SingletonSerializeAsToken
 import net.corda.core.utilities.MAX_HASH_HEX_SIZE
 import net.corda.node.services.identity.PersistentIdentityService
+import net.corda.node.services.persistence.WritablePublicKeyToOwningIdentityCache
 import net.corda.node.utilities.AppendOnlyPersistentMap
+import net.corda.nodeapi.internal.KeyOwningIdentity
 import net.corda.nodeapi.internal.persistence.CordaPersistence
 import net.corda.nodeapi.internal.persistence.NODE_DATABASE_PREFIX
 import org.apache.commons.lang3.ArrayUtils.EMPTY_BYTE_ARRAY
@@ -26,8 +27,10 @@ import javax.persistence.*
  * This class needs database transactions to be in-flight during method calls and init.
  */
 @Deprecated("Superseded by net.corda.node.services.keys.BasicHSMKeyManagementService")
-class PersistentKeyManagementService(cacheFactory: NamedCacheFactory, val identityService: PersistentIdentityService,
+class PersistentKeyManagementService(cacheFactory: NamedCacheFactory,
-                                     private val database: CordaPersistence) : SingletonSerializeAsToken(), KeyManagementServiceInternal {
+                                     override val identityService: PersistentIdentityService,
+                                     private val database: CordaPersistence,
+                                     private val pkToIdCache: WritablePublicKeyToOwningIdentityCache) : SingletonSerializeAsToken(), KeyManagementServiceInternal {
     @Entity
     @Table(name = "${NODE_DATABASE_PREFIX}our_key_pairs")
     class PersistentKey(
@@ -76,31 +79,16 @@ class PersistentKeyManagementService(cacheFactory: NamedCacheFactory, val identi
         identityService.stripNotOurKeys(candidateKeys)
     }
 
-    override fun freshKey(): PublicKey {
+    override fun freshKeyInternal(externalId: UUID?): PublicKey {
         val keyPair = generateKeyPair()
         database.transaction {
             keysMap[keyPair.public] = keyPair.private
+            pkToIdCache[keyPair.public] = KeyOwningIdentity.fromUUID(externalId)
         }
         return keyPair.public
     }
 
-    override fun freshKey(externalId: UUID): PublicKey {
+    override fun getSigner(publicKey: PublicKey): ContentSigner = getSigner(getSigningKeyPair(publicKey))
-        val newKey = freshKey()
-        database.transaction { session.persist(PublicKeyHashToExternalId(externalId, newKey)) }
-        return newKey
-    }
-
-    override fun freshKeyAndCert(identity: PartyAndCertificate, revocationEnabled: Boolean): PartyAndCertificate {
-        return freshCertificate(identityService, freshKey(), identity, getSigner(identity.owningKey))
-    }
-
-    override fun freshKeyAndCert(identity: PartyAndCertificate, revocationEnabled: Boolean, externalId: UUID): PartyAndCertificate {
-        val newKeyWithCert = freshKeyAndCert(identity, revocationEnabled)
-        database.transaction { session.persist(PublicKeyHashToExternalId(externalId, newKeyWithCert.owningKey)) }
-        return newKeyWithCert
-    }
-
-    private fun getSigner(publicKey: PublicKey): ContentSigner = getSigner(getSigningKeyPair(publicKey))
 
     //It looks for the PublicKey in the (potentially) CompositeKey that is ours, and then returns the associated PrivateKey to use in signing
     private fun getSigningKeyPair(publicKey: PublicKey): KeyPair {

node/src/main/kotlin/net/corda/node/services/persistence/PublicKeyHashToExternalId.kt

@@ -0,0 +1,26 @@
+package net.corda.node.services.persistence
+
+import net.corda.core.crypto.toStringShort
+import org.hibernate.annotations.Type
+import java.security.PublicKey
+import java.util.*
+import javax.persistence.*
+
+@Entity
+@Table(name = "pk_hash_to_ext_id_map", indexes = [Index(name = "pk_hash_to_xid_idx", columnList = "public_key_hash")])
+class PublicKeyHashToExternalId(
+        @Id
+        @GeneratedValue
+        @Column(name = "id", unique = true, nullable = false)
+        val key: Long?,
+
+        @Column(name = "external_id", nullable = false)
+        @Type(type = "uuid-char")
+        val externalId: UUID,
+
+        @Column(name = "public_key_hash", nullable = false)
+        val publicKeyHash: String
+) {
+    constructor(accountId: UUID, publicKey: PublicKey)
+            : this(null, accountId, publicKey.toStringShort())
+}

node/src/main/kotlin/net/corda/node/services/persistence/PublicKeyToOwningIdentityCacheImpl.kt

@@ -0,0 +1,89 @@
+package net.corda.node.services.persistence
+
+import com.github.benmanes.caffeine.cache.Caffeine
+import net.corda.core.crypto.toStringShort
+import net.corda.core.internal.NamedCacheFactory
+import net.corda.core.utilities.contextLogger
+import net.corda.core.utilities.debug
+import net.corda.node.services.keys.BasicHSMKeyManagementService
+import net.corda.nodeapi.internal.KeyOwningIdentity
+import net.corda.nodeapi.internal.persistence.CordaPersistence
+import java.security.PublicKey
+import java.util.*
+
+/**
+ * The [PublicKeyToOwningIdentityCacheImpl] provides a caching layer over the pk_hash_to_external_id table. Gets will attempt to read an
+ * external identity from the database if it is not present in memory, while sets will write external identity UUIDs to this database table.
+ */
+class PublicKeyToOwningIdentityCacheImpl(private val database: CordaPersistence,
+                                         cacheFactory: NamedCacheFactory) : WritablePublicKeyToOwningIdentityCache {
+    companion object {
+        val log = contextLogger()
+    }
+
+    private val cache = cacheFactory.buildNamed<PublicKey, KeyOwningIdentity>(Caffeine.newBuilder(), "PublicKeyToOwningIdentityCache_cache")
+
+    private fun isKeyBelongingToNode(key: PublicKey): Boolean {
+        return database.transaction {
+            val criteriaBuilder = session.criteriaBuilder
+            val criteriaQuery = criteriaBuilder.createQuery(Long::class.java)
+            val queryRoot = criteriaQuery.from(BasicHSMKeyManagementService.PersistentKey::class.java)
+            criteriaQuery.select(criteriaBuilder.count(queryRoot))
+            criteriaQuery.where(
+                    criteriaBuilder.equal(queryRoot.get<String>(BasicHSMKeyManagementService.PersistentKey::publicKeyHash.name), key.toStringShort())
+            )
+            val query = session.createQuery(criteriaQuery)
+            query.uniqueResult() > 0
+        }
+    }
+
+    /**
+     * Return the owning identity associated with a given key.
+     *
+     * This method caches the result of a database lookup to prevent multiple database accesses for the same key. This assumes that once a
+     * key is generated, the UUID assigned to it is never changed.
+     */
+    override operator fun get(key: PublicKey): KeyOwningIdentity? {
+        return cache.asMap().computeIfAbsent(key) {
+            database.transaction {
+                val criteriaBuilder = session.criteriaBuilder
+                val criteriaQuery = criteriaBuilder.createQuery(UUID::class.java)
+                val queryRoot = criteriaQuery.from(PublicKeyHashToExternalId::class.java)
+                criteriaQuery.select(queryRoot.get(PublicKeyHashToExternalId::externalId.name))
+                criteriaQuery.where(
+                        criteriaBuilder.equal(queryRoot.get<String>(PublicKeyHashToExternalId::publicKeyHash.name), key.toStringShort())
+                )
+                val query = session.createQuery(criteriaQuery)
+                val uuid = query.uniqueResult()
+                if (uuid != null || isKeyBelongingToNode(key)) {
+                    val signingEntity = KeyOwningIdentity.fromUUID(uuid)
+                    log.debug { "Database lookup for public key ${key.toStringShort()}, found signing entity $signingEntity" }
+                    signingEntity
+                } else {
+                    log.debug { "Attempted to find owning identity for public key ${key.toStringShort()}, but key is unknown to node" }
+                    null
+                }
+            }
+        }
+    }
+
+    /**
+     * Associate a key with a given [KeyOwningIdentity].
+     *
+     * This will write to the pk_hash_to_external_id table if the key belongs to an external id. If a key is created within a transaction
+     * that is rolled back in the future, the cache may contain stale entries. However, that key should be missing from the
+     * KeyManagementService in that case, and so querying it from this cache should not occur (as the key is inaccessible).
+     *
+     * The same key should not be written twice.
+     */
+    override operator fun set(key: PublicKey, value: KeyOwningIdentity) {
+        when (value) {
+            is KeyOwningIdentity.MappedIdentity -> {
+                database.transaction { session.persist(PublicKeyHashToExternalId(value.uuid, key)) }
+            }
+            is KeyOwningIdentity.UnmappedIdentity -> {
+            }
+        }
+        cache.asMap()[key] = value
+    }
+}

node/src/main/kotlin/net/corda/node/services/persistence/WritablePublicKeyToOwningIdentityCache.kt

@@ -0,0 +1,16 @@
+package net.corda.node.services.persistence
+
+import net.corda.nodeapi.internal.KeyOwningIdentity
+import net.corda.nodeapi.internal.PublicKeyToOwningIdentityCache
+import java.security.PublicKey
+
+/**
+ * Internal only version of a [PublicKeyToOwningIdentityCache] that allows writing to the cache and underlying database table
+ */
+interface WritablePublicKeyToOwningIdentityCache : PublicKeyToOwningIdentityCache {
+
+    /**
+     * Assign a public key to an owning identity.
+     */
+    operator fun set(key: PublicKey, value: KeyOwningIdentity)
+}

node/src/main/kotlin/net/corda/node/services/schema/NodeSchemaService.kt

@@ -15,11 +15,11 @@ import net.corda.node.services.events.NodeSchedulerService
 import net.corda.node.services.identity.PersistentIdentityService
 import net.corda.node.services.keys.BasicHSMKeyManagementService
 import net.corda.node.services.keys.PersistentKeyManagementService
-import net.corda.node.services.keys.PublicKeyHashToExternalId
 import net.corda.node.services.messaging.P2PMessageDeduplicator
 import net.corda.node.services.persistence.DBCheckpointStorage
 import net.corda.node.services.persistence.DBTransactionStorage
 import net.corda.node.services.persistence.NodeAttachmentService
+import net.corda.node.services.persistence.PublicKeyHashToExternalId
 import net.corda.node.services.upgrade.ContractUpgradeServiceImpl
 import net.corda.node.services.vault.VaultSchemaV1
 

node/src/main/kotlin/net/corda/node/utilities/NodeNamedCache.kt

@@ -59,6 +59,7 @@ open class DefaultNamedCacheFactory protected constructor(private val metricRegi
                 name == "RaftUniquenessProvider_transactions" -> caffeine.maximumSize(defaultCacheSize)
                 name == "BasicHSMKeyManagementService_keys" -> caffeine.maximumSize(defaultCacheSize)
                 name == "NodeParametersStorage_networkParametersByHash" -> caffeine.maximumSize(defaultCacheSize)
+                name == "PublicKeyToOwningIdentityCache_cache" -> caffeine.maximumSize(defaultCacheSize)
                 else -> throw IllegalArgumentException("Unexpected cache name $name. Did you add a new cache?")
             }
         }

node/src/test/kotlin/net/corda/node/services/network/NetworkMapUpdaterTest.kt

@@ -33,6 +33,7 @@ import net.corda.testing.internal.DEV_ROOT_CA
 import net.corda.testing.internal.TestNodeInfoBuilder
 import net.corda.testing.internal.createNodeInfoAndSigned
 import net.corda.testing.node.internal.MockKeyManagementService
+import net.corda.testing.node.internal.MockPublicKeyToOwningIdentityCache
 import net.corda.testing.node.internal.network.NetworkMapServer
 import net.corda.testing.node.makeTestIdentityService
 import org.assertj.core.api.Assertions.assertThat
@@ -102,7 +103,7 @@ class NetworkMapUpdaterTest {
                 server.networkParameters.serialize().hash,
                 ourNodeInfo,
                 networkParameters,
-                MockKeyManagementService(makeTestIdentityService(), ourKeyPair),
+                MockKeyManagementService(makeTestIdentityService(), ourKeyPair, pkToIdCache = MockPublicKeyToOwningIdentityCache()),
                 NetworkParameterAcceptanceSettings(autoAcceptNetworkParameters, excludedAutoAcceptNetworkParameters))
     }
 

node/src/test/kotlin/net/corda/node/services/network/NodeInfoWatcherTest.kt

@@ -2,17 +2,18 @@ package net.corda.node.services.network
 
 import com.google.common.jimfs.Configuration
 import com.google.common.jimfs.Jimfs
+import net.corda.core.internal.NODE_INFO_DIRECTORY
 import net.corda.core.internal.createDirectories
 import net.corda.core.internal.div
 import net.corda.core.internal.size
 import net.corda.core.node.services.KeyManagementService
-import net.corda.core.internal.NODE_INFO_DIRECTORY
 import net.corda.nodeapi.internal.NodeInfoAndSigned
 import net.corda.nodeapi.internal.network.NodeInfoFilesCopier
 import net.corda.testing.core.ALICE_NAME
 import net.corda.testing.core.SerializationEnvironmentRule
 import net.corda.testing.internal.createNodeInfoAndSigned
 import net.corda.testing.node.internal.MockKeyManagementService
+import net.corda.testing.node.internal.MockPublicKeyToOwningIdentityCache
 import net.corda.testing.node.makeTestIdentityService
 import org.assertj.core.api.Assertions.assertThat
 import org.junit.Before
@@ -49,7 +50,7 @@ class NodeInfoWatcherTest {
     fun start() {
         nodeInfoAndSigned = createNodeInfoAndSigned(ALICE_NAME)
         val identityService = makeTestIdentityService()
-        keyManagementService = MockKeyManagementService(identityService)
+        keyManagementService = MockKeyManagementService(identityService, pkToIdCache = MockPublicKeyToOwningIdentityCache())
         nodeInfoWatcher = NodeInfoWatcher(tempFolder.root.toPath(), scheduler)
         nodeInfoPath = tempFolder.root.toPath() / NODE_INFO_DIRECTORY
     }

node/src/test/kotlin/net/corda/node/services/persistence/PublicKeyToOwningIdentityCacheImplTest.kt

@@ -0,0 +1,121 @@
+package net.corda.node.services.persistence
+
+import junit.framework.TestCase.assertEquals
+import net.corda.core.crypto.generateKeyPair
+import net.corda.core.node.services.KeyManagementService
+import net.corda.core.utilities.getOrThrow
+import net.corda.nodeapi.internal.KeyOwningIdentity
+import net.corda.nodeapi.internal.persistence.CordaPersistence
+import net.corda.nodeapi.internal.persistence.withoutDatabaseAccess
+import net.corda.testing.common.internal.testNetworkParameters
+import net.corda.testing.core.ALICE_NAME
+import net.corda.testing.core.TestIdentity
+import net.corda.testing.internal.TestingNamedCacheFactory
+import net.corda.testing.node.MockServices
+import org.junit.After
+import org.junit.Before
+import org.junit.Test
+import java.security.PublicKey
+import java.util.*
+import java.util.concurrent.ExecutorService
+import java.util.concurrent.Executors
+
+class PublicKeyToOwningIdentityCacheImplTest {
+
+    private lateinit var database: CordaPersistence
+    private lateinit var testCache: PublicKeyToOwningIdentityCacheImpl
+    private lateinit var keyManagementService: KeyManagementService
+    private val testKeys = mutableListOf<Pair<KeyOwningIdentity, PublicKey>>()
+    private val alice = TestIdentity(ALICE_NAME, 20)
+    private lateinit var executor: ExecutorService
+
+    @Before
+    fun setUp() {
+        val databaseAndServices = MockServices.makeTestDatabaseAndPersistentServices(
+                listOf(),
+                alice,
+                testNetworkParameters(),
+                emptySet(),
+                emptySet()
+        )
+        database = databaseAndServices.first
+        testCache = PublicKeyToOwningIdentityCacheImpl(database, TestingNamedCacheFactory())
+        keyManagementService = databaseAndServices.second.keyManagementService
+        createTestKeys()
+        executor = Executors.newFixedThreadPool(2)
+    }
+
+    @After
+    fun tearDown() {
+        database.close()
+        executor.shutdown()
+    }
+
+    private fun createTestKeys() {
+        val duplicatedUUID = UUID.randomUUID()
+        val uuids = listOf(UUID.randomUUID(), UUID.randomUUID(), null, null, duplicatedUUID, duplicatedUUID)
+        uuids.forEach {
+            val key = if (it != null) {
+                keyManagementService.freshKey(it)
+            } else {
+                keyManagementService.freshKey()
+            }
+            testKeys.add(Pair(KeyOwningIdentity.fromUUID(it), key))
+        }
+    }
+
+    private fun performTestRun() {
+        for ((keyOwningIdentity, key) in testKeys) {
+            assertEquals(keyOwningIdentity, testCache[key])
+        }
+    }
+
+    @Test
+    fun `cache returns right key for each UUID`() {
+        performTestRun()
+    }
+
+    @Test
+    fun `querying for key twice does not go to database the second time`() {
+        performTestRun()
+
+        withoutDatabaseAccess {
+            performTestRun()
+        }
+    }
+
+    @Test
+    fun `entries can be fetched if cache invalidated`() {
+        testCache = PublicKeyToOwningIdentityCacheImpl(database, TestingNamedCacheFactory(sizeOverride = 0))
+
+        performTestRun()
+    }
+
+    @Test
+    fun `cache access is thread safe`() {
+        val executor = Executors.newFixedThreadPool(2)
+        val f1 = executor.submit { performTestRun() }
+        val f2 = executor.submit { performTestRun() }
+        f2.getOrThrow()
+        f1.getOrThrow()
+    }
+
+    private fun createAndAddKeys() {
+        keyManagementService.freshKey(UUID.randomUUID())
+    }
+
+    @Test
+    fun `can set multiple keys across threads`() {
+        val executor = Executors.newFixedThreadPool(2)
+        val f1 = executor.submit { repeat(5) { createAndAddKeys() } }
+        val f2 = executor.submit { repeat(5) { createAndAddKeys() } }
+        f2.getOrThrow()
+        f1.getOrThrow()
+    }
+
+    @Test
+    fun `requesting a key unknown to the node returns null`() {
+        val keys = generateKeyPair()
+        assertEquals(null, testCache[keys.public])
+    }
+}

testing/node-driver/src/main/kotlin/net/corda/testing/node/MockServices.kt

@@ -21,16 +21,17 @@ import net.corda.core.serialization.SerializeAsToken
 import net.corda.core.transactions.SignedTransaction
 import net.corda.core.utilities.NetworkHostAndPort
 import net.corda.node.VersionInfo
-import net.corda.nodeapi.internal.cordapp.CordappLoader
 import net.corda.node.internal.ServicesForResolutionImpl
 import net.corda.node.internal.cordapp.JarScanningCordappLoader
 import net.corda.node.services.api.*
 import net.corda.node.services.identity.InMemoryIdentityService
 import net.corda.node.services.identity.PersistentIdentityService
 import net.corda.node.services.keys.PersistentKeyManagementService
+import net.corda.node.services.persistence.PublicKeyToOwningIdentityCacheImpl
 import net.corda.node.services.schema.NodeSchemaService
 import net.corda.node.services.transactions.InMemoryTransactionVerifierService
 import net.corda.node.services.vault.NodeVaultService
+import net.corda.nodeapi.internal.cordapp.CordappLoader
 import net.corda.nodeapi.internal.persistence.CordaPersistence
 import net.corda.nodeapi.internal.persistence.DatabaseConfig
 import net.corda.nodeapi.internal.persistence.contextTransaction
@@ -75,7 +76,11 @@ open class MockServices private constructor(
         private val initialNetworkParameters: NetworkParameters,
         private val initialIdentity: TestIdentity,
         private val moreKeys: Array<out KeyPair>,
-        override val keyManagementService: KeyManagementService = MockKeyManagementService(identityService, *arrayOf(initialIdentity.keyPair) + moreKeys)
+        override val keyManagementService: KeyManagementService = MockKeyManagementService(
+                identityService,
+                *arrayOf(initialIdentity.keyPair) + moreKeys,
+                pkToIdCache = MockPublicKeyToOwningIdentityCache()
+        )
 ) : ServiceHub {
 
     companion object {
@@ -120,7 +125,11 @@ open class MockServices private constructor(
             val dataSourceProps = makeTestDataSourceProperties()
             val schemaService = NodeSchemaService(cordappLoader.cordappSchemas)
             val database = configureDatabase(dataSourceProps, DatabaseConfig(), identityService::wellKnownPartyFromX500Name, identityService::wellKnownPartyFromAnonymous, schemaService, schemaService.internalSchemas())
-            val keyManagementService = MockKeyManagementService(identityService, *arrayOf(initialIdentity.keyPair) + moreKeys)
+            val keyManagementService = MockKeyManagementService(
+                    identityService,
+                    *arrayOf(initialIdentity.keyPair) + moreKeys,
+                    pkToIdCache = MockPublicKeyToOwningIdentityCache()
+            )
             val mockService = database.transaction {
                 makeMockMockServices(cordappLoader, identityService, networkParameters, initialIdentity, moreKeys.toSet(), keyManagementService, schemaService, database)
             }
@@ -163,7 +172,8 @@ open class MockServices private constructor(
             // Create a persistent key management service and add the key pair which was created for the TestIdentity.
             // We only add the keypair for the initial identity and any other keys which this node may control. Note: We don't add the keys
             // for the other identities.
-            val keyManagementService = PersistentKeyManagementService(TestingNamedCacheFactory(), identityService, persistence)
+            val pkToIdCache = PublicKeyToOwningIdentityCacheImpl(persistence, TestingNamedCacheFactory())
+            val keyManagementService = PersistentKeyManagementService(TestingNamedCacheFactory(), identityService, persistence, pkToIdCache)
             persistence.transaction { keyManagementService.start(moreKeys + initialIdentity.keyPair) }
 
             val mockService = persistence.transaction {

testing/node-driver/src/main/kotlin/net/corda/testing/node/internal/InternalMockNetwork.kt

@@ -36,7 +36,6 @@ import net.corda.node.services.config.*
 import net.corda.node.services.identity.PersistentIdentityService
 import net.corda.node.services.keys.BasicHSMKeyManagementService
 import net.corda.node.services.keys.KeyManagementServiceInternal
-import net.corda.nodeapi.internal.cryptoservice.bouncycastle.BCCryptoService
 import net.corda.node.services.messaging.Message
 import net.corda.node.services.messaging.MessagingService
 import net.corda.node.services.persistence.NodeAttachmentService
@@ -46,6 +45,7 @@ import net.corda.node.utilities.AffinityExecutor.ServiceAffinityExecutor
 import net.corda.node.utilities.DefaultNamedCacheFactory
 import net.corda.nodeapi.internal.DevIdentityGenerator
 import net.corda.nodeapi.internal.config.User
+import net.corda.nodeapi.internal.cryptoservice.bouncycastle.BCCryptoService
 import net.corda.nodeapi.internal.network.NetworkParametersCopier
 import net.corda.nodeapi.internal.persistence.CordaPersistence
 import net.corda.nodeapi.internal.persistence.DatabaseConfig
@@ -373,7 +373,7 @@ open class InternalMockNetwork(cordappPackages: List<String> = emptyList(),
         }
 
         override fun makeKeyManagementService(identityService: PersistentIdentityService): KeyManagementServiceInternal {
-            return BasicHSMKeyManagementService(cacheFactory, identityService, database, cryptoService)
+            return BasicHSMKeyManagementService(cacheFactory, identityService, database, cryptoService, pkToIdCache)
         }
 
         override fun startShell() {

testing/node-driver/src/main/kotlin/net/corda/testing/node/internal/MockKeyManagementService.kt

@@ -1,63 +1,46 @@
 package net.corda.testing.node.internal
 
 import net.corda.core.crypto.*
-import net.corda.core.identity.PartyAndCertificate
 import net.corda.core.node.services.IdentityService
 import net.corda.core.node.services.KeyManagementService
 import net.corda.core.serialization.SingletonSerializeAsToken
-import net.corda.node.services.keys.freshCertificate
+import net.corda.node.services.keys.KeyManagementServiceInternal
+import net.corda.node.services.persistence.WritablePublicKeyToOwningIdentityCache
+import net.corda.nodeapi.internal.KeyOwningIdentity
 import org.bouncycastle.operator.ContentSigner
 import java.security.KeyPair
 import java.security.PrivateKey
 import java.security.PublicKey
 import java.util.*
-import java.util.concurrent.ConcurrentHashMap
 
 /**
  * A class which provides an implementation of [KeyManagementService] which is used in [MockServices]
  *
  * @property identityService The [IdentityService] which contains the given identities.
  */
-class MockKeyManagementService(val identityService: IdentityService,
+class MockKeyManagementService(override val identityService: IdentityService,
-                               vararg initialKeys: KeyPair) : SingletonSerializeAsToken(), KeyManagementService {
+                               vararg initialKeys: KeyPair,
+                               private val pkToIdCache: WritablePublicKeyToOwningIdentityCache) : SingletonSerializeAsToken(), KeyManagementServiceInternal {
     private val keyStore: MutableMap<PublicKey, PrivateKey> = initialKeys.associateByTo(HashMap(), { it.public }, { it.private })
 
     override val keys: Set<PublicKey> get() = keyStore.keys
 
     private val nextKeys = LinkedList<KeyPair>()
 
-    val keysById: MutableMap<UUID, Set<PublicKey>> = ConcurrentHashMap()
+    override fun freshKeyInternal(externalId: UUID?): PublicKey {
-
-    override fun freshKey(): PublicKey {
         val k = nextKeys.poll() ?: generateKeyPair()
         keyStore[k.public] = k.private
+        pkToIdCache[k.public] = KeyOwningIdentity.fromUUID(externalId)
         return k.public
     }
 
-    private fun mapKeyToId(publicKey: PublicKey, externalId: UUID) {
-        val keysForId = keysById.getOrPut(externalId) { emptySet() }
-        keysById[externalId] = keysForId + publicKey
-    }
-
-    override fun freshKey(externalId: UUID): PublicKey {
-        val key = freshKey()
-        mapKeyToId(key, externalId)
-        return key
-    }
-
-    override fun freshKeyAndCert(identity: PartyAndCertificate, revocationEnabled: Boolean, externalId: UUID): PartyAndCertificate {
-        val keyAndCert = freshKeyAndCert(identity, revocationEnabled)
-        mapKeyToId(keyAndCert.owningKey, externalId)
-        return keyAndCert
-    }
-
     override fun filterMyKeys(candidateKeys: Iterable<PublicKey>): Iterable<PublicKey> = candidateKeys.filter { it in this.keys }
 
-    override fun freshKeyAndCert(identity: PartyAndCertificate, revocationEnabled: Boolean): PartyAndCertificate {
+    override fun getSigner(publicKey: PublicKey): ContentSigner = net.corda.node.services.keys.getSigner(getSigningKeyPair(publicKey))
-        return freshCertificate(identityService, freshKey(), identity, getSigner(identity.owningKey))
-    }
 
-    private fun getSigner(publicKey: PublicKey): ContentSigner = net.corda.node.services.keys.getSigner(getSigningKeyPair(publicKey))
+    override fun start(initialKeyPairs: Set<KeyPair>) {
+        initialKeyPairs.forEach { keyStore[it.public] = it.private }
+    }
 
     private fun getSigningKeyPair(publicKey: PublicKey): KeyPair {
         val pk = publicKey.keys.firstOrNull { keyStore.containsKey(it) }

testing/node-driver/src/main/kotlin/net/corda/testing/node/internal/MockPublicKeyToOwningIdentityCache.kt

@@ -0,0 +1,23 @@
+package net.corda.testing.node.internal
+
+import net.corda.core.internal.toSynchronised
+import net.corda.node.services.persistence.WritablePublicKeyToOwningIdentityCache
+import net.corda.nodeapi.internal.KeyOwningIdentity
+import java.security.PublicKey
+
+/**
+ * A mock implementation of [WritablePublicKeyToOwningIdentityCache] that stores all key mappings in memory. Used in testing scenarios that do not
+ * require database access.
+ */
+class MockPublicKeyToOwningIdentityCache : WritablePublicKeyToOwningIdentityCache {
+
+    private val cache: MutableMap<PublicKey, KeyOwningIdentity> = mutableMapOf<PublicKey, KeyOwningIdentity>().toSynchronised()
+
+    override fun get(key: PublicKey): KeyOwningIdentity? {
+        return cache[key]
+    }
+
+    override fun set(key: PublicKey, value: KeyOwningIdentity) {
+        cache[key] = value
+    }
+}