CORDA-3485 Restore CollectSignaturesFlow.kt to allow multiple collections from Well Known sessions (#5800)

* modify CollectSignaturesFlow.kt to allow multiple collections from wellKnown party initiated sessions

* detekt fixes

* review comments

* move require lambdas back outside of the function definition of requires

* address review comments

* fix detekt

* fix api scanner

.ci/api-current.txt

@@ -2876,7 +2876,7 @@ public static final class net.corda.core.flows.WithReferencedStatesFlow$Companio
 ##
 @DoNotImplement
 @CordaSerializable
-public abstract class net.corda.core.identity.AbstractParty extends java.lang.Object
+public abstract class net.corda.core.identity.AbstractParty extends java.lang.Object implements net.corda.core.flows.Destination
   public <init>(java.security.PublicKey)
   public boolean equals(Object)
   @NotNull

core-tests/src/test/kotlin/net/corda/coretests/flows/CollectSignaturesFlowTests.kt

@@ -5,13 +5,31 @@ import com.natpryce.hamkrest.assertion.assertThat
 import net.corda.core.contracts.Command
 import net.corda.core.contracts.StateAndContract
 import net.corda.core.contracts.requireThat
-import net.corda.core.flows.*
-import net.corda.core.identity.*
+import net.corda.core.flows.CollectSignaturesFlow
+import net.corda.core.flows.Destination
+import net.corda.core.flows.FinalityFlow
+import net.corda.core.flows.FlowLogic
+import net.corda.core.flows.FlowSession
+import net.corda.core.flows.InitiatedBy
+import net.corda.core.flows.InitiatingFlow
+import net.corda.core.flows.ReceiveFinalityFlow
+import net.corda.core.flows.SignTransactionFlow
+import net.corda.core.identity.AnonymousParty
+import net.corda.core.identity.CordaX500Name
+import net.corda.core.identity.Party
+import net.corda.core.identity.PartyAndCertificate
+import net.corda.core.identity.excludeHostNode
+import net.corda.core.identity.groupAbstractPartyByWellKnownParty
 import net.corda.core.node.services.IdentityService
 import net.corda.core.transactions.SignedTransaction
 import net.corda.core.transactions.TransactionBuilder
+import net.corda.core.utilities.getOrThrow
 import net.corda.testing.contracts.DummyContract
-import net.corda.testing.core.*
+import net.corda.testing.core.ALICE_NAME
+import net.corda.testing.core.BOB_NAME
+import net.corda.testing.core.CHARLIE_NAME
+import net.corda.testing.core.TestIdentity
+import net.corda.testing.core.singleIdentity
 import net.corda.testing.internal.matchers.flow.willReturn
 import net.corda.testing.internal.matchers.flow.willThrow
 import net.corda.testing.internal.rigorousMock
@@ -109,6 +127,36 @@ class CollectSignaturesFlowTests : WithContracts {
         Assert.assertThat(missingSigners, `is`(emptySet()))
     }
 
+    @Test(expected = IllegalArgumentException::class)
+    fun `throws exception when extra sessions are initiated`() {
+        bobNode.registerInitiatedFlow(ExtraSessionsFlowResponder::class.java)
+        charlieNode.registerInitiatedFlow(ExtraSessionsFlowResponder::class.java)
+        val future = aliceNode.startFlow(ExtraSessionsFlow(
+                listOf(
+                        bobNode.info.singleIdentity(),
+                        charlieNode.info.singleIdentity()
+                ),
+                listOf(bobNode.info.singleIdentity(), alice)))
+                .resultFuture
+        mockNet.runNetwork()
+        future.getOrThrow()
+    }
+
+    @Test
+    fun `it is possible to collect from multiple well known sessions`() {
+        bobNode.registerInitiatedFlow(ExtraSessionsFlowResponder::class.java)
+        charlieNode.registerInitiatedFlow(ExtraSessionsFlowResponder::class.java)
+        val future = aliceNode.startFlow(ExtraSessionsFlow(listOf(
+                bobNode.info.singleIdentity(),
+                bobNode.info.singleIdentity(),
+                bobNode.info.singleIdentity(),
+                bobNode.info.singleIdentity()),
+                listOf(bobNode.info.singleIdentity(), alice))).resultFuture
+        mockNet.runNetwork()
+        val signedTx = future.getOrThrow()
+        Assert.assertThat(signedTx.getMissingSigners(), `is`(emptySet()))
+    }
+
     @Test
     fun `no need to collect any signatures`() {
         val ptx = aliceNode.signDummyContract(alice.ref(1))
@@ -270,3 +318,33 @@ class MixAndMatchAnonymousSessionTestFlowResponder(private val otherSideSession:
         subFlow(signFlow)
     }
 }
+
+@InitiatingFlow
+class ExtraSessionsFlow(private val openFor: List<Party>, private val involve: List<Party>) : FlowLogic<SignedTransaction>() {
+    @Suspendable
+    override fun call(): SignedTransaction {
+
+        val sessions = openFor.map { initiateFlow(it) }
+        val state = DummyContract.MultiOwnerState(owners = involve.map { AnonymousParty(it.owningKey) })
+        val create = net.corda.testing.contracts.DummyContract.Commands.Create()
+        val txBuilder = TransactionBuilder(notary = serviceHub.networkMapCache.notaryIdentities.first())
+                .addOutputState(state)
+                .addCommand(create, involve.map { it.owningKey })
+
+        val signedByUsTx = serviceHub.signInitialTransaction(txBuilder)
+        return subFlow(CollectSignaturesFlow(signedByUsTx, sessions))
+    }
+}
+
+@InitiatedBy(ExtraSessionsFlow::class)
+class ExtraSessionsFlowResponder(private val otherSideSession: FlowSession) : FlowLogic<Unit>() {
+    @Suspendable
+    override fun call() {
+        val signFlow = object : SignTransactionFlow(otherSideSession) {
+            @Suspendable
+            override fun checkTransaction(stx: SignedTransaction) = requireThat {
+            }
+        }
+        subFlow(signFlow)
+    }
+}

core/src/main/kotlin/net/corda/core/flows/CollectSignaturesFlow.kt

@@ -4,10 +4,10 @@ import co.paralleluniverse.fibers.Suspendable
 import net.corda.core.crypto.TransactionSignature
 import net.corda.core.crypto.isFulfilledBy
 import net.corda.core.crypto.toStringShort
+import net.corda.core.identity.AbstractParty
 import net.corda.core.identity.AnonymousParty
 import net.corda.core.identity.Party
 import net.corda.core.identity.groupPublicKeysByWellKnownParty
-import net.corda.core.internal.toMultiMap
 import net.corda.core.node.ServiceHub
 import net.corda.core.transactions.SignedTransaction
 import net.corda.core.transactions.WireTransaction
@@ -112,63 +112,72 @@ class CollectSignaturesFlow @JvmOverloads constructor(val partiallySignedTx: Sig
         // If the unsigned counterparties list is empty then we don't need to collect any more signatures here.
         if (unsigned.isEmpty()) return partiallySignedTx
 
-        val setOfAllSessionKeys: Map<PublicKey, FlowSession> = sessionsToCollectFrom
-                .groupBy {
-                    val destination = it.destination
-                    when (destination) {
-                        is Party -> destination.owningKey
-                        is AnonymousParty -> destination.owningKey
-                        else -> throw IllegalArgumentException("Signatures can only be collected from Party or AnonymousParty, not $destination")
-                    }
-                }
-                .mapValues {
-                    require(it.value.size == 1) { "There are multiple sessions initiated for party key ${it.key.toStringShort()}" }
-                    it.value.first()
+        val wellKnownSessions = sessionsToCollectFrom.filter { it.destination is Party }
+        val anonymousSessions = sessionsToCollectFrom.filter { it.destination is AnonymousParty }
+
+        require(wellKnownSessions.size + anonymousSessions.size == sessionsToCollectFrom.size) {
+            "Unrecognized Destination type used to initiate a flow session"
         }
 
-        val partyToKeysItSignsFor: Map<Party, List<PublicKey>> = groupPublicKeysByWellKnownParty(serviceHub, unsigned)
-        val keyToSigningParty: Map<PublicKey, Party> = partyToKeysItSignsFor
-                .flatMap { (wellKnown, allKeysItSignsFor) -> allKeysItSignsFor.map { it to wellKnown } }
-                .toMap()
+        val wellKnownPartyToSessionMap: Map<Party, List<FlowSession>> = wellKnownSessions.groupBy { (it.destination as Party) }
+        val anonymousPartyToSessionMap: Map<AnonymousParty, List<FlowSession>> = anonymousSessions
+                .groupBy { (it.destination as AnonymousParty) }
 
-        val unrelatedSessions = sessionsToCollectFrom.filterNot {
-            if (it.destination is Party) {
-                // The session must have a corresponding unsigned.
-                it.destination in partyToKeysItSignsFor
-            } else {
-                // setOfAllSessionKeys has already checked for valid destination types so we can safely cast to AnonoymousParty here.
-                // This session was not initiated by a wellKnownParty so must directly exist in the unsigned.
-                (it.destination as AnonymousParty).owningKey in unsigned
+        //check that there is at most one session for each not well known party
+        for (entry in anonymousPartyToSessionMap) {
+            require(entry.value.size == 1) {
+                "There are multiple sessions initiated for Anonymous Party ${entry.key.owningKey.toStringShort()}"
             }
         }
 
-        val keyToSessionList = unsigned.map {
-            val session = setOfAllSessionKeys[it]
-            if (session != null) {
-                // The unsigned key exists directly as a sessionKey, so use that session
-                it to session
-            } else {
-                // It might be delegated to a wellKnownParty
-                val wellKnownParty = checkNotNull(keyToSigningParty[it]) { "Could not find a session or wellKnown party for key ${it.toStringShort()}" }
-                // There is a wellKnownParty for this key, check if it has a session, and if so - use that session
-                val sessionForWellKnownParty = checkNotNull(setOfAllSessionKeys[wellKnownParty.owningKey]) {
-                    "No session available to request signature for key: ${it.toStringShort()}"
+        //all keys that were used to initate a session must be sent to that session
+        val keysToSendToAnonymousSessions: Set<PublicKey> = unsigned.intersect(anonymousPartyToSessionMap.keys.map { it.owningKey })
+
+        //all keys that are left over MUST map back to a
+        val keysThatMustMapToAWellKnownSession: Set<PublicKey> = unsigned - keysToSendToAnonymousSessions
+        //if a key does not have a well known identity associated with it, it does not map to a wellKnown session
+        val keysThatDoNotMapToAWellKnownSession: List<PublicKey> = keysThatMustMapToAWellKnownSession
+                .filter { serviceHub.identityService.wellKnownPartyFromAnonymous(AnonymousParty(it)) == null }
+        //ensure that no keys are impossible to map to a session
+        require(keysThatDoNotMapToAWellKnownSession.isEmpty()) {
+            " Unable to match key(s): $keysThatDoNotMapToAWellKnownSession to a session to collect signatures from"
         }
-                it to sessionForWellKnownParty
+
+        //we now know that all the keys are either related to a specific session due to being used as a Destination for that session
+        //OR map back to a wellKnown party
+        //now we must check that each wellKnown party has a session passed for it
+        val groupedByPartyKeys = groupPublicKeysByWellKnownParty(serviceHub, keysThatMustMapToAWellKnownSession)
+        for (entry in groupedByPartyKeys) {
+            require(wellKnownPartyToSessionMap.contains(entry.key)) {
+                "${entry.key} is a required signer, but no session has been passed in for them"
             }
         }
 
-        // Now invert the map to find the keys per session
-        val sessionToKeysMap = keyToSessionList.map { it.second to it.first }.toMultiMap()
+        //so we now know that all keys are linked to a session in some way
+        //we need to check that there are no extra sessions
+        val extraNotWellKnownSessions = anonymousSessions.filterNot { (it.destination as AnonymousParty).owningKey in unsigned }
+        val extraWellKnownSessions = wellKnownSessions.filterNot { it.counterparty in groupedByPartyKeys }
 
-        require(unrelatedSessions.isEmpty()) {
-            "The Initiator of CollectSignaturesFlow must pass in exactly the sessions required to sign the transaction."
+        require(extraNotWellKnownSessions.isEmpty() && extraWellKnownSessions.isEmpty()) {
+            "The Initiator of CollectSignaturesFlow must pass in exactly the sessions required to sign the transaction, " +
+                    "the following extra sessions were passed in: " +
+                    (extraWellKnownSessions.map { it.counterparty.name.toString() } +
+                            extraNotWellKnownSessions.map { (it.destination as AbstractParty).owningKey.toString() })
         }
-        // Collect signatures from all counterparties and append them to the partially signed transaction.
-        val counterpartySignatures = sessionToKeysMap.flatMap { (session, keys) ->
-            subFlow(CollectSignatureFlow(partiallySignedTx, session, keys))
+
+        //OK let's collect some signatures!
+
+        val sigsFromNotWellKnownSessions = anonymousSessions.flatMap { flowSession ->
+            //anonymous sessions will only ever sign for their own key
+            subFlow(CollectSignatureFlow(partiallySignedTx, flowSession, (flowSession.destination as AbstractParty).owningKey))
         }
-        val stx = partiallySignedTx + counterpartySignatures
+
+        val sigsFromWellKnownSessions = wellKnownSessions.flatMap { flowSession ->
+            val keysToAskThisSessionFor = groupedByPartyKeys[flowSession.counterparty] ?: emptyList()
+            subFlow(CollectSignatureFlow(partiallySignedTx, flowSession, keysToAskThisSessionFor))
+        }
+
+        val stx = partiallySignedTx + (sigsFromNotWellKnownSessions + sigsFromWellKnownSessions).toSet()
 
         // Verify all but the notary's signature if the transaction requires a notary, otherwise verify all signatures.
         progressTracker.currentStep = VERIFYING

core/src/main/kotlin/net/corda/core/identity/AbstractParty.kt

@@ -2,6 +2,7 @@ package net.corda.core.identity
 
 import net.corda.core.DoNotImplement
 import net.corda.core.contracts.PartyAndReference
+import net.corda.core.flows.Destination
 import net.corda.core.serialization.CordaSerializable
 import net.corda.core.utilities.OpaqueBytes
 import java.security.PublicKey
@@ -12,7 +13,7 @@ import java.security.PublicKey
  */
 @CordaSerializable
 @DoNotImplement
-abstract class AbstractParty(val owningKey: PublicKey) {
+abstract class AbstractParty(val owningKey: PublicKey): Destination {
     /** Anonymised parties do not include any detail apart from owning key, so equality is dependent solely on the key */
     override fun equals(other: Any?): Boolean = other === this || other is AbstractParty && other.owningKey == owningKey