CORDA-2920 Hash to Signature Constraint automatic propagation

Allow Hash Constraints to propagate to Signature Constraints. When
propagating, the new jar is added to the transaction instead of the
original contract attachment (jar).

The following requirements must be met to do so:
- System property "net.corda.node.disableHashConstraints" must be set
  to `true`
- The min platform version must be 4 or greater
- There must be an input state with a hash constraint
- There must be no output states with a hash constraint
- The new jar must be signed

If these requirements are not met, the original Hash
Constraint will be kept and the contract attachment related to it is
also used.

This transition is done at the start of `handleContract` as it is not
the normal path a transition would follow. It is considered a backdoor
and should be treated separately from the rest of the attachment and
constraint logic. Furthermore, it will only work in private network
since all nodes must set the special `disableHashConstraints` flag.

core-tests/src/test/kotlin/net/corda/coretests/transactions/TransactionBuilderTest.kt

@@ -146,7 +146,7 @@ class TransactionBuilderTest {
     }, DummyContract.PROGRAM_ID)
 
     private fun signedAttachment(vararg parties: Party) = ContractAttachment.create(object : AbstractAttachment({ byteArrayOf() }, "test") {
-        override val id: SecureHash get() = throw UnsupportedOperationException()
+        override val id: SecureHash get() = contractAttachmentId
 
         override val signerKeys: List<PublicKey> get() = parties.map { it.owningKey }
     }, DummyContract.PROGRAM_ID, signerKeys = parties.map { it.owningKey })

core/src/main/kotlin/net/corda/core/transactions/TransactionBuilder.kt

@@ -16,13 +16,10 @@ import net.corda.core.node.services.KeyManagementService
 import net.corda.core.serialization.SerializationContext
 import net.corda.core.serialization.SerializationFactory
 import net.corda.core.utilities.contextLogger
-import java.io.NotSerializableException
-import java.lang.Exception
 import java.security.PublicKey
 import java.time.Duration
 import java.time.Instant
-import java.util.ArrayDeque
+import java.util.*
-import java.util.UUID
 import java.util.regex.Pattern
 import kotlin.collections.ArrayList
 import kotlin.collections.component1
