Serialization of large contract attachments causes OOM exception (#1991)

* Don't serialize contract attachment, only hash and contract class name if we are checkpointing

node-api/src/main/kotlin/net/corda/nodeapi/internal/serialization/kryo/DefaultKryoCustomizer.kt

@@ -14,7 +14,10 @@ import de.javakaffee.kryoserializers.guava.*
 import net.corda.core.contracts.ContractAttachment
 import net.corda.core.contracts.PrivacySalt
 import net.corda.core.crypto.CompositeKey
+import net.corda.core.crypto.SecureHash
 import net.corda.core.identity.PartyAndCertificate
+import net.corda.core.internal.AbstractAttachment
+import net.corda.core.serialization.MissingAttachmentsException
 import net.corda.core.serialization.SerializationWhitelist
 import net.corda.core.serialization.SerializeAsToken
 import net.corda.core.serialization.SerializedBytes
@@ -196,15 +199,38 @@ object DefaultKryoCustomizer {
 
     private object ContractAttachmentSerializer : Serializer<ContractAttachment>() {
         override fun write(kryo: Kryo, output: Output, obj: ContractAttachment) {
-            val buffer = ByteArrayOutputStream()
-            obj.attachment.open().use { it.copyTo(buffer) }
-            output.writeBytesWithLength(buffer.toByteArray())
+            if (kryo.serializationContext() != null) {
+                output.writeBytes(obj.attachment.id.bytes)
+            } else {
+                val buffer = ByteArrayOutputStream()
+                obj.attachment.open().use { it.copyTo(buffer) }
+                output.writeBytesWithLength(buffer.toByteArray())
+            }
             output.writeString(obj.contract)
         }
 
         override fun read(kryo: Kryo, input: Input, type: Class<ContractAttachment>): ContractAttachment {
-            val attachment = GeneratedAttachment(input.readBytesWithLength())
-            return ContractAttachment(attachment, input.readString())
+            if (kryo.serializationContext() != null) {
+                val attachmentHash = SecureHash.SHA256(input.readBytes(32))
+                val contract = input.readString()
+
+                val context = kryo.serializationContext()!!
+                val attachmentStorage = context.serviceHub.attachments
+
+                val lazyAttachment = object : AbstractAttachment({
+                    val attachment = attachmentStorage.openAttachment(attachmentHash) ?: throw MissingAttachmentsException(listOf(attachmentHash))
+                    attachment.open().readBytes()
+                }) {
+                    override val id = attachmentHash
+                }
+
+                return ContractAttachment(lazyAttachment, contract)
+            } else {
+                val attachment = GeneratedAttachment(input.readBytesWithLength())
+                val contract = input.readString()
+
+                return ContractAttachment(attachment, contract)
+            }
         }
     }
-}
+}

node-api/src/test/kotlin/net/corda/nodeapi/internal/serialization/ContractAttachmentSerializerTest.kt

@@ -0,0 +1,103 @@
+package net.corda.nodeapi.internal.serialization
+
+import net.corda.core.contracts.ContractAttachment
+import net.corda.core.serialization.*
+import net.corda.testing.SerializationEnvironmentRule
+import net.corda.testing.contracts.DummyContract
+import net.corda.testing.node.MockServices
+import org.assertj.core.api.Assertions.assertThat
+import org.assertj.core.api.Assertions.assertThatThrownBy
+import org.junit.Assert.assertArrayEquals
+import org.junit.Before
+import org.junit.Rule
+import org.junit.Test
+import kotlin.test.assertEquals
+
+class ContractAttachmentSerializerTest {
+
+    @Rule
+    @JvmField
+    val testSerialization = SerializationEnvironmentRule()
+
+    private lateinit var factory: SerializationFactory
+    private lateinit var context: SerializationContext
+    private lateinit var contextWithToken: SerializationContext
+
+    private val mockServices = MockServices()
+
+    @Before
+    fun setup() {
+        factory = testSerialization.env.SERIALIZATION_FACTORY
+        context = testSerialization.env.CHECKPOINT_CONTEXT
+
+        contextWithToken = context.withTokenContext(SerializeAsTokenContextImpl(Any(), factory, context, mockServices))
+    }
+
+    @Test
+    fun `write contract attachment and read it back`() {
+        val contractAttachment = ContractAttachment(GeneratedAttachment(ByteArray(0)), DummyContract.PROGRAM_ID)
+        // no token context so will serialize the whole attachment
+        val serialized = contractAttachment.serialize(factory, context)
+        val deserialized = serialized.deserialize(factory, context)
+
+        assertEquals(contractAttachment.id, deserialized.attachment.id)
+        assertEquals(contractAttachment.contract, deserialized.contract)
+        assertArrayEquals(contractAttachment.open().readBytes(), deserialized.open().readBytes())
+    }
+
+    @Test
+    fun `write contract attachment and read it back using token context`() {
+        val attachment = GeneratedAttachment("test".toByteArray())
+
+        mockServices.attachments.importAttachment(attachment.open())
+
+        val contractAttachment = ContractAttachment(attachment, DummyContract.PROGRAM_ID)
+        val serialized = contractAttachment.serialize(factory, contextWithToken)
+        val deserialized = serialized.deserialize(factory, contextWithToken)
+
+        assertEquals(contractAttachment.id, deserialized.attachment.id)
+        assertEquals(contractAttachment.contract, deserialized.contract)
+        assertArrayEquals(contractAttachment.open().readBytes(), deserialized.open().readBytes())
+    }
+
+    @Test
+    fun `check only serialize attachment id and contract class name when using token context`() {
+        val largeAttachmentSize = 1024 * 1024
+        val attachment = GeneratedAttachment(ByteArray(largeAttachmentSize))
+
+        mockServices.attachments.importAttachment(attachment.open())
+
+        val contractAttachment = ContractAttachment(attachment, DummyContract.PROGRAM_ID)
+        val serialized = contractAttachment.serialize(factory, contextWithToken)
+
+        assertThat(serialized.size).isLessThan(largeAttachmentSize)
+    }
+
+    @Test
+    fun `throws when missing attachment when using token context`() {
+        val attachment = GeneratedAttachment("test".toByteArray())
+
+        // don't importAttachment in mockService
+
+        val contractAttachment = ContractAttachment(attachment, DummyContract.PROGRAM_ID)
+        val serialized = contractAttachment.serialize(factory, contextWithToken)
+        val deserialized = serialized.deserialize(factory, contextWithToken)
+
+        assertThatThrownBy { deserialized.attachment.open() }.isInstanceOf(MissingAttachmentsException::class.java)
+    }
+
+    @Test
+    fun `check attachment in deserialize is lazy loaded when using token context`() {
+        val attachment = GeneratedAttachment(ByteArray(0))
+
+        // don't importAttachment in mockService
+
+        val contractAttachment = ContractAttachment(attachment, DummyContract.PROGRAM_ID)
+        val serialized = contractAttachment.serialize(factory, contextWithToken)
+        serialized.deserialize(factory, contextWithToken)
+
+        // MissingAttachmentsException thrown if we try to open attachment
+    }
+}
+
+