Use a wrapper over StateRef for references (#4286)

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

@@ -13,6 +13,7 @@ import net.corda.core.internal.notary.generateSignature
 import net.corda.core.internal.notary.validateSignatures
 import net.corda.core.internal.pushToLoggingContext
 import net.corda.core.transactions.ContractUpgradeWireTransaction
+import net.corda.core.transactions.ReferenceStateRef
 import net.corda.core.transactions.SignedTransaction
 import net.corda.core.transactions.WireTransaction
 import net.corda.core.utilities.ProgressTracker
@@ -98,7 +99,7 @@ class NotaryFlow {
             val ctx = stx.coreTransaction
             val tx = when (ctx) {
                 is ContractUpgradeWireTransaction -> ctx.buildFilteredTransaction()
-                is WireTransaction -> ctx.buildFilteredTransaction(Predicate { it is StateRef || it is TimeWindow || it == notaryParty })
+                is WireTransaction -> ctx.buildFilteredTransaction(Predicate { it is StateRef || it is ReferenceStateRef || it is TimeWindow || it == notaryParty })
                 else -> ctx
             }
             return session.sendAndReceiveWithRetry(NotarisationPayload(tx, signature))

core/src/main/kotlin/net/corda/core/internal/CordaUtils.kt

@@ -15,7 +15,7 @@ import net.corda.core.transactions.TransactionBuilder
 import net.corda.core.transactions.WireTransaction
 import org.slf4j.MDC
 
-// *Internal* Corda-specific utilities
+// *Internal* Corda-specific utilities.
 
 const val PLATFORM_VERSION = 4
 
@@ -33,13 +33,13 @@ fun checkMinimumPlatformVersion(minimumPlatformVersion: Int, requiredMinPlatform
     }
 }
 
-/** Provide access to internal method for AttachmentClassLoaderTests */
+/** Provide access to internal method for AttachmentClassLoaderTests. */
 @DeleteForDJVM
 fun TransactionBuilder.toWireTransaction(services: ServicesForResolution, serializationContext: SerializationContext): WireTransaction {
     return toWireTransactionWithContext(services, serializationContext)
 }
 
-/** Provide access to internal method for AttachmentClassLoaderTests */
+/** Provide access to internal method for AttachmentClassLoaderTests. */
 @DeleteForDJVM
 fun TransactionBuilder.toLedgerTransaction(services: ServicesForResolution, serializationContext: SerializationContext): LedgerTransaction {
     return toLedgerTransactionWithContext(services, serializationContext)

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

@@ -147,7 +147,9 @@ class FilteredTransaction internal constructor(
                 wtx.attachments.forEachIndexed { internalIndex, it -> filter(it, ATTACHMENTS_GROUP.ordinal, internalIndex) }
                 if (wtx.notary != null) filter(wtx.notary, NOTARY_GROUP.ordinal, 0)
                 if (wtx.timeWindow != null) filter(wtx.timeWindow, TIMEWINDOW_GROUP.ordinal, 0)
-                wtx.references.forEachIndexed { internalIndex, it -> filter(it, REFERENCES_GROUP.ordinal, internalIndex) }
+                // Note that because [inputs] and [references] share the same type [StateRef], we use a wrapper for references [ReferenceStateRef],
+                // when filtering. Thus, to filter-in all [references] based on type, one should use the wrapper type [ReferenceStateRef] and not [StateRef].
+                wtx.references.forEachIndexed { internalIndex, it -> filter(ReferenceStateRef(it), REFERENCES_GROUP.ordinal, internalIndex) }
                 // It is highlighted that because there is no a signers property in TraversableTransaction,
                 // one cannot specifically filter them in or out.
                 // The above is very important to ensure someone won't filter out the signers component group if at least one
@@ -344,3 +346,8 @@ class ComponentVisibilityException(val id: SecureHash, val reason: String) : Cor
 @KeepForDJVM
 @CordaSerializable
 class FilteredTransactionVerificationException(val id: SecureHash, val reason: String) : CordaException("Transaction with id:$id cannot be verified. Reason: $reason")
+
+/** Wrapper over [StateRef] to be used when filtering reference states. */
+@KeepForDJVM
+@CordaSerializable
+data class ReferenceStateRef(val stateRef: StateRef)

core/src/test/kotlin/net/corda/core/crypto/PartialMerkleTreeTest.kt

@@ -8,11 +8,13 @@ import net.corda.core.identity.CordaX500Name
 import net.corda.core.identity.Party
 import net.corda.core.serialization.deserialize
 import net.corda.core.serialization.serialize
+import net.corda.core.transactions.ReferenceStateRef
 import net.corda.core.transactions.WireTransaction
 import net.corda.finance.DOLLARS
 import net.corda.finance.`issued by`
 import net.corda.finance.contracts.asset.Cash
 import net.corda.node.services.api.IdentityServiceInternal
+import net.corda.testing.common.internal.testNetworkParameters
 import net.corda.testing.core.DUMMY_NOTARY_NAME
 import net.corda.testing.core.SerializationEnvironmentRule
 import net.corda.testing.core.TestIdentity
@@ -59,9 +61,14 @@ class PartialMerkleTreeTest {
         hashed = nodes.map { it.serialize().sha256() }
         expectedRoot = MerkleTree.getMerkleTree(hashed.toMutableList() + listOf(zeroHash, zeroHash)).hash
         merkleTree = MerkleTree.getMerkleTree(hashed)
-        testLedger = MockServices(emptyList(), MEGA_CORP.name, rigorousMock<IdentityServiceInternal>().also {
-            doReturn(MEGA_CORP).whenever(it).partyFromKey(MEGA_CORP_PUBKEY)
-        }).ledger(DUMMY_NOTARY) {
+
+        testLedger = MockServices(
+                cordappPackages = emptyList(),
+                initialIdentity = TestIdentity(MEGA_CORP.name),
+                identityService = rigorousMock<IdentityServiceInternal>().also {
+                    doReturn(MEGA_CORP).whenever(it).partyFromKey(MEGA_CORP_PUBKEY) },
+                networkParameters = testNetworkParameters(minimumPlatformVersion = 4)
+        ).ledger(DUMMY_NOTARY) {
             unverifiedTransaction {
                 attachments(Cash.PROGRAM_ID)
                 output(Cash.PROGRAM_ID, "MEGA_CORP cash",
@@ -76,6 +83,7 @@ class PartialMerkleTreeTest {
             transaction {
                 attachments(Cash.PROGRAM_ID)
                 input("MEGA_CORP cash")
+                reference("dummy cash 1")
                 output(Cash.PROGRAM_ID, "MEGA_CORP cash".output<Cash.State>().copy(owner = MINI_CORP))
                 command(MEGA_CORP_PUBKEY, Cash.Commands.Move())
                 timeWindow(TEST_TX_TIME)
@@ -148,6 +156,7 @@ class PartialMerkleTreeTest {
         // the signers component is also sent (required for visibility purposes).
         assertEquals(5, ftx.filteredComponentGroups.size)
         assertEquals(1, ftx.inputs.size)
+        assertEquals(0, ftx.references.size)
         assertEquals(0, ftx.attachments.size)
         assertEquals(1, ftx.outputs.size)
         assertEquals(1, ftx.commands.size)
@@ -173,6 +182,7 @@ class PartialMerkleTreeTest {
         assertTrue(ftxNothing.attachments.isEmpty())
         assertTrue(ftxNothing.commands.isEmpty())
         assertTrue(ftxNothing.inputs.isEmpty())
+        assertTrue(ftxNothing.references.isEmpty())
         assertTrue(ftxNothing.outputs.isEmpty())
         assertNull(ftxNothing.timeWindow)
         assertTrue(ftxNothing.availableComponentGroups.flatten().isEmpty())
@@ -321,4 +331,21 @@ class PartialMerkleTreeTest {
         // The provided hash is not in the tree (using a leaf that didn't exist in the original Merkle tree).
         assertFailsWith<MerkleTreeException> { pmtAllIncluded.leafIndex(SecureHash.sha256("30")) }
     }
+
+    @Test
+    fun `building Merkle for reference states only`() {
+        fun filtering(elem: Any): Boolean {
+            return when (elem) {
+                is ReferenceStateRef -> true
+                else -> false
+            }
+        }
+
+        val ftx = testTx.buildFilteredTransaction(Predicate(::filtering))
+
+        assertEquals(1, ftx.filteredComponentGroups.size)
+        assertEquals(0, ftx.inputs.size)
+        assertEquals(1, ftx.references.size)
+        ftx.verify()
+    }
 }

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

@@ -23,7 +23,7 @@ import org.junit.Rule
 import org.junit.Test
 import kotlin.test.assertFailsWith
 
-val CONTRACT_ID = "net.corda.core.transactions.ReferenceStateTests\$ExampleContract"
+const val CONTRACT_ID = "net.corda.core.transactions.ReferenceStateTests\$ExampleContract"
 
 class ReferenceStateTests {
     private companion object {

docs/source/changelog.rst

@@ -277,6 +277,11 @@ Unreleased
   normal state when it occurs in an input or output position. *This feature is only available on Corda networks running
   with a minimum platform version of 4.*
 
+* A new wrapper class over ``StateRef`` is introduced, called ``ReferenceStateRef``. Although "reference input states" are stored as
+  ``StateRef`` objects in ``WireTransaction``, we needed a way to distinguish between "input states" and "reference input states" when
+  required to filter by object type. Thus, when one wants to filter-in all "reference input states" in a ``FilteredTransaction``
+  then he/she should check if it is of type ``ReferenceStateRef``.
+
 * Removed type parameter `U` from `tryLockFungibleStatesForSpending` to allow the function to be used with `FungibleState`
   as well as `FungibleAsset`. This _might_ cause a compile failure in some obscure cases due to the removal of the type
   parameter from the method. If your CorDapp does specify types explicitly when using this method then updating the types