@@ -273,16 +270,17 @@ open class TransactionBuilder(
         val refStateContractAttachments: List<AttachmentId> = referenceStateGroups
             .filterNot { it.key in allContracts }
             .map { refStateEntry ->
-                    selectAttachmentThatSatisfiesConstraints(true, refStateEntry.key, refStateEntry.value, emptySet(), services)
+                getInstalledContractAttachmentId(
+                    refStateEntry.key,
+                    refStateEntry.value,
+                    services
+                )
             }
 
-        val contractClassNameToInputStateRef: Map<ContractClassName, Set<StateRef>> = inputsWithTransactionState.map { Pair(it.state.contract, it.ref) }
-                .groupBy { it.first }.mapValues { it.value.map { e -> e.second }.toSet() }
-
         // For each contract, resolve the AutomaticPlaceholderConstraint, and select the attachment.
         val contractAttachmentsAndResolvedOutputStates: List<Pair<AttachmentId, List<TransactionState<ContractState>>?>> = allContracts.toSet()
                 .map { ctr ->
-                    handleContract(ctr, inputContractGroups[ctr], contractClassNameToInputStateRef[ctr], outputContractGroups[ctr], explicitAttachmentContractsMap[ctr], services)
+                    handleContract(ctr, inputContractGroups[ctr], outputContractGroups[ctr], explicitAttachmentContractsMap[ctr], services)
                 }
 
         val resolvedStates: List<TransactionState<ContractState>> = contractAttachmentsAndResolvedOutputStates.mapNotNull { it.second }
@@ -316,13 +314,45 @@ open class TransactionBuilder(
     private fun handleContract(
             contractClassName: ContractClassName,
             inputStates: List<TransactionState<ContractState>>?,
-            inputStateRefs: Set<StateRef>?,
             outputStates: List<TransactionState<ContractState>>?,
             explicitContractAttachment: AttachmentId?,
             services: ServicesForResolution
     ): Pair<AttachmentId, List<TransactionState<ContractState>>?> {
         val inputsAndOutputs = (inputStates ?: emptyList()) + (outputStates ?: emptyList())
 
+        fun selectAttachment() = getInstalledContractAttachmentId(
+            contractClassName,
+            inputsAndOutputs.filterNot { it.constraint in automaticConstraints },
+            services
+        )
+
+        /*
+        This block handles the very specific code path where a [HashAttachmentConstraint] can
+        migrate to a [SignatureAttachmentConstraint]. If all the criteria is met, this function
+        will return early as the rest of the logic is no longer required.
+
+        This can only happen in a private network where all nodes have started with
+        a system parameter that disables the hash constraint check.
+        */
+        if (canMigrateFromHashToSignatureConstraint(inputStates, outputStates, services)) {
+            val attachmentId = selectAttachment()
+            val attachment = services.attachments.openAttachment(attachmentId)
+            require(attachment != null) { "Contract attachment $attachmentId for $contractClassName is missing." }
+            if ((attachment as ContractAttachment).isSigned && (explicitContractAttachment == null || explicitContractAttachment == attachment.id)) {
+                val signatureConstraint =
+                    makeSignatureAttachmentConstraint(attachment.signerKeys)
+                require(signatureConstraint.isSatisfiedBy(attachment)) { "Selected output constraint: $signatureConstraint not satisfying ${attachment.id}" }
+                val resolvedOutputStates = outputStates?.map {
+                    if (it.constraint in automaticConstraints) {
+                        it.copy(constraint = signatureConstraint)
+                    } else {
+                        it
+                    }
+                }
+                return attachment.id to resolvedOutputStates
+            }
+        }
+
         // Determine if there are any HashConstraints that pin the version of a contract. If there are, check if we trust them.
         val hashAttachments = inputsAndOutputs
                 .filter { it.constraint is HashAttachmentConstraint }
@@ -349,13 +379,6 @@ open class TransactionBuilder(
         // This will contain the hash of the JAR that *has* to be used by this Transaction, because it is explicit. Or null if none.
         val forcedAttachmentId = explicitContractAttachment ?: hashAttachments.singleOrNull()?.id
 
-        fun selectAttachment() = selectAttachmentThatSatisfiesConstraints(
-                false,
-                contractClassName,
-                inputsAndOutputs.filterNot { it.constraint in automaticConstraints },
-                inputStateRefs,
-                services)
-
         // This will contain the hash of the JAR that will be used by this Transaction.
         val selectedAttachmentId = forcedAttachmentId ?: selectAttachment()
 
@@ -402,6 +425,25 @@ open class TransactionBuilder(
         return Pair(selectedAttachmentId, resolvedOutputStates)
     }
 
+    /**
+     * Checks whether the current transaction can migrate from a [HashAttachmentConstraint] to a
+     * [SignatureAttachmentConstraint]. This is only possible in very specific scenarios. Most
+     * importantly, [HashAttachmentConstraint.disableHashConstraints] must be set to `true` for
+     * any possibility of transition off of existing [HashAttachmentConstraint]s.
+     */
+    private fun canMigrateFromHashToSignatureConstraint(
+        inputStates: List<TransactionState<ContractState>>?,
+        outputStates: List<TransactionState<ContractState>>?,
+        services: ServicesForResolution
+    ): Boolean {
+        return HashAttachmentConstraint.disableHashConstraints
+                && services.networkParameters.minimumPlatformVersion >= 4
+                // `disableHashConstraints == true` therefore it does not matter if there are
+                // multiple input states with different hash constraints
+                && inputStates?.any { it.constraint is HashAttachmentConstraint } == true
+                && outputStates?.none { it.constraint is HashAttachmentConstraint } == true
+    }
+
     /**
      * If there are multiple input states with different constraints then run the constraint intersection logic to determine the resulting output constraint.
      * For issuing transactions where the attachmentToUse is JarSigned, then default to the SignatureConstraint with all the signatures.
@@ -412,7 +454,7 @@ open class TransactionBuilder(
             inputStates: List<TransactionState<ContractState>>?,
             attachmentToUse: ContractAttachment,
             services: ServicesForResolution): AttachmentConstraint = when {
-        inputStates != null -> attachmentConstraintsTransition(inputStates.groupBy { it.constraint }.keys, attachmentToUse)
+        inputStates != null -> attachmentConstraintsTransition(inputStates.groupBy { it.constraint }.keys, attachmentToUse, services)
         attachmentToUse.signerKeys.isNotEmpty() && services.networkParameters.minimumPlatformVersion < 4 -> {
             log.warnOnce("Signature constraints not available on network requiring a minimum platform version of 4. Current is: ${services.networkParameters.minimumPlatformVersion}.")
             if (useWhitelistedByZoneAttachmentConstraint(contractClassName, services.networkParameters)) {
@@ -437,7 +479,8 @@ open class TransactionBuilder(
      */
     private fun attachmentConstraintsTransition(
         constraints: Set<AttachmentConstraint>,
-        attachmentToUse: ContractAttachment
+        attachmentToUse: ContractAttachment,
+        services: ServicesForResolution
     ): AttachmentConstraint = when {
 
         // Sanity check.
@@ -452,7 +495,8 @@ open class TransactionBuilder(
             throw IllegalArgumentException("Cannot mix HashConstraints with different hashes in the same transaction.")
 
         // The HashAttachmentConstraint is the strongest constraint, so it wins when mixed with anything. As long as the actual constraints pass.
-        // TODO - this could change if we decide to introduce a way to gracefully migrate from the Hash Constraint to the Signature Constraint.
+        // Migration from HashAttachmentConstraint to SignatureAttachmentConstraint is handled in [TransactionBuilder.handleContract]
+        // If we have reached this point, then no migration is possible and the existing HashAttachmentConstraint must be used
         constraints.any { it is HashAttachmentConstraint } -> constraints.find { it is HashAttachmentConstraint }!!
 
         // TODO, we don't currently support mixing signature constraints with different signers. This will change once we introduce third party signers.
@@ -460,14 +504,7 @@ open class TransactionBuilder(
             throw IllegalArgumentException("Cannot mix SignatureAttachmentConstraints signed by different parties in the same transaction.")
 
         // This ensures a smooth migration from a Whitelist Constraint to a Signature Constraint
-        constraints.any { it is WhitelistedByZoneAttachmentConstraint } && attachmentToUse.isSigned -> {
+        constraints.any { it is WhitelistedByZoneAttachmentConstraint } && attachmentToUse.isSigned && services.networkParameters.minimumPlatformVersion >= 4 -> transitionToSignatureConstraint(constraints, attachmentToUse)
-            val signatureConstraint = constraints.singleOrNull { it is SignatureAttachmentConstraint }
-            // If there were states transitioned already used in the current transaction use that signature constraint, otherwise create a new one.
-            when {
-                signatureConstraint != null -> signatureConstraint
-                else -> makeSignatureAttachmentConstraint(attachmentToUse.signerKeys)
-            }
-        }
 
         // This condition is hit when the current node has not installed the latest signed version but has already received states that have been migrated
         constraints.any { it is SignatureAttachmentConstraint } && !attachmentToUse.isSigned ->
@@ -479,17 +516,23 @@ open class TransactionBuilder(
         else -> throw IllegalArgumentException("Unexpected constraints $constraints.")
     }
 
+    private fun transitionToSignatureConstraint(constraints: Set<AttachmentConstraint>, attachmentToUse: ContractAttachment): SignatureAttachmentConstraint {
+        val signatureConstraint = constraints.singleOrNull { it is SignatureAttachmentConstraint } as? SignatureAttachmentConstraint
+        // If there were states transitioned already used in the current transaction use that signature constraint, otherwise create a new one.
+        return when {
+            signatureConstraint != null -> signatureConstraint
+            else -> makeSignatureAttachmentConstraint(attachmentToUse.signerKeys)
+        }
+    }
+
     private fun makeSignatureAttachmentConstraint(attachmentSigners: List<PublicKey>) =
             SignatureAttachmentConstraint(CompositeKey.Builder().addKeys(attachmentSigners).build())
 
-    /**
+    private fun getInstalledContractAttachmentId(
-     * This method should only be called for upgradeable contracts.
+        contractClassName: String,
-     */
+        states: List<TransactionState<ContractState>>,
-    private fun selectAttachmentThatSatisfiesConstraints(isReference: Boolean, contractClassName: String, states: List<TransactionState<ContractState>>, stateRefs: Set<StateRef>?, services: ServicesForResolution): AttachmentId {
+        services: ServicesForResolution
-        val constraints = states.map { it.constraint }
+    ): AttachmentId {
-        require(constraints.none { it in automaticConstraints })
-        require(isReference || constraints.none { it is HashAttachmentConstraint })
-
         return services.cordappProvider.getContractAttachmentID(contractClassName)
             ?: throw MissingContractAttachments(states, contractClassName)
     }

node/src/integration-test/kotlin/net/corda/contracts/SignatureConstraintVersioningTests.kt

@@ -7,10 +7,12 @@ import net.corda.core.contracts.*
 import net.corda.core.crypto.sha256
 import net.corda.core.flows.FlowLogic
 import net.corda.core.flows.StartableByRPC
+import net.corda.core.identity.AbstractParty
 import net.corda.core.identity.CordaX500Name
 import net.corda.core.identity.Party
 import net.corda.core.internal.*
 import net.corda.core.messaging.startFlow
+import net.corda.core.serialization.CordaSerializable
 import net.corda.core.transactions.CoreTransaction
 import net.corda.core.transactions.LedgerTransaction
 import net.corda.core.transactions.SignedTransaction
@@ -19,8 +21,6 @@ import net.corda.core.utilities.getOrThrow
 import net.corda.node.flows.isQuasarAgentSpecified
 import net.corda.node.services.Permissions.Companion.invokeRpc
 import net.corda.node.services.Permissions.Companion.startFlow
-import net.corda.testMessage.Message
-import net.corda.testMessage.MessageState
 import net.corda.testing.common.internal.testNetworkParameters
 import net.corda.testing.core.singleIdentity
 import net.corda.testing.driver.DriverDSL
@@ -46,6 +46,7 @@ class SignatureConstraintVersioningTests {
     private val oldUnsignedCordapp = baseUnsigned.copy(versionId = 2)
     private val oldCordapp = base.copy(versionId = 2)
     private val newCordapp = base.copy(versionId = 3)
+    private val newUnsignedCordapp = baseUnsigned.copy(versionId = 3)
     private val user = User("mark", "dadada", setOf(startFlow<CreateMessage>(), startFlow<ConsumeMessage>(), invokeRpc("vaultQuery")))
     private val message = Message("Hello world!")
     private val transformedMessage = Message(message.value + "A")
@@ -82,7 +83,7 @@ class SignatureConstraintVersioningTests {
                     page.states.singleOrNull()
                 }
                 CordaRPCClient(nodeHandle.rpcAddress).start(user.username, user.password).use {
-                    it.proxy.startFlow(::ConsumeMessage, result!!, defaultNotaryIdentity).returnValue.getOrThrow()
+                    it.proxy.startFlow(::ConsumeMessage, result!!, defaultNotaryIdentity, false, false).returnValue.getOrThrow()
                 }
                 result = CordaRPCClient(nodeHandle.rpcAddress).start(user.username, user.password).use {
                     val page = it.proxy.vaultQuery(MessageState::class.java)
@@ -100,18 +101,182 @@ class SignatureConstraintVersioningTests {
     @Test
     fun `auto migration from WhitelistConstraint to SignatureConstraint`() {
         assumeFalse(System.getProperty("os.name").toLowerCase().startsWith("win")) // See NodeStatePersistenceTests.kt.
-        val transaction =
+        val (issuanceTransaction, consumingTransaction) = upgradeCorDappBetweenTransactions(
-            upgradeCorDappBetweenTransactions(oldUnsignedCordapp, newCordapp, listOf(oldUnsignedCordapp, newCordapp))
+            cordapp = oldUnsignedCordapp,
-        assertEquals(1, transaction.outputs.size)
+            newCordapp = newCordapp,
-        assertTrue(transaction.outputs.single().constraint is SignatureAttachmentConstraint)
+            whiteListedCordapps = mapOf(
+                TEST_MESSAGE_CONTRACT_PROGRAM_ID to listOf(
+                    oldUnsignedCordapp,
+                    newCordapp
+                )
+            ),
+            systemProperties = emptyMap(),
+            startNodesInProcess = false
+        )
+        assertEquals(1, issuanceTransaction.outputs.size)
+        assertTrue(issuanceTransaction.outputs.single().constraint is WhitelistedByZoneAttachmentConstraint)
+        assertEquals(1, consumingTransaction.outputs.size)
+        assertTrue(consumingTransaction.outputs.single().constraint is SignatureAttachmentConstraint)
     }
 
     @Test
-    fun `auto migration from WhitelistConstraint to SignatureConstraint fail for not whitelisted signed JAR`() {
+    fun `WhitelistConstraint cannot be migrated to SignatureConstraint if platform version is not 4 or greater`() {
+        assumeFalse(System.getProperty("os.name").toLowerCase().startsWith("win")) // See NodeStatePersistenceTests.kt.
+        val (issuanceTransaction, consumingTransaction) = upgradeCorDappBetweenTransactions(
+            cordapp = oldUnsignedCordapp,
+            newCordapp = newCordapp,
+            whiteListedCordapps = mapOf(
+                TEST_MESSAGE_CONTRACT_PROGRAM_ID to listOf(
+                    oldUnsignedCordapp,
+                    newCordapp
+                )
+            ),
+            systemProperties = emptyMap(),
+            startNodesInProcess = false,
+            minimumPlatformVersion = 3
+        )
+        assertEquals(1, issuanceTransaction.outputs.size)
+        assertTrue(issuanceTransaction.outputs.single().constraint is WhitelistedByZoneAttachmentConstraint)
+        assertEquals(1, consumingTransaction.outputs.size)
+        assertTrue(consumingTransaction.outputs.single().constraint is WhitelistedByZoneAttachmentConstraint)
+    }
+
+    @Test
+    fun `WhitelistConstraint cannot be migrated to SignatureConstraint if signed JAR is not whitelisted`() {
         assumeFalse(System.getProperty("os.name").toLowerCase().startsWith("win")) // See NodeStatePersistenceTests.kt.
         assertThatExceptionOfType(CordaRuntimeException::class.java).isThrownBy {
-            upgradeCorDappBetweenTransactions(oldUnsignedCordapp, newCordapp, emptyList())
+            upgradeCorDappBetweenTransactions(
-        }.withMessageContaining("Selected output constraint: $WhitelistedByZoneAttachmentConstraint not satisfying")
+                cordapp = oldUnsignedCordapp,
+                newCordapp = newCordapp,
+                whiteListedCordapps = mapOf(TEST_MESSAGE_CONTRACT_PROGRAM_ID to emptyList()),
+                systemProperties = emptyMap(),
+                startNodesInProcess = true
+            )
+        }
+            .withMessageContaining("Selected output constraint: $WhitelistedByZoneAttachmentConstraint not satisfying")
+    }
+
+    @Test
+    fun `auto migration from WhitelistConstraint to SignatureConstraint will only transition states that do not have a constraint specified`() {
+        assumeFalse(System.getProperty("os.name").toLowerCase().startsWith("win")) // See NodeStatePersistenceTests.kt.
+        val (issuanceTransaction, consumingTransaction) = upgradeCorDappBetweenTransactions(
+            cordapp = oldUnsignedCordapp,
+            newCordapp = newCordapp,
+            whiteListedCordapps = mapOf(
+                TEST_MESSAGE_CONTRACT_PROGRAM_ID to listOf(
+                    oldUnsignedCordapp,
+                    newCordapp
+                )
+            ),
+            systemProperties = emptyMap(),
+            startNodesInProcess = true,
+            specifyExistingConstraint = true,
+            addAnotherAutomaticConstraintState = true
+        )
+        assertEquals(1, issuanceTransaction.outputs.size)
+        assertTrue(issuanceTransaction.outputs.single().constraint is WhitelistedByZoneAttachmentConstraint)
+        assertEquals(2, consumingTransaction.outputs.size)
+        assertTrue(consumingTransaction.outputs[0].constraint is WhitelistedByZoneAttachmentConstraint)
+        assertTrue(consumingTransaction.outputs[1].constraint is SignatureAttachmentConstraint)
+        assertEquals(
+            issuanceTransaction.outputs.single().constraint,
+            consumingTransaction.outputs.first().constraint,
+            "The constraint from the issuance transaction should be the same constraint used in the consuming transaction for the first state"
+        )
+    }
+
+    @Test
+    fun `auto migration from HashConstraint to SignatureConstraint`() {
+        assumeFalse(System.getProperty("os.name").toLowerCase().startsWith("win")) // See NodeStatePersistenceTests.kt.
+        val (issuanceTransaction, consumingTransaction) = upgradeCorDappBetweenTransactions(
+            cordapp = oldUnsignedCordapp,
+            newCordapp = newCordapp,
+            whiteListedCordapps = emptyMap(),
+            systemProperties = mapOf("net.corda.node.disableHashConstraints" to true.toString()),
+            startNodesInProcess = false
+        )
+        assertEquals(1, issuanceTransaction.outputs.size)
+        assertTrue(issuanceTransaction.outputs.single().constraint is HashAttachmentConstraint)
+        assertEquals(1, consumingTransaction.outputs.size)
+        assertTrue(consumingTransaction.outputs.single().constraint is SignatureAttachmentConstraint)
+    }
+
+    @Test
+    fun `HashConstraint cannot be migrated if 'disableHashConstraints' system property is not set to true`() {
+        assumeFalse(System.getProperty("os.name").toLowerCase().startsWith("win")) // See NodeStatePersistenceTests.kt.
+        val (issuanceTransaction, consumingTransaction) = upgradeCorDappBetweenTransactions(
+            cordapp = oldUnsignedCordapp,
+            newCordapp = newCordapp,
+            whiteListedCordapps = emptyMap(),
+            systemProperties = emptyMap(),
+            startNodesInProcess = false
+        )
+        assertEquals(1, issuanceTransaction.outputs.size)
+        assertTrue(issuanceTransaction.outputs.single().constraint is HashAttachmentConstraint)
+        assertEquals(1, consumingTransaction.outputs.size)
+        assertTrue(consumingTransaction.outputs.single().constraint is HashAttachmentConstraint)
+    }
+
+    @Test
+    fun `HashConstraint cannot be migrated to SignatureConstraint if new jar is not signed`() {
+        assumeFalse(System.getProperty("os.name").toLowerCase().startsWith("win")) // See NodeStatePersistenceTests.kt.
+        val (issuanceTransaction, consumingTransaction) = upgradeCorDappBetweenTransactions(
+            cordapp = oldUnsignedCordapp,
+            newCordapp = newUnsignedCordapp,
+            whiteListedCordapps = emptyMap(),
+            systemProperties = mapOf("net.corda.node.disableHashConstraints" to true.toString()),
+            startNodesInProcess = false
+        )
+        assertEquals(1, issuanceTransaction.outputs.size)
+        assertTrue(issuanceTransaction.outputs.single().constraint is HashAttachmentConstraint)
+        assertEquals(1, consumingTransaction.outputs.size)
+        assertTrue(consumingTransaction.outputs.single().constraint is HashAttachmentConstraint)
+    }
+
+    @Test
+    fun `HashConstraint cannot be migrated to SignatureConstraint if platform version is not 4 or greater`() {
+        assumeFalse(System.getProperty("os.name").toLowerCase().startsWith("win")) // See NodeStatePersistenceTests.kt.
+        val (issuanceTransaction, consumingTransaction) = upgradeCorDappBetweenTransactions(
+            cordapp = oldUnsignedCordapp,
+            newCordapp = newCordapp,
+            whiteListedCordapps = emptyMap(),
+            systemProperties = mapOf("net.corda.node.disableHashConstraints" to true.toString()),
+            startNodesInProcess = false,
+            minimumPlatformVersion = 3
+        )
+        assertEquals(1, issuanceTransaction.outputs.size)
+        assertTrue(issuanceTransaction.outputs.single().constraint is HashAttachmentConstraint)
+        assertEquals(1, consumingTransaction.outputs.size)
+        assertTrue(consumingTransaction.outputs.single().constraint is HashAttachmentConstraint)
+    }
+
+    @Test
+    fun `HashConstraint cannot be migrated to SignatureConstraint if a HashConstraint is specified for one state and another uses an AutomaticPlaceholderConstraint`() {
+        assumeFalse(System.getProperty("os.name").toLowerCase().startsWith("win")) // See NodeStatePersistenceTests.kt.
+        val (issuanceTransaction, consumingTransaction) = upgradeCorDappBetweenTransactions(
+            cordapp = oldUnsignedCordapp,
+            newCordapp = newCordapp,
+            whiteListedCordapps = emptyMap(),
+            systemProperties = mapOf("net.corda.node.disableHashConstraints" to true.toString()),
+            startNodesInProcess = false,
+            specifyExistingConstraint = true,
+            addAnotherAutomaticConstraintState = true
+        )
+        assertEquals(1, issuanceTransaction.outputs.size)
+        assertTrue(issuanceTransaction.outputs.single().constraint is HashAttachmentConstraint)
+        assertEquals(2, consumingTransaction.outputs.size)
+        assertTrue(consumingTransaction.outputs[0].constraint is HashAttachmentConstraint)
+        assertTrue(consumingTransaction.outputs[1].constraint is HashAttachmentConstraint)
+        assertEquals(
+            issuanceTransaction.outputs.single().constraint,
+            consumingTransaction.outputs.first().constraint,
+            "The constraint from the issuance transaction should be the same constraint used in the consuming transaction"
+        )
+        assertEquals(
+            consumingTransaction.outputs[0].constraint,
+            consumingTransaction.outputs[1].constraint,
+            "The AutomaticPlaceholderConstraint of the second state should become the same HashConstraint used in other state"
+        )
     }
 
     /**
@@ -121,44 +286,73 @@ class SignatureConstraintVersioningTests {
     private fun upgradeCorDappBetweenTransactions(
         cordapp: CustomCordapp,
         newCordapp: CustomCordapp,
-        whiteListedCordapps: List<CustomCordapp>
+        whiteListedCordapps: Map<ContractClassName, List<CustomCordapp>>,
-    ): CoreTransaction {
+        systemProperties: Map<String, String>,
+        startNodesInProcess: Boolean,
+        minimumPlatformVersion: Int = 4,
+        specifyExistingConstraint: Boolean = false,
+        addAnotherAutomaticConstraintState: Boolean = false
+    ): Pair<CoreTransaction, CoreTransaction> {
 
-        val attachmentHashes = whiteListedCordapps.map { Files.newInputStream(it.jarFile).readFully().sha256() }
+        val whitelistedAttachmentHashes = whiteListedCordapps.mapValues { (_, cordapps) ->
+            cordapps.map {
+                Files.newInputStream(it.jarFile).readFully().sha256()
+            }
+        }
 
         return internalDriver(
             inMemoryDB = false,
-            startNodesInProcess = isQuasarAgentSpecified(),
+            startNodesInProcess = startNodesInProcess,
             networkParameters = testNetworkParameters(
                 notaries = emptyList(),
-                minimumPlatformVersion = 4, whitelistedContractImplementations = mapOf(TEST_MESSAGE_CONTRACT_PROGRAM_ID to attachmentHashes)
+                minimumPlatformVersion = minimumPlatformVersion,
-            )
+                whitelistedContractImplementations = whitelistedAttachmentHashes
+            ),
+            systemProperties = systemProperties
         ) {
             // create transaction using first Cordapp
-            val (nodeName, baseDirectory) = createIssuanceTransaction(cordapp)
+            val (nodeName, baseDirectory, issuanceTransaction) = createIssuanceTransaction(cordapp)
             // delete the first cordapp
             deleteCorDapp(baseDirectory, cordapp)
             // create transaction using the upgraded cordapp resuing   input for transaction
-            createConsumingTransaction(nodeName, newCordapp).coreTransaction
+            val consumingTransaction = createConsumingTransaction(
+                nodeName,
+                newCordapp,
+                specifyExistingConstraint,
+                addAnotherAutomaticConstraintState
+            ).coreTransaction
+            issuanceTransaction to consumingTransaction
         }
     }
 
-    private fun DriverDSL.createIssuanceTransaction(cordapp: CustomCordapp): Pair<CordaX500Name, Path> {
+    private fun DriverDSL.createIssuanceTransaction(cordapp: CustomCordapp): Triple<CordaX500Name, Path, CoreTransaction> {
-        val nodeHandle = startNode(NodeParameters(rpcUsers = listOf(user), additionalCordapps = listOf(cordapp))).getOrThrow()
+        val nodeHandle = startNode(
+            NodeParameters(
+                rpcUsers = listOf(user),
+                additionalCordapps = listOf(cordapp)
+            )
+        ).getOrThrow()
         val nodeName = nodeHandle.nodeInfo.singleIdentity().name
-        CordaRPCClient(nodeHandle.rpcAddress).start(user.username, user.password).use {
+        val tx = CordaRPCClient(nodeHandle.rpcAddress).start(user.username, user.password).use {
-            it.proxy.startFlow(::CreateMessage, message, defaultNotaryIdentity).returnValue.getOrThrow()
+            it.proxy.startFlow(::CreateMessage, message, defaultNotaryIdentity)
+                .returnValue.getOrThrow().coreTransaction
         }
         nodeHandle.stop()
-        return Pair(nodeName, nodeHandle.baseDirectory)
+        return Triple(nodeName, nodeHandle.baseDirectory, tx)
     }
 
     private fun deleteCorDapp(baseDirectory: Path, cordapp: CustomCordapp) {
-        val cordappPath = baseDirectory.resolve(Paths.get("cordapps")).resolve(cordapp.jarFile.fileName)
+        val cordappPath =
+            baseDirectory.resolve(Paths.get("cordapps")).resolve(cordapp.jarFile.fileName)
         cordappPath.delete()
     }
 
-    private fun DriverDSL.createConsumingTransaction(nodeName: CordaX500Name, cordapp: CustomCordapp): SignedTransaction {
+    private fun DriverDSL.createConsumingTransaction(
+        nodeName: CordaX500Name,
+        cordapp: CustomCordapp,
+        specifyExistingConstraint: Boolean,
+        addAnotherAutomaticConstraintState: Boolean
+    ): SignedTransaction {
         val nodeHandle = startNode(
             NodeParameters(
                 providedName = nodeName,
@@ -166,12 +360,21 @@ class SignatureConstraintVersioningTests {
                 additionalCordapps = listOf(cordapp)
             )
         ).getOrThrow()
-        val result: StateAndRef<MessageState>? = CordaRPCClient(nodeHandle.rpcAddress).start(user.username, user.password).use {
+        val result: StateAndRef<MessageState>? =
+            CordaRPCClient(nodeHandle.rpcAddress).start(user.username, user.password).use {
                 val page = it.proxy.vaultQuery(MessageState::class.java)
                 page.states.singleOrNull()
             }
-        val transaction = CordaRPCClient(nodeHandle.rpcAddress).start(user.username, user.password).use {
+        val transaction =
-            it.proxy.startFlow(::ConsumeMessage, result!!, defaultNotaryIdentity).returnValue.getOrThrow()
+            CordaRPCClient(nodeHandle.rpcAddress).start(user.username, user.password).use {
+                it.proxy.startFlow(
+                    ::ConsumeMessage,
+                    result!!,
+                    defaultNotaryIdentity,
+                    specifyExistingConstraint,
+                    addAnotherAutomaticConstraintState
+                )
+                    .returnValue.getOrThrow()
             }
         nodeHandle.stop()
         return transaction
@@ -179,13 +382,21 @@ class SignatureConstraintVersioningTests {
 }
 
 @StartableByRPC
-class CreateMessage(private val message: Message, private val notary: Party) : FlowLogic<SignedTransaction>() {
+class CreateMessage(private val message: Message, private val notary: Party) :
+    FlowLogic<SignedTransaction>() {
     @Suspendable
     override fun call(): SignedTransaction {
         val messageState = MessageState(message = message, by = ourIdentity)
-        val txCommand = Command(DummyMessageContract.Commands.Send(), messageState.participants.map { it.owningKey })
+        val txCommand = Command(
-        val txBuilder = TransactionBuilder(notary).withItems(StateAndContract(messageState, TEST_MESSAGE_CONTRACT_PROGRAM_ID), txCommand)
+            DummyMessageContract.Commands.Send(),
-        txBuilder.toWireTransaction(serviceHub).toLedgerTransaction(serviceHub).verify()
+            messageState.participants.map { it.owningKey })
+        val txBuilder = TransactionBuilder(notary).withItems(
+            StateAndContract(
+                messageState,
+                TEST_MESSAGE_CONTRACT_PROGRAM_ID
+            ), txCommand
+        )
+        txBuilder.toLedgerTransaction(serviceHub).verify()
         val signedTx = serviceHub.signInitialTransaction(txBuilder)
         serviceHub.recordTransactions(signedTx)
         return signedTx
@@ -194,14 +405,41 @@ class CreateMessage(private val message: Message, private val notary: Party) : F
 
 //TODO merge both flows?
 @StartableByRPC
-class ConsumeMessage(private val stateRef: StateAndRef<MessageState>, private val notary: Party) : FlowLogic<SignedTransaction>() {
+class ConsumeMessage(
+    private val stateRef: StateAndRef<MessageState>,
+    private val notary: Party,
+    private val specifyExistingConstraint: Boolean,
+    private val addAnotherAutomaticConstraintState: Boolean
+) : FlowLogic<SignedTransaction>() {
     @Suspendable
     override fun call(): SignedTransaction {
         val oldMessageState = stateRef.state.data
-        val messageState = MessageState(Message(oldMessageState.message.value + "A"), ourIdentity, stateRef.state.data.linearId)
+        val messageState = MessageState(
-        val txCommand = Command(DummyMessageContract.Commands.Send(), messageState.participants.map { it.owningKey })
+            Message(oldMessageState.message.value + "A"),
-        val txBuilder =
+            ourIdentity,
-            TransactionBuilder(notary).withItems(StateAndContract(messageState, TEST_MESSAGE_CONTRACT_PROGRAM_ID), txCommand, stateRef)
+            stateRef.state.data.linearId
+        )
+        val txCommand = Command(
+            DummyMessageContract.Commands.Send(),
+            messageState.participants.map { it.owningKey })
+        val txBuilder = TransactionBuilder(notary).apply {
+            if (specifyExistingConstraint) {
+                addOutputState(messageState, stateRef.state.constraint)
+            } else {
+                addOutputState(messageState)
+            }
+            if (addAnotherAutomaticConstraintState) {
+                addOutputState(
+                    MessageState(
+                        Message("Another message"),
+                        ourIdentity,
+                        UniqueIdentifier()
+                    )
+                )
+            }
+            addInputState(stateRef)
+            addCommand(txCommand)
+        }
         txBuilder.toWireTransaction(serviceHub).toLedgerTransaction(serviceHub).verify()
         val signedTx = serviceHub.signInitialTransaction(txBuilder)
         serviceHub.recordTransactions(signedTx)
@@ -209,6 +447,18 @@ class ConsumeMessage(private val stateRef: StateAndRef<MessageState>, private va
     }
 }
 
+@CordaSerializable
+data class Message(val value: String)
+
+@BelongsToContract(DummyMessageContract::class)
+data class MessageState(
+    val message: Message,
+    val by: Party,
+    override val linearId: UniqueIdentifier = UniqueIdentifier()
+) : LinearState {
+    override val participants: List<AbstractParty> = listOf(by)
+}
+
 //TODO enrich original MessageContract for new command
 const val TEST_MESSAGE_CONTRACT_PROGRAM_ID = "net.corda.contracts.DummyMessageContract"
 
@@ -216,9 +466,7 @@ open class DummyMessageContract : Contract {
     override fun verify(tx: LedgerTransaction) {
         val command = tx.commands.requireSingleCommand<Commands.Send>()
         requireThat {
-            // Generic constraints around the IOU transaction.
+            val out = tx.outputsOfType<MessageState>().first()
-            "Only one output state should be created." using (tx.outputs.size == 1)
-            val out = tx.outputsOfType<MessageState>().single()
             "Message sender must sign." using (command.signers.containsAll(out.participants.map { it.owningKey }))
             "Message value must not be empty." using (out.message.value.isNotBlank())
         }