CORDA-2563: Assign external IDs when creating new keys (#4727)

* First pass at fixing 2563.

* In memory KMS now maps keys to IDs.

* CreateDatabaseAndMockServices now creates a persistent key management service and a can take a persistent identity service, so now the external id mapping works for mock services.

* * Created a helper for mock services which allows the creation of a mock services with persistent identity management service key management service and vault.
* MockNode now uses persistent key management service - not sure why it didn't do before?

* * MockNode now uses BasicHSMKeyManagementService
* Updated api-current file

* Little fix required after rebase to master.

* Fixed broken test.

* Added informative error messages to UnsupportedOperationExceptions thrown by E2ETestKeyManagementService.

* Removed redundant private constructor for mock services from api-current.txt.

* Addressed Rick's comments.

.ci/api-current.txt

@@ -7745,7 +7745,7 @@ public class net.corda.testing.node.MockServices extends java.lang.Object implem
   public <init>(net.corda.core.identity.CordaX500Name, net.corda.core.node.services.IdentityService)
   public <init>(net.corda.core.identity.CordaX500Name, net.corda.core.node.services.IdentityService, int, kotlin.jvm.internal.DefaultConstructorMarker)
   public <init>(net.corda.core.identity.CordaX500Name, net.corda.core.node.services.IdentityService, java.security.KeyPair, java.security.KeyPair...)
-  public <init>(net.corda.core.identity.CordaX500Name, net.corda.core.node.services.IdentityService, java.security.KeyPair, java.security.KeyPair[], int, kotlin.jvm.internal.DefaultConstructorMarker)
+  public <init>(net.corda.node.cordapp.CordappLoader, net.corda.core.node.services.IdentityService, net.corda.core.node.NetworkParameters, net.corda.testing.core.TestIdentity, java.security.KeyPair[], net.corda.core.node.services.KeyManagementService, kotlin.jvm.internal.DefaultConstructorMarker)
   public <init>(net.corda.testing.core.TestIdentity, net.corda.core.node.NetworkParameters, net.corda.testing.core.TestIdentity...)
   public <init>(net.corda.testing.core.TestIdentity, net.corda.testing.core.TestIdentity...)
   public final void addMockCordapp(String)

core/src/main/kotlin/net/corda/core/node/services/KeyManagementService.kt

@@ -7,10 +7,12 @@ import net.corda.core.crypto.DigitalSignature
 import net.corda.core.crypto.SignableData
 import net.corda.core.crypto.TransactionSignature
 import net.corda.core.identity.PartyAndCertificate
+import net.corda.core.node.ServiceHub
 import java.security.KeyPair
 import java.security.PrivateKey
 import java.security.PublicKey
 import java.security.cert.X509Certificate
+import java.util.*
 
 /**
  * The KMS is responsible for storing and using private keys to sign things. An implementation of this may, for example,
@@ -31,8 +33,15 @@ interface KeyManagementService {
     fun freshKey(): PublicKey
 
     /**
-     * Generates a new random [KeyPair], adds it to the internal key storage, then generates a corresponding
+     * Generates a new random [KeyPair] and adds it to the internal key storage. Associates the public key to an external ID. Returns the
-     * [X509Certificate] and adds it to the identity service.
+     * public key part of the pair.
+     */
+    @Suspendable
+    fun freshKey(externalId: UUID): PublicKey
+
+    /**
+     * Generates a new random [KeyPair], adds it to the internal key storage, then generates a corresponding [X509Certificate] and adds it
+     * to the identity service. Associates the public key to an external ID. Returns the public part of the pair.
      *
      * @param identity identity to generate a key and certificate for. Must be an identity this node has CA privileges for.
      * @param revocationEnabled whether to check revocation status of certificates in the certificate path.
@@ -41,6 +50,18 @@ interface KeyManagementService {
     @Suspendable
     fun freshKeyAndCert(identity: PartyAndCertificate, revocationEnabled: Boolean): PartyAndCertificate
 
+    /**
+     * Generates a new random [KeyPair], adds it to the internal key storage, then generates a corresponding
+     * [X509Certificate] and adds it to the identity service.
+     *
+     * @param identity identity to generate a key and certificate for. Must be an identity this node has CA privileges for.
+     * @param revocationEnabled whether to check revocation status of certificates in the certificate path.
+     * @param externalId ID to associate the newly created [PublicKey] with.
+     * @return X.509 certificate and path to the trust root.
+     */
+    @Suspendable
+    fun freshKeyAndCert(identity: PartyAndCertificate, revocationEnabled: Boolean, externalId: UUID): PartyAndCertificate
+
     /**
      * Filter some keys down to the set that this node owns (has private keys for).
      *
@@ -69,4 +90,4 @@ interface KeyManagementService {
      */
     @Suspendable
     fun sign(signableData: SignableData, publicKey: PublicKey): TransactionSignature
-}
+}

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

