[CORDA-738] Ensure encumbrances are bi-directional (#4089)

core/src/main/kotlin/net/corda/core/contracts/TransactionVerificationException.kt

@@ -116,13 +116,32 @@ sealed class TransactionVerificationException(val txId: SecureHash, message: Str
     class TransactionMissingEncumbranceException(txId: SecureHash, val missing: Int, val inOut: Direction)
         : TransactionVerificationException(txId, "Missing required encumbrance $missing in $inOut", null)
 
+    /**
+     * If two or more states refer to another state (as their encumbrance), then the bi-directionality property cannot
+     * be satisfied.
+     */
+    @KeepForDJVM
+    class TransactionDuplicateEncumbranceException(txId: SecureHash, index: Int)
+        : TransactionVerificationException(txId, "The bi-directionality property of encumbered output states " +
+            "is not satisfied. Index $index is referenced more than once", null)
+
+    /**
+     * An encumbered state should also be referenced as the encumbrance of another state in order to satisfy the
+     * bi-directionality property (a full cycle should be present).
+     */
+    @KeepForDJVM
+    class TransactionNonMatchingEncumbranceException(txId: SecureHash, nonMatching: Collection<Int>)
+        : TransactionVerificationException(txId, "The bi-directionality property of encumbered output states " +
+            "is not satisfied. Encumbered states should also be referenced as an encumbrance of another state to form " +
+            "a full cycle. Offending indices $nonMatching", null)
+
     /** Whether the inputs or outputs list contains an encumbrance issue, see [TransactionMissingEncumbranceException]. */
     @CordaSerializable
     @KeepForDJVM
     enum class Direction {
-        /** Issue in the inputs list */
+        /** Issue in the inputs list. */
         INPUT,
-        /** Issue in the outputs list */
+        /** Issue in the outputs list. */
         OUTPUT
     }
 

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

@@ -62,6 +62,7 @@ abstract class AbstractStateReplacementFlow {
         @Throws(StateReplacementException::class)
         override fun call(): StateAndRef<T> {
             val (stx) = assembleTx()
+            stx.verify(serviceHub, checkSufficientSignatures = false)
             val participantSessions = getParticipantSessions()
             progressTracker.currentStep = SIGNING
 

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

@@ -46,14 +46,14 @@ class NotaryChangeFlow<out T : ContractState>(
         return AbstractStateReplacementFlow.UpgradeTx(stx)
     }
 
-    /** Resolves the encumbrance state chain for the given [state] */
+    /** Resolves the encumbrance state chain for the given [state]. */
     private fun resolveEncumbrances(state: StateAndRef<T>): List<StateAndRef<T>> {
-        val states = mutableListOf(state)
+        val states = mutableSetOf(state)
         while (states.last().state.encumbrance != null) {
             val encumbranceStateRef = StateRef(states.last().ref.txhash, states.last().state.encumbrance!!)
             val encumbranceState = serviceHub.toStateAndRef<T>(encumbranceStateRef)
-            states.add(encumbranceState)
+            if (!states.add(encumbranceState)) break // Stop if there is a cycle.
         }
-        return states
+        return states.toList()
     }
 }

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

@@ -169,8 +169,20 @@ data class LedgerTransaction @JvmOverloads constructor(
 
     private fun checkEncumbrancesValid() {
         // Validate that all encumbrances exist within the set of input states.
-        val encumberedInputs = inputs.filter { it.state.encumbrance != null }
+        inputs.filter { it.state.encumbrance != null }
-        encumberedInputs.forEach { (state, ref) ->
+                .forEach { (state, ref) -> checkInputEncumbranceStateExists(state, ref) }
+
+        // Check that in the outputs,
+        // a) an encumbered state does not refer to itself as the encumbrance
+        // b) the number of outputs can contain the encumbrance
+        // c) the bi-directionality (full cycle) property is satisfied.
+        val statesAndEncumbrance = outputs.withIndex().filter { it.value.encumbrance != null }.map { Pair(it.index, it.value.encumbrance!!) }
+        if (!statesAndEncumbrance.isEmpty()) {
+            checkOutputEncumbrances(statesAndEncumbrance)
+        }
+    }
+
+    private fun checkInputEncumbranceStateExists(state: TransactionState<ContractState>, ref: StateRef) {
         val encumbranceStateExists = inputs.any {
             it.ref.txhash == ref.txhash && it.ref.index == state.encumbrance
         }
@@ -183,18 +195,55 @@ data class LedgerTransaction @JvmOverloads constructor(
         }
     }
 
-        // Check that, in the outputs, an encumbered state does not refer to itself as the encumbrance,
+    // Using basic graph theory, a full cycle of encumbered (co-dependent) states should exist to achieve bi-directional
-        // and that the number of outputs can contain the encumbrance.
+    // encumbrances. This property is important to ensure that no states involved in an encumbrance-relationship
-        for ((i, output) in outputs.withIndex()) {
+    // can be spent on their own. Briefly, if any of the states is having more than one encumbrance references by
-            val encumbranceIndex = output.encumbrance ?: continue
+    // other states, a full cycle detection will fail. As a result, all of the encumbered states must be present
-            if (encumbranceIndex == i || encumbranceIndex >= outputs.size) {
+    // as "from" and "to" only once (or zero times if no encumbrance takes place). For instance,
+    // a -> b
+    // c -> b    and     a -> b
+    // b -> a            b -> c
+    // do not satisfy the bi-directionality (full cycle) property.
+    //
+    // In the first example "b" appears twice in encumbrance ("to") list and "c" exists in the encumbered ("from") list only.
+    // Due the above, one could consume "a" and "b" in the same transaction and then, because "b" is already consumed, "c" cannot be spent.
+    //
+    // Similarly, the second example does not form a full cycle because "a" and "c" exist in one of the lists only.
+    // As a result, one can consume "b" and "c" in the same transactions, which will make "a" impossible to be spent.
+    //
+    // On other hand the following are valid constructions:
+    // a -> b            a -> c
+    // b -> c    and     c -> b
+    // c -> a            b -> a
+    // and form a full cycle, meaning that the bi-directionality property is satisfied.
+    private fun checkOutputEncumbrances(statesAndEncumbrance: List<Pair<Int, Int>>) {
+        // [Set] of "from" (encumbered states).
+        val encumberedSet = mutableSetOf<Int>()
+        // [Set] of "to" (encumbrance states).
+        val encumbranceSet = mutableSetOf<Int>()
+        // Update both [Set]s.
+        statesAndEncumbrance.forEach { (statePosition, encumbrance) ->
+            // Check it does not refer to itself.
+            if (statePosition == encumbrance || encumbrance >= outputs.size) {
                 throw TransactionVerificationException.TransactionMissingEncumbranceException(
                         id,
-                        encumbranceIndex,
+                        encumbrance,
                         TransactionVerificationException.Direction.OUTPUT)
+            } else {
+                encumberedSet.add(statePosition) // Guaranteed to have unique elements.
+                if (!encumbranceSet.add(encumbrance)) {
+                    throw TransactionVerificationException.TransactionDuplicateEncumbranceException(id, encumbrance)
                 }
             }
         }
+        // At this stage we have ensured that "from" and "to" [Set]s are equal in size, but we should check their
+        // elements do indeed match. If they don't match, we return their symmetric difference (disjunctive union).
+        val symmetricDifference = (encumberedSet union encumbranceSet).subtract(encumberedSet intersect encumbranceSet)
+        if (symmetricDifference.isNotEmpty()) {
+            // At least one encumbered state is not in the [encumbranceSet] and vice versa.
+            throw TransactionVerificationException.TransactionNonMatchingEncumbranceException(id, symmetricDifference)
+        }
+    }
 
     /**
      * Given a type and a function that returns a grouping key, associates inputs and outputs together so that they

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

@@ -102,13 +102,21 @@ data class NotaryChangeLedgerTransaction(
 
     override val references: List<StateAndRef<ContractState>> = emptyList()
 
-    /** We compute the outputs on demand by applying the notary field modification to the inputs */
+    /** We compute the outputs on demand by applying the notary field modification to the inputs. */
     override val outputs: List<TransactionState<ContractState>>
-        get() = inputs.mapIndexed { pos, (state) ->
+        get() = computeOutputs()
+
+    private fun computeOutputs(): List<TransactionState<ContractState>> {
+        val inputPositionIndex: Map<StateRef, Int> = inputs.mapIndexed { index, stateAndRef -> stateAndRef.ref to index }.toMap()
+        return inputs.map { (state, ref) ->
             if (state.encumbrance != null) {
-                state.copy(notary = newNotary, encumbrance = pos + 1)
+                val encumbranceStateRef = StateRef(ref.txhash, state.encumbrance)
+                val encumbrancePosition = inputPositionIndex[encumbranceStateRef]
+                        ?: throw IllegalStateException("Unable to generate output states – transaction not constructed correctly.")
+                state.copy(notary = newNotary, encumbrance = encumbrancePosition)
             } else state.copy(notary = newNotary)
         }
+    }
 
     override val requiredSigningKeys: Set<PublicKey>
         get() = inputs.flatMap { it.state.data.participants }.map { it.owningKey }.toSet() + notary.owningKey
@@ -118,18 +126,16 @@ data class NotaryChangeLedgerTransaction(
     }
 
     /**
-     * Check that encumbrances have been included in the inputs. The [NotaryChangeFlow] guarantees that an encumbrance
+     * Check that encumbrances have been included in the inputs.
-     * will follow its encumbered state in the inputs.
      */
     private fun checkEncumbrances() {
-        inputs.forEachIndexed { i, (state, ref) ->
+        val encumberedStates = inputs.asSequence().filter { it.state.encumbrance != null }.associateBy { it.ref }
-            state.encumbrance?.let {
+        if (encumberedStates.isNotEmpty()) {
-                val nextIndex = i + 1
+            inputs.forEach { (state, ref) ->
-                fun nextStateIsEncumbrance() = (inputs[nextIndex].ref.txhash == ref.txhash) && (inputs[nextIndex].ref.index == it)
+                if (StateRef(ref.txhash, state.encumbrance!!) !in encumberedStates) {
-                if (nextIndex >= inputs.size || !nextStateIsEncumbrance()) {
                     throw TransactionVerificationException.TransactionMissingEncumbranceException(
                             id,
-                            it,
+                            state.encumbrance,
                             TransactionVerificationException.Direction.INPUT)
                 }
             }

core/src/main/kotlin/net/corda/core/utilities/KotlinUtils.kt

@@ -134,4 +134,3 @@ fun <V> Future<V>.getOrThrow(timeout: Duration? = null): V = try {
 } catch (e: ExecutionException) {
     throw e.cause!!
 }
-

core/src/test/kotlin/net/corda/core/transactions/TransactionEncumbranceTests.kt

@@ -4,6 +4,7 @@ import com.nhaarman.mockito_kotlin.doReturn
 import com.nhaarman.mockito_kotlin.whenever
 import net.corda.core.contracts.Contract
 import net.corda.core.contracts.ContractState
+import net.corda.core.contracts.TransactionVerificationException
 import net.corda.core.contracts.requireThat
 import net.corda.core.identity.AbstractParty
 import net.corda.core.identity.CordaX500Name
@@ -21,33 +22,43 @@ import org.junit.Rule
 import org.junit.Test
 import java.time.Instant
 import java.time.temporal.ChronoUnit
+import kotlin.test.assertFailsWith
 
 const val TEST_TIMELOCK_ID = "net.corda.core.transactions.TransactionEncumbranceTests\$DummyTimeLock"
 
 class TransactionEncumbranceTests {
+    @Rule
+    @JvmField
+    val testSerialization = SerializationEnvironmentRule()
+
     private companion object {
         val DUMMY_NOTARY = TestIdentity(DUMMY_NOTARY_NAME, 20).party
         val megaCorp = TestIdentity(CordaX500Name("MegaCorp", "London", "GB"))
         val MINI_CORP = TestIdentity(CordaX500Name("MiniCorp", "London", "GB")).party
         val MEGA_CORP get() = megaCorp.party
         val MEGA_CORP_PUBKEY get() = megaCorp.publicKey
-    }
 
-    @Rule
-    @JvmField
-    val testSerialization = SerializationEnvironmentRule()
         val defaultIssuer = MEGA_CORP.ref(1)
 
         val state = Cash.State(
                 amount = 1000.DOLLARS `issued by` defaultIssuer,
                 owner = MEGA_CORP
         )
+
         val stateWithNewOwner = state.copy(owner = MINI_CORP)
+        val extraCashState = state.copy(amount = state.amount * 3)
 
         val FOUR_PM: Instant = Instant.parse("2015-04-17T16:00:00.00Z")
         val FIVE_PM: Instant = FOUR_PM.plus(1, ChronoUnit.HOURS)
         val timeLock = DummyTimeLock.State(FIVE_PM)
 
+
+        val ledgerServices = MockServices(listOf("net.corda.core.transactions", "net.corda.finance.contracts.asset"), MEGA_CORP.name,
+                rigorousMock<IdentityServiceInternal>().also {
+                    doReturn(MEGA_CORP).whenever(it).partyFromKey(MEGA_CORP_PUBKEY)
+                })
+    }
+
     class DummyTimeLock : Contract {
         override fun verify(tx: LedgerTransaction) {
             val timeLockInput = tx.inputsOfType<State>().singleOrNull() ?: return
@@ -65,18 +76,78 @@ class TransactionEncumbranceTests {
         }
     }
 
-    private val ledgerServices = MockServices(listOf("net.corda.core.transactions", "net.corda.finance.contracts.asset"), MEGA_CORP.name,
-            rigorousMock<IdentityServiceInternal>().also {
-                doReturn(MEGA_CORP).whenever(it).partyFromKey(MEGA_CORP_PUBKEY)
-            })
-
     @Test
-    fun `state can be encumbered`() {
+    fun `states can be bi-directionally encumbered`() {
+        // Basic encumbrance example for encumbrance index links 0 -> 1 and 1 -> 0
         ledgerServices.ledger(DUMMY_NOTARY) {
             transaction {
                 attachments(Cash.PROGRAM_ID, TEST_TIMELOCK_ID)
                 input(Cash.PROGRAM_ID, state)
-                output(Cash.PROGRAM_ID, encumbrance = 1, contractState = stateWithNewOwner)
+                output(Cash.PROGRAM_ID, "state encumbered by 5pm time-lock", encumbrance = 1, contractState = stateWithNewOwner)
+                output(TEST_TIMELOCK_ID, "5pm time-lock", 0, timeLock)
+                command(MEGA_CORP.owningKey, Cash.Commands.Move())
+                verifies()
+            }
+        }
+
+        // Full cycle example with 4 elements 0 -> 1, 1 -> 2, 2 -> 3 and 3 -> 0
+        // All 3 Cash states and the TimeLock are linked and should be consumed in the same transaction.
+        // Note that all of the Cash states are encumbered both together and with time lock.
+        ledgerServices.ledger(DUMMY_NOTARY) {
+            transaction {
+                attachments(Cash.PROGRAM_ID, TEST_TIMELOCK_ID)
+                input(Cash.PROGRAM_ID, extraCashState)
+                output(Cash.PROGRAM_ID, "state encumbered by state 1", encumbrance = 1, contractState = stateWithNewOwner)
+                output(Cash.PROGRAM_ID, "state encumbered by state 2", encumbrance = 2, contractState = stateWithNewOwner)
+                output(Cash.PROGRAM_ID, "state encumbered by state 3", encumbrance = 3, contractState = stateWithNewOwner)
+                output(TEST_TIMELOCK_ID, "5pm time-lock", 0, timeLock)
+                command(MEGA_CORP.owningKey, Cash.Commands.Move())
+                verifies()
+            }
+        }
+
+        // A transaction that includes multiple independent encumbrance chains.
+        // Each Cash state is encumbered with its own TimeLock.
+        // Note that all of the Cash states are encumbered both together and with time lock.
+        ledgerServices.ledger(DUMMY_NOTARY) {
+            transaction {
+                attachments(Cash.PROGRAM_ID, TEST_TIMELOCK_ID)
+                input(Cash.PROGRAM_ID, extraCashState)
+                output(Cash.PROGRAM_ID, "state encumbered by 5pm time-lock A", encumbrance = 3, contractState = stateWithNewOwner)
+                output(Cash.PROGRAM_ID, "state encumbered by 5pm time-lock B", encumbrance = 4, contractState = stateWithNewOwner)
+                output(Cash.PROGRAM_ID, "state encumbered by 5pm time-lock C", encumbrance = 5, contractState = stateWithNewOwner)
+                output(TEST_TIMELOCK_ID, "5pm time-lock A", 0, timeLock)
+                output(TEST_TIMELOCK_ID, "5pm time-lock B", 1, timeLock)
+                output(TEST_TIMELOCK_ID, "5pm time-lock C", 2, timeLock)
+                command(MEGA_CORP.owningKey, Cash.Commands.Move())
+                verifies()
+            }
+        }
+
+        // Full cycle example with 4 elements (different combination) 0 -> 3, 1 -> 2, 2 -> 0 and 3 -> 1
+        ledgerServices.ledger(DUMMY_NOTARY) {
+            transaction {
+                attachments(Cash.PROGRAM_ID, TEST_TIMELOCK_ID)
+                input(Cash.PROGRAM_ID, extraCashState)
+                output(Cash.PROGRAM_ID, "state encumbered by state 3", encumbrance = 3, contractState = stateWithNewOwner)
+                output(Cash.PROGRAM_ID, "state encumbered by state 2", encumbrance = 2, contractState = stateWithNewOwner)
+                output(Cash.PROGRAM_ID, "state encumbered by state 0", encumbrance = 0, contractState = stateWithNewOwner)
+                output(TEST_TIMELOCK_ID, "5pm time-lock", 1, timeLock)
+                command(MEGA_CORP.owningKey, Cash.Commands.Move())
+                verifies()
+            }
+        }
+    }
+
+    @Test
+    fun `non bi-directional encumbrance will fail`() {
+        // Single encumbrance with no back link.
+        assertFailsWith<TransactionVerificationException.TransactionNonMatchingEncumbranceException> {
+            ledgerServices.ledger(DUMMY_NOTARY) {
+                transaction {
+                    attachments(Cash.PROGRAM_ID, TEST_TIMELOCK_ID)
+                    input(Cash.PROGRAM_ID, state)
+                    output(Cash.PROGRAM_ID, "state encumbered by 5pm time-lock", encumbrance = 1, contractState = stateWithNewOwner)
                     output(TEST_TIMELOCK_ID, "5pm time-lock", timeLock)
                     command(MEGA_CORP.owningKey, Cash.Commands.Move())
                     verifies()
@@ -84,6 +155,59 @@ class TransactionEncumbranceTests {
             }
         }
 
+        // Full cycle fails due to duplicate encumbrance reference.
+        // 0 -> 1, 1 -> 3, 2 -> 3 (thus 3 is referenced two times).
+        assertFailsWith<TransactionVerificationException.TransactionDuplicateEncumbranceException> {
+            ledgerServices.ledger(DUMMY_NOTARY) {
+                transaction {
+                    attachments(Cash.PROGRAM_ID, TEST_TIMELOCK_ID)
+                    input(Cash.PROGRAM_ID, state)
+                    output(Cash.PROGRAM_ID, "state encumbered by state 1", encumbrance = 1, contractState = stateWithNewOwner)
+                    output(Cash.PROGRAM_ID, "state encumbered by state 3", encumbrance = 3, contractState = stateWithNewOwner)
+                    output(Cash.PROGRAM_ID, "state encumbered by state 3 again", encumbrance = 3, contractState = stateWithNewOwner)
+                    output(TEST_TIMELOCK_ID, "5pm time-lock", timeLock)
+                    command(MEGA_CORP.owningKey, Cash.Commands.Move())
+                    verifies()
+                }
+            }
+        }
+
+        // No Full cycle due to non-matching encumbered-encumbrance elements.
+        // 0 -> 1, 1 -> 3, 2 -> 0 (thus offending indices [2, 3], because 2 is not referenced and 3 is not encumbered).
+        assertFailsWith<TransactionVerificationException.TransactionNonMatchingEncumbranceException> {
+            ledgerServices.ledger(DUMMY_NOTARY) {
+                transaction {
+                    attachments(Cash.PROGRAM_ID, TEST_TIMELOCK_ID)
+                    input(Cash.PROGRAM_ID, state)
+                    output(Cash.PROGRAM_ID, "state encumbered by state 1", encumbrance = 1, contractState = stateWithNewOwner)
+                    output(Cash.PROGRAM_ID, "state encumbered by state 3", encumbrance = 3, contractState = stateWithNewOwner)
+                    output(Cash.PROGRAM_ID, "state encumbered by state 0", encumbrance = 0, contractState = stateWithNewOwner)
+                    output(TEST_TIMELOCK_ID, "5pm time-lock", timeLock)
+                    command(MEGA_CORP.owningKey, Cash.Commands.Move())
+                    verifies()
+                }
+            }
+        }
+
+        // No Full cycle in one of the encumbrance chains due to non-matching encumbered-encumbrance elements.
+        // 0 -> 2, 2 -> 0 is valid. On the other hand, there is 1 -> 3 only and 3 -> 1 does not exist.
+        // (thus offending indices [1, 3], because 1 is not referenced and 3 is not encumbered).
+        assertFailsWith<TransactionVerificationException.TransactionNonMatchingEncumbranceException> {
+            ledgerServices.ledger(DUMMY_NOTARY) {
+                transaction {
+                    attachments(Cash.PROGRAM_ID, TEST_TIMELOCK_ID)
+                    input(Cash.PROGRAM_ID, state)
+                    output(Cash.PROGRAM_ID, "state encumbered by 5pm time-lock A", encumbrance = 2, contractState = stateWithNewOwner)
+                    output(Cash.PROGRAM_ID, "state encumbered by 5pm time-lock B", encumbrance = 3, contractState = stateWithNewOwner)
+                    output(TEST_TIMELOCK_ID, "5pm time-lock A", 0, timeLock)
+                    output(TEST_TIMELOCK_ID, "5pm time-lock B", timeLock)
+                    command(MEGA_CORP.owningKey, Cash.Commands.Move())
+                    verifies()
+                }
+            }
+        }
+    }
+
     @Test
     fun `state can transition if encumbrance rules are met`() {
         ledgerServices.ledger(DUMMY_NOTARY) {
@@ -132,7 +256,7 @@ class TransactionEncumbranceTests {
             unverifiedTransaction {
                 attachments(Cash.PROGRAM_ID, TEST_TIMELOCK_ID)
                 output(Cash.PROGRAM_ID, "state encumbered by 5pm time-lock", encumbrance = 1, contractState = state)
-                output(TEST_TIMELOCK_ID, "5pm time-lock", timeLock)
+                output(TEST_TIMELOCK_ID, "5pm time-lock",0, timeLock)
             }
             transaction {
                 attachments(Cash.PROGRAM_ID)
@@ -151,7 +275,7 @@ class TransactionEncumbranceTests {
             transaction {
                 attachments(Cash.PROGRAM_ID)
                 input(Cash.PROGRAM_ID, state)
-                output(Cash.PROGRAM_ID, encumbrance = 0, contractState = stateWithNewOwner)
+                output(Cash.PROGRAM_ID, "state encumbered by itself", encumbrance = 0, contractState = stateWithNewOwner)
                 command(MEGA_CORP.owningKey, Cash.Commands.Move())
                 this `fails with` "Missing required encumbrance 0 in OUTPUT"
             }
@@ -164,7 +288,7 @@ class TransactionEncumbranceTests {
             transaction {
                 attachments(Cash.PROGRAM_ID, TEST_TIMELOCK_ID)
                 input(Cash.PROGRAM_ID, state)
-                output(TEST_TIMELOCK_ID, encumbrance = 2, contractState = stateWithNewOwner)
+                output(TEST_TIMELOCK_ID, "state encumbered by state 2 which does not exist", encumbrance = 2, contractState = stateWithNewOwner)
                 output(TEST_TIMELOCK_ID, timeLock)
                 command(MEGA_CORP.owningKey, Cash.Commands.Move())
                 this `fails with` "Missing required encumbrance 2 in OUTPUT"
@@ -178,7 +302,7 @@ class TransactionEncumbranceTests {
             unverifiedTransaction {
                 attachments(Cash.PROGRAM_ID, TEST_TIMELOCK_ID)
                 output(Cash.PROGRAM_ID, "state encumbered by some other state", encumbrance = 1, contractState = state)
-                output(Cash.PROGRAM_ID, "some other state", state)
+                output(Cash.PROGRAM_ID, "some other state", encumbrance = 0, contractState = state)
                 output(TEST_TIMELOCK_ID, "5pm time-lock", timeLock)
             }
             transaction {

node/src/test/kotlin/net/corda/node/services/NotaryChangeTests.kt

@@ -109,20 +109,13 @@ class NotaryChangeTests {
         // Check that all encumbrances have been propagated to the outputs
         val originalOutputs = issueTx.outputStates
         val newOutputs = notaryChangeTx.outputStates
-        assertTrue(originalOutputs.minus(newOutputs).isEmpty())
+        assertTrue(originalOutputs.size == newOutputs.size && originalOutputs.containsAll(newOutputs))
 
-        // Check that encumbrance links aren't broken after notary change
+        // Check if encumbrance linking between states has not changed.
-        val encumbranceLink = HashMap<ContractState, ContractState?>()
+        val originalLinkedStates = issueTx.outputs.asSequence().filter { it.encumbrance != null }.map { Pair(it.data, issueTx.outputs[it.encumbrance!!].data) }.toSet()
-        issueTx.outputs.forEach {
+        val notaryChangeLinkedStates = notaryChangeTx.outputs.asSequence().filter { it.encumbrance != null }.map { Pair(it.data, notaryChangeTx.outputs[it.encumbrance!!].data) }.toSet()
-            val currentState = it.data
+
-            val encumbranceState = it.encumbrance?.let { issueTx.outputs[it].data }
+        assertTrue { originalLinkedStates.size == notaryChangeLinkedStates.size && originalLinkedStates.containsAll(notaryChangeLinkedStates) }
-            encumbranceLink[currentState] = encumbranceState
-        }
-        notaryChangeTx.outputs.forEach {
-            val currentState = it.data
-            val encumbranceState = it.encumbrance?.let { notaryChangeTx.outputs[it].data }
-            assertEquals(encumbranceLink[currentState], encumbranceState)
-        }
     }
 
     @Test
@@ -172,10 +165,11 @@ class NotaryChangeTests {
         val stateB = DummyContract.SingleOwnerState(Random().nextInt(), owner.party)
         val stateC = DummyContract.SingleOwnerState(Random().nextInt(), owner.party)
 
+        // Ensure encumbrances form a cycle.
         val tx = TransactionBuilder(null).apply {
             addCommand(Command(DummyContract.Commands.Create(), owner.party.owningKey))
             addOutputState(stateA, DummyContract.PROGRAM_ID, notaryIdentity, encumbrance = 2) // Encumbered by stateB
-            addOutputState(stateC, DummyContract.PROGRAM_ID, notaryIdentity)
+            addOutputState(stateC, DummyContract.PROGRAM_ID, notaryIdentity, encumbrance = 0) // Encumbered by stateA
             addOutputState(stateB, DummyContract.PROGRAM_ID, notaryIdentity, encumbrance = 1) // Encumbered by stateC
         }
         val stx = services.signInitialTransaction(tx)