@@ -49,26 +49,7 @@ class BasicHSMKeyManagementService(cacheFactory: NamedCacheFactory, val identity
             : this(publicKey.toStringShort(), publicKey.encoded, privateKey.encoded)
     }
 
-    @Entity
+    private companion object {
-    @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())
-    }
-
-    companion object {
         fun createKeyMap(cacheFactory: NamedCacheFactory): AppendOnlyPersistentMap<PublicKey, PrivateKey, PersistentKey, String> {
             return AppendOnlyPersistentMap(
                     cacheFactory = cacheFactory,
@@ -116,10 +97,22 @@ class BasicHSMKeyManagementService(cacheFactory: NamedCacheFactory, val identity
         return keyPair.public
     }
 
+    override fun freshKey(externalId: UUID): 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 {
         val signingPublicKey = getSigningPublicKey(publicKey)
         return if (signingPublicKey in originalKeysMap) {

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

@@ -61,10 +61,18 @@ class E2ETestKeyManagementService(val identityService: IdentityService, private
         return keyPair.public
     }
 
+    override fun freshKey(externalId: UUID): 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 {

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

@@ -1,8 +1,32 @@
 package net.corda.node.services.keys
 
+import net.corda.core.crypto.toStringShort
 import net.corda.core.node.services.KeyManagementService
+import org.hibernate.annotations.Type
 import java.security.KeyPair
+import java.security.PublicKey
+import java.util.*
+import javax.persistence.*
 
 interface KeyManagementServiceInternal : KeyManagementService {
     fun start(initialKeyPairs: Set<KeyPair>)
 }
+
+@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

@@ -14,6 +14,7 @@ import org.bouncycastle.operator.ContentSigner
 import java.security.KeyPair
 import java.security.PrivateKey
 import java.security.PublicKey
+import java.util.*
 import javax.persistence.*
 
 /**
@@ -82,10 +83,22 @@ class PersistentKeyManagementService(cacheFactory: NamedCacheFactory, val identi
         return keyPair.public
     }
 
+    override fun freshKey(externalId: UUID): 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

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

@@ -15,6 +15,7 @@ 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
@@ -45,7 +46,7 @@ class NodeSchemaService(private val extraSchemas: Set<MappedSchema> = emptySet()
                     PersistentIdentityService.PersistentIdentityNames::class.java,
                     ContractUpgradeServiceImpl.DBContractUpgrade::class.java,
                     DBNetworkParametersStorage.PersistentNetworkParameters::class.java,
-                    BasicHSMKeyManagementService.PublicKeyHashToExternalId::class.java
+                    PublicKeyHashToExternalId::class.java
             )) {
         override val migrationResource = "node-core.changelog-master"
     }
@@ -97,3 +98,5 @@ class NodeSchemaService(private val extraSchemas: Set<MappedSchema> = emptySet()
             schemaOptions.keys.map { schema -> crossReferencesToOtherMappedSchema(schema) }.flatMap { it.toList() }
 
 }
+
+

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

@@ -2,7 +2,10 @@ package net.corda.node.services.persistence
 
 import co.paralleluniverse.fibers.Suspendable
 import net.corda.core.contracts.Amount
+import net.corda.core.identity.CordaX500Name
 import net.corda.core.identity.Party
+import net.corda.core.node.services.KeyManagementService
+import net.corda.core.transactions.SignedTransaction
 import net.corda.core.transactions.TransactionBuilder
 import net.corda.core.utilities.OpaqueBytes
 import net.corda.core.utilities.getOrThrow
@@ -11,52 +14,56 @@ import net.corda.finance.`issued by`
 import net.corda.finance.contracts.asset.Cash
 import net.corda.finance.flows.AbstractCashFlow
 import net.corda.finance.issuedBy
+import net.corda.node.migration.VaultStateMigrationTest.Companion.bankOfCorda
 import net.corda.node.services.identity.PersistentIdentityService
+import net.corda.node.services.keys.BasicHSMKeyManagementService
 import net.corda.node.services.keys.E2ETestKeyManagementService
+import net.corda.nodeapi.internal.persistence.CordaPersistence
+import net.corda.testing.common.internal.testNetworkParameters
 import net.corda.testing.core.BOC_NAME
+import net.corda.testing.core.SerializationEnvironmentRule
+import net.corda.testing.core.TestIdentity
 import net.corda.testing.internal.TestingNamedCacheFactory
 import net.corda.testing.node.InMemoryMessagingNetwork.ServicePeerAllocationStrategy.RoundRobin
 import net.corda.testing.node.MockNetwork
 import net.corda.testing.node.MockNetworkParameters
+import net.corda.testing.node.MockServices
 import net.corda.testing.node.StartedMockNode
 import net.corda.testing.node.internal.FINANCE_CORDAPPS
+import net.corda.testing.node.internal.InternalMockNetwork
+import net.corda.testing.node.internal.TestStartedNode
 import org.junit.After
 import org.junit.Before
+import org.junit.Rule
 import org.junit.Test
 import java.util.*
 import kotlin.test.assertEquals
 
 class HibernateColumnConverterTests {
-    private lateinit var mockNet: MockNetwork
-    private lateinit var bankOfCordaNode: StartedMockNode
-    private lateinit var bankOfCorda: Party
-    private lateinit var notary: Party
 
-    class TestCashIssueFlow(private val amount: Amount<Currency>,
+    @Rule
-                            private val issuerBankPartyRef: OpaqueBytes,
+    @JvmField
-                            private val notary: Party) : AbstractCashFlow<AbstractCashFlow.Result>(tracker()) {
+    val testSerialization = SerializationEnvironmentRule()
-        @Suspendable
+
-        override fun call(): AbstractCashFlow.Result {
+    private val cordapps = listOf("net.corda.finance")
-            val builder = TransactionBuilder(notary)
+
-            val issuer = ourIdentity.ref(issuerBankPartyRef)
+    private val myself = TestIdentity(CordaX500Name("Me", "London", "GB"))
-            val signers = Cash().generateIssue(builder, amount.issuedBy(issuer), ourIdentity, notary)
+    private val notary = TestIdentity(CordaX500Name("NotaryService", "London", "GB"), 1337L)
-            val tx = serviceHub.signInitialTransaction(builder, signers)
+
-            serviceHub.recordTransactions(tx)
+    lateinit var services: MockServices
-            return Result(tx, ourIdentity)
+    lateinit var database: CordaPersistence
-        }
+
-    }
-    
     @Before
-    fun start() {
+    fun setUp() {
-        mockNet = MockNetwork(MockNetworkParameters(servicePeerAllocationStrategy = RoundRobin(), cordappsForAllNodes = FINANCE_CORDAPPS))
+        val (db, mockServices) = MockServices.makeTestDatabaseAndPersistentServices(
-        bankOfCordaNode = mockNet.createPartyNode(BOC_NAME)
+                cordappPackages = cordapps,
-        bankOfCorda = bankOfCordaNode.info.identityFromX500Name(BOC_NAME)
+                initialIdentity = myself,
-        notary = mockNet.defaultNotaryIdentity
+                networkParameters = testNetworkParameters(minimumPlatformVersion = 4),
-    }
+                moreIdentities = setOf(notary.identity),
-
+                moreKeys = emptySet()
-    @After
+        )
-    fun cleanUp() {
+        services = mockServices
-        mockNet.stopNodes()
+        database = db
     }
 
     // AbstractPartyToX500NameAsStringConverter could cause circular flush of Hibernate session because it is invoked during flush, and a
@@ -68,20 +75,24 @@ class HibernateColumnConverterTests {
 
         // Create parallel set of key and identity services so that the values are not cached, forcing the node caches to do a lookup.
         val identityService = PersistentIdentityService(TestingNamedCacheFactory())
-        val originalIdentityService: PersistentIdentityService = bankOfCordaNode.services.identityService as PersistentIdentityService
+        val originalIdentityService: PersistentIdentityService = services.identityService as PersistentIdentityService
         identityService.database = originalIdentityService.database
         identityService.start(originalIdentityService.trustRoot)
         val keyService = E2ETestKeyManagementService(identityService)
-        keyService.start((bankOfCordaNode.services.keyManagementService as E2ETestKeyManagementService).keyPairs)
+        keyService.start(setOf(myself.keyPair))
 
         // New identity for a notary (doesn't matter that it's for Bank Of Corda... since not going to use it as an actual notary etc).
-        val newKeyAndCert = keyService.freshKeyAndCert(bankOfCordaNode.info.legalIdentitiesAndCerts[0], false)
+        val newKeyAndCert = keyService.freshKeyAndCert(services.myInfo.legalIdentitiesAndCerts[0], false)
-        val randomNotary = Party(BOC_NAME, newKeyAndCert.owningKey)
+        val randomNotary = Party(myself.name, newKeyAndCert.owningKey)
 
-        val future = bankOfCordaNode.startFlow(TestCashIssueFlow(expected, ref, randomNotary))
+        val ourIdentity = services.myInfo.legalIdentities.first()
-        mockNet.runNetwork()
+        val builder = TransactionBuilder(notary.party)
-        val issueTx = future.getOrThrow().stx
+        val issuer = services.myInfo.legalIdentities.first().ref(ref)
-        val output = issueTx.tx.outputsOfType<Cash.State>().single()
+        val signers = Cash().generateIssue(builder, expected.issuedBy(issuer), ourIdentity, randomNotary)
-        assertEquals(expected.`issued by`(bankOfCorda.ref(ref)), output.amount)
+        val tx: SignedTransaction = services.signInitialTransaction(builder, signers)
+        services.recordTransactions(tx)
+
+        val output = tx.tx.outputsOfType<Cash.State>().single()
+        assertEquals(expected.`issued by`(ourIdentity.ref(ref)), output.amount)
     }
 }

node/src/test/kotlin/net/corda/node/services/vault/ExternalIdMappingTest.kt

@@ -1,8 +1,5 @@
 package net.corda.node.services.vault
 
-import com.nhaarman.mockito_kotlin.doReturn
-import com.nhaarman.mockito_kotlin.mock
-import com.nhaarman.mockito_kotlin.whenever
 import net.corda.core.identity.AbstractParty
 import net.corda.core.identity.AnonymousParty
 import net.corda.core.identity.CordaX500Name
@@ -10,15 +7,12 @@ import net.corda.core.node.services.queryBy
 import net.corda.core.node.services.vault.QueryCriteria
 import net.corda.core.node.services.vault.builder
 import net.corda.core.transactions.TransactionBuilder
-import net.corda.node.services.api.IdentityServiceInternal
-import net.corda.node.services.keys.BasicHSMKeyManagementService
 import net.corda.nodeapi.internal.persistence.CordaPersistence
 import net.corda.testing.common.internal.testNetworkParameters
 import net.corda.testing.contracts.DummyContract
 import net.corda.testing.contracts.DummyState
 import net.corda.testing.core.SerializationEnvironmentRule
 import net.corda.testing.core.TestIdentity
-import net.corda.testing.internal.rigorousMock
 import net.corda.testing.node.MockServices
 import org.junit.Before
 import org.junit.Rule
@@ -45,41 +39,17 @@ class ExternalIdMappingTest {
 
     @Before
     fun setUp() {
-        val (db, mockServices) = MockServices.makeTestDatabaseAndMockServices(
+        val (db, mockServices) = MockServices.makeTestDatabaseAndPersistentServices(
                 cordappPackages = cordapps,
-                identityService = mock<IdentityServiceInternal>().also {
-                    doReturn(notary.party).whenever(it).partyFromKey(notary.publicKey)
-                    doReturn(notary.party).whenever(it).wellKnownPartyFromAnonymous(notary.party)
-                    doReturn(notary.party).whenever(it).wellKnownPartyFromX500Name(notary.name)
-                },
                 initialIdentity = myself,
-                networkParameters = testNetworkParameters(minimumPlatformVersion = 4)
+                networkParameters = testNetworkParameters(minimumPlatformVersion = 4),
+                moreIdentities = setOf(notary.identity),
+                moreKeys = emptySet()
         )
         services = mockServices
         database = db
     }
 
-    private fun freshKeyForExternalId(externalId: UUID): AnonymousParty {
-        val anonymousParty = freshKey()
-        database.transaction {
-            services.withEntityManager {
-                val mapping = BasicHSMKeyManagementService.PublicKeyHashToExternalId(externalId, anonymousParty.owningKey)
-                persist(mapping)
-                flush()
-            }
-        }
-        return anonymousParty
-    }
-
-    private fun freshKey(): AnonymousParty {
-        val key = services.keyManagementService.freshKey()
-        val anonymousParty = AnonymousParty(key)
-        // Add behaviour to the mock identity management service for dealing with the new key.
-        // It won't be able to resolve it as it's just an anonymous key that is not linked to an identity.
-        services.identityService.also { doReturn(null).whenever(it).wellKnownPartyFromAnonymous(anonymousParty) }
-        return anonymousParty
-    }
-
     private fun createDummyState(participants: List<AbstractParty>): DummyState {
         val tx = TransactionBuilder(notary = notary.party).apply {
             addOutputState(DummyState(1, participants), DummyContract.PROGRAM_ID)
@@ -95,11 +65,11 @@ class ExternalIdMappingTest {
         val vaultService = services.vaultService
         // Create new external ID and two keys mapped to it.
         val id = UUID.randomUUID()
-        val keyOne = freshKeyForExternalId(id)
+        val keyOne = services.keyManagementService.freshKeyAndCert(myself.identity, false, id)
-        val keyTwo = freshKeyForExternalId(id)
+        val keyTwo = services.keyManagementService.freshKeyAndCert(myself.identity, false, id)
         // Create states with a public key assigned to the new external ID.
-        val dummyStateOne = createDummyState(listOf(keyOne))
+        val dummyStateOne = createDummyState(listOf(AnonymousParty(keyOne.owningKey)))
-        val dummyStateTwo = createDummyState(listOf(keyTwo))
+        val dummyStateTwo = createDummyState(listOf(AnonymousParty(keyTwo.owningKey)))
         // This query should return two states!
         val result = database.transaction {
             val externalId = builder { VaultSchemaV1.StateToExternalId::externalId.`in`(listOf(id)) }
@@ -122,11 +92,11 @@ class ExternalIdMappingTest {
         val vaultService = services.vaultService
         // Create new external ID.
         val idOne = UUID.randomUUID()
-        val keyOne = freshKeyForExternalId(idOne)
+        val keyOne = services.keyManagementService.freshKeyAndCert(myself.identity, false, idOne)
         val idTwo = UUID.randomUUID()
-        val keyTwo = freshKeyForExternalId(idTwo)
+        val keyTwo = services.keyManagementService.freshKeyAndCert(myself.identity, false, idTwo)
         // Create state with a public key assigned to the new external ID.
-        val dummyState = createDummyState(listOf(keyOne, keyTwo))
+        val dummyState = createDummyState(listOf(AnonymousParty(keyOne.owningKey), AnonymousParty(keyTwo.owningKey)))
         // This query should return one state!
         val result = database.transaction {
             val externalId = builder { VaultSchemaV1.StateToExternalId::externalId.`in`(listOf(idOne, idTwo)) }

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

@@ -26,6 +26,8 @@ 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.schema.NodeSchemaService
 import net.corda.node.services.transactions.InMemoryTransactionVerifierService
 import net.corda.node.services.vault.NodeVaultService
@@ -36,6 +38,7 @@ import net.corda.testing.common.internal.testNetworkParameters
 import net.corda.testing.core.TestIdentity
 import net.corda.testing.internal.DEV_ROOT_CA
 import net.corda.testing.internal.MockCordappProvider
+import net.corda.testing.internal.TestingNamedCacheFactory
 import net.corda.testing.internal.configureDatabase
 import net.corda.testing.node.internal.*
 import net.corda.testing.services.MockAttachmentStorage
@@ -71,7 +74,8 @@ open class MockServices private constructor(
         override val identityService: IdentityService,
         private val initialNetworkParameters: NetworkParameters,
         private val initialIdentity: TestIdentity,
-        private val moreKeys: Array<out KeyPair>
+        private val moreKeys: Array<out KeyPair>,
+        override val keyManagementService: KeyManagementService = MockKeyManagementService(identityService, *arrayOf(initialIdentity.keyPair) + moreKeys)
 ) : ServiceHub {
 
     companion object {
@@ -116,32 +120,91 @@ 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 mockService = database.transaction {
-                object : MockServices(cordappLoader, identityService, networkParameters, initialIdentity, moreKeys) {
+                makeMockMockServices(cordappLoader, identityService, networkParameters, initialIdentity, moreKeys.toSet(), keyManagementService, schemaService, database)
-                    override var networkParametersService: NetworkParametersService = MockNetworkParametersStorage(networkParameters)
-                    override val vaultService: VaultService = makeVaultService(schemaService, database, cordappLoader)
-                    override fun recordTransactions(statesToRecord: StatesToRecord, txs: Iterable<SignedTransaction>) {
-                        ServiceHubInternal.recordTransactions(statesToRecord, txs,
-                                validatedTransactions as WritableTransactionStorage,
-                                mockStateMachineRecordedTransactionMappingStorage,
-                                vaultService as VaultServiceInternal,
-                                database)
-                    }
-
-                    override fun jdbcSession(): Connection = database.createSession()
-
-                    override fun <T : Any> withEntityManager(block: EntityManager.() -> T): T {
-                        return block(contextTransaction.restrictedEntityManager)
-                    }
-
-                    override fun withEntityManager(block: Consumer<EntityManager>) {
-                        return block.accept(contextTransaction.restrictedEntityManager)
-                    }
-                }
             }
             return Pair(database, mockService)
         }
 
+        /**
+         * Makes database and persistent services appropriate for unit tests which require persistence across the vault, identity service
+         * and key managment service.
+         *
+         * @param cordappPackages A [List] of cordapp packages to scan for any cordapp code, e.g. contract verification code, flows and services.
+         * @param initialIdentity The first (typically sole) identity the services will represent.
+         * @param moreKeys A list of additional [KeyPair] instances to be used by [MockServices].
+         * @param moreIdentities A list of additional [KeyPair] instances to be used by [MockServices].
+         * @return A pair where the first element is the instance of [CordaPersistence] and the second is [MockServices].
+         */
+        @JvmStatic
+        @JvmOverloads
+        fun makeTestDatabaseAndPersistentServices(
+                cordappPackages: List<String>,
+                initialIdentity: TestIdentity,
+                networkParameters: NetworkParameters = testNetworkParameters(modifiedTime = Instant.MIN),
+                moreKeys: Set<KeyPair>,
+                moreIdentities: Set<PartyAndCertificate>
+        ): Pair<CordaPersistence, MockServices> {
+            val cordappLoader = cordappLoaderForPackages(cordappPackages)
+            val dataSourceProps = makeTestDataSourceProperties()
+            val schemaService = NodeSchemaService(cordappLoader.cordappSchemas)
+            val identityService = PersistentIdentityService(TestingNamedCacheFactory())
+            val persistence = configureDatabase(dataSourceProps, DatabaseConfig(), identityService::wellKnownPartyFromX500Name, identityService::wellKnownPartyFromAnonymous, schemaService, schemaService.internalSchemas())
+
+            // Create a persistent identity service and add all the supplied identities.
+            identityService.apply {
+                ourNames = setOf(initialIdentity.name)
+                database = persistence
+                start(DEV_ROOT_CA.certificate)
+                persistence.transaction { identityService.loadIdentities(moreIdentities + initialIdentity.identity) }
+            }
+
+            // 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)
+            persistence.transaction { keyManagementService.start(moreKeys + initialIdentity.keyPair) }
+
+            val mockService = persistence.transaction {
+                makeMockMockServices(cordappLoader, identityService, networkParameters, initialIdentity, moreKeys, keyManagementService, schemaService, persistence)
+            }
+            return Pair(persistence, mockService)
+        }
+
+        private fun makeMockMockServices(
+                cordappLoader: CordappLoader,
+                identityService: IdentityService,
+                networkParameters: NetworkParameters,
+                initialIdentity: TestIdentity,
+                moreKeys: Set<KeyPair>,
+                keyManagementService: KeyManagementService,
+                schemaService: SchemaService,
+                persistence: CordaPersistence
+        ): MockServices {
+            return object : MockServices(cordappLoader, identityService, networkParameters, initialIdentity, moreKeys.toTypedArray(), keyManagementService) {
+                override var networkParametersService: NetworkParametersService = MockNetworkParametersStorage(networkParameters)
+                override val vaultService: VaultService = makeVaultService(schemaService, persistence, cordappLoader)
+                override fun recordTransactions(statesToRecord: StatesToRecord, txs: Iterable<SignedTransaction>) {
+                    ServiceHubInternal.recordTransactions(statesToRecord, txs,
+                            validatedTransactions as WritableTransactionStorage,
+                            mockStateMachineRecordedTransactionMappingStorage,
+                            vaultService as VaultServiceInternal,
+                            persistence)
+                }
+
+                override fun jdbcSession(): Connection = persistence.createSession()
+
+                override fun <T : Any> withEntityManager(block: EntityManager.() -> T): T {
+                    return block(contextTransaction.restrictedEntityManager)
+                }
+
+                override fun withEntityManager(block: Consumer<EntityManager>) {
+                    return block.accept(contextTransaction.restrictedEntityManager)
+                }
+            }
+        }
+
         // Because Kotlin is dumb and makes not publicly visible objects public, thus changing the public API.
         private val mockStateMachineRecordedTransactionMappingStorage = MockStateMachineRecordedTransactionMappingStorage()
 
@@ -175,8 +238,18 @@ open class MockServices private constructor(
     }
 
     private constructor(cordappLoader: CordappLoader, identityService: IdentityService, networkParameters: NetworkParameters,
-                        initialIdentity: TestIdentity, moreKeys: Array<out KeyPair>)
+                        initialIdentity: TestIdentity, moreKeys: Array<out KeyPair>, keyManagementService: KeyManagementService)
-            : this(cordappLoader, MockTransactionStorage(), identityService, networkParameters, initialIdentity, moreKeys)
+            : this(cordappLoader, MockTransactionStorage(), identityService, networkParameters, initialIdentity, moreKeys, keyManagementService)
+
+    private constructor(cordappLoader: CordappLoader, identityService: IdentityService, networkParameters: NetworkParameters,
+                        initialIdentity: TestIdentity, moreKeys: Array<out KeyPair>) : this(
+            cordappLoader,
+            MockTransactionStorage(),
+            identityService,
+            networkParameters,
+            initialIdentity,
+            moreKeys
+    )
 
     /**
      * Create a mock [ServiceHub] that looks for app code in the given package names, uses the provided identity service
@@ -196,6 +269,14 @@ open class MockServices private constructor(
                 vararg moreKeys: KeyPair) :
             this(cordappLoaderForPackages(cordappPackages), identityService, networkParameters, initialIdentity, moreKeys)
 
+    constructor(cordappPackages: Iterable<String>,
+                initialIdentity: TestIdentity,
+                identityService: IdentityService,
+                networkParameters: NetworkParameters,
+                vararg moreKeys: KeyPair,
+                keyManagementService: KeyManagementService) :
+            this(cordappLoaderForPackages(cordappPackages), identityService, networkParameters, initialIdentity, moreKeys, keyManagementService)
+
     /**
      * Create a mock [ServiceHub] that looks for app code in the given package names, uses the provided identity service
      * (you can get one from [makeTestIdentityService]) and represents the given identity.
@@ -298,7 +379,6 @@ open class MockServices private constructor(
         get() = networkParametersService.run { lookup(currentHash)!! }
 
     final override val attachments = MockAttachmentStorage()
-    override val keyManagementService: KeyManagementService by lazy { MockKeyManagementService(identityService, *arrayOf(initialIdentity.keyPair) + moreKeys) }
     override val vaultService: VaultService get() = throw UnsupportedOperationException()
     override val contractUpgradeService: ContractUpgradeService get() = throw UnsupportedOperationException()
     override val networkMapCache: NetworkMapCache get() = throw UnsupportedOperationException()

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

@@ -34,7 +34,7 @@ import net.corda.node.services.api.ServiceHubInternal
 import net.corda.node.services.api.StartedNodeServices
 import net.corda.node.services.config.*
 import net.corda.node.services.identity.PersistentIdentityService
-import net.corda.node.services.keys.E2ETestKeyManagementService
+import net.corda.node.services.keys.BasicHSMKeyManagementService
 import net.corda.node.services.keys.KeyManagementServiceInternal
 import net.corda.node.services.keys.cryptoservice.BCCryptoService
 import net.corda.node.services.messaging.Message
@@ -373,7 +373,7 @@ open class InternalMockNetwork(cordappPackages: List<String> = emptyList(),
         }
 
         override fun makeKeyManagementService(identityService: PersistentIdentityService): KeyManagementServiceInternal {
-            return E2ETestKeyManagementService(identityService, cryptoService)
+            return BasicHSMKeyManagementService(cacheFactory, identityService, database, cryptoService)
         }
 
         override fun startShell() {

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

@@ -11,6 +11,7 @@ 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]
@@ -25,12 +26,31 @@ class MockKeyManagementService(val identityService: IdentityService,
 
     private val nextKeys = LinkedList<KeyPair>()
 
+    val keysById: MutableMap<UUID, Set<PublicKey>> = ConcurrentHashMap()
+
     override fun freshKey(): PublicKey {
         val k = nextKeys.poll() ?: generateKeyPair()
         keyStore[k.public] = k.private
         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 {