Several tests were corrupting Kryo which was then returned to the common pool. (#378)

* We were leaving trailing attachmentStorage on pooled kryo instances after some tests.  Changed attachment storage logic to make it impossible to leave it behind.

* Some low level tests corrupt the Kryo config, so do not return to pool when this is the case.  Also, we discovered that Kryo is caching class name to class resolution.  We don't want to do this where attachments are involved.  The errors raised highlighted a class missing from the whitelist.  Need to write a unit test to test the class loader issue.

* Unit test for attachment class loading with kryo.

core/src/main/kotlin/net/corda/core/serialization/CordaClassResolver.kt

@@ -92,6 +92,24 @@ class CordaClassResolver(val whitelist: ClassWhitelist) : DefaultClassResolver()
         return type.interfaces.any { it.isAnnotationPresent(CordaSerializable::class.java) || hasAnnotationOnInterface(it) }
                 || (type.superclass != null && hasAnnotationOnInterface(type.superclass))
     }
+
+    // Need to clear out class names from attachments.
+    override fun reset() {
+        super.reset()
+        // Kryo creates a cache of class name to Class<*> which does not work so well with multiple class loaders.
+        // TODO: come up with a more efficient way.  e.g. segregate the name space by class loader.
+        if(nameToClass != null) {
+            val classesToRemove: MutableList<String> = ArrayList(nameToClass.size)
+            for (entry in nameToClass.entries()) {
+                if (entry.value.classLoader is AttachmentsClassLoader) {
+                    classesToRemove += entry.key
+                }
+            }
+            for (className in classesToRemove) {
+                nameToClass.remove(className)
+            }
+        }
+    }
 }
 
 interface ClassWhitelist {

core/src/main/kotlin/net/corda/core/serialization/Kryo.kt

@@ -527,11 +527,18 @@ object ReferencesAwareJavaSerializer : JavaSerializer() {
 
 val ATTACHMENT_STORAGE = "ATTACHMENT_STORAGE"
 
-var Kryo.attachmentStorage: AttachmentStorage?
+val Kryo.attachmentStorage: AttachmentStorage?
     get() = this.context.get(ATTACHMENT_STORAGE, null) as AttachmentStorage?
-    set(value) {
+
-        this.context.put(ATTACHMENT_STORAGE, value)
+fun <T> Kryo.withAttachmentStorage(attachmentStorage: AttachmentStorage?, block: () -> T): T {
+    val priorAttachmentStorage = this.attachmentStorage
+    this.context.put(ATTACHMENT_STORAGE, attachmentStorage)
+    try {
+        return block()
+    } finally {
+        this.context.put(ATTACHMENT_STORAGE, priorAttachmentStorage)
     }
+}
 
 object OrderedSerializer : Serializer<HashMap<Any, Any>>() {
     override fun write(kryo: Kryo, output: Output, obj: HashMap<Any, Any>) {

core/src/test/kotlin/net/corda/core/node/AttachmentClassLoaderTests.kt

@@ -83,16 +83,11 @@ class AttachmentClassLoaderTests {
 
     @Before
     fun setup() {
+        // Do not release these back to the pool, since we do some unorthodox modifications to them below.
         kryo = p2PKryo().borrow()
         kryo2 = p2PKryo().borrow()
     }
 
-    @After
-    fun teardown() {
-        p2PKryo().release(kryo)
-        p2PKryo().release(kryo2)
-    }
-
     @Test
     fun `dynamically load AnotherDummyContract from isolated contracts jar`() {
         val child = ClassLoaderForTests()
@@ -258,8 +253,19 @@ class AttachmentClassLoaderTests {
         val state2 = bytes.deserialize(kryo)
         assertEquals(cl, state2.contract.javaClass.classLoader)
         assertNotNull(state2)
+
+        // We should be able to load same class from a different class loader and have them be distinct.
+        val cl2 = AttachmentsClassLoader(arrayOf(att0, att1, att2).map { storage.openAttachment(it)!! }, FilteringClassLoader)
+
+        kryo.classLoader = cl2
+        kryo.addToWhitelist(Class.forName("net.corda.contracts.isolated.AnotherDummyContract", true, cl2))
+
+        val state3 = bytes.deserialize(kryo)
+        assertEquals(cl2, state3.contract.javaClass.classLoader)
+        assertNotNull(state3)
     }
 
+
     @Test
     fun `test serialization of WireTransaction with statically loaded contract`() {
         val tx = ATTACHMENT_TEST_PROGRAM_ID.generateInitial(MEGA_CORP.ref(0), 42, DUMMY_NOTARY)
@@ -283,24 +289,25 @@ class AttachmentClassLoaderTests {
         kryo.addToWhitelist(Class.forName("net.corda.contracts.isolated.AnotherDummyContract\$Commands\$Create", true, child))
 
         // todo - think about better way to push attachmentStorage down to serializer
-        kryo.attachmentStorage = storage
+        val bytes = kryo.withAttachmentStorage(storage) {
 
-        val attachmentRef = importJar(storage)
+            val attachmentRef = importJar(storage)
 
-        tx.addAttachment(storage.openAttachment(attachmentRef)!!.id)
+            tx.addAttachment(storage.openAttachment(attachmentRef)!!.id)
 
-        val wireTransaction = tx.toWireTransaction()
+            val wireTransaction = tx.toWireTransaction()
-
-        val bytes = wireTransaction.serialize(kryo)
 
+            wireTransaction.serialize(kryo)
+        }
         // use empty attachmentStorage
-        kryo2.attachmentStorage = storage
+        kryo2.withAttachmentStorage(storage) {
 
-        val copiedWireTransaction = bytes.deserialize(kryo2)
+            val copiedWireTransaction = bytes.deserialize(kryo2)
 
-        assertEquals(1, copiedWireTransaction.outputs.size)
+            assertEquals(1, copiedWireTransaction.outputs.size)
-        val contract2 = copiedWireTransaction.outputs[0].data.contract as DummyContractBackdoor
+            val contract2 = copiedWireTransaction.outputs[0].data.contract as DummyContractBackdoor
-        assertEquals(42, contract2.inspectState(copiedWireTransaction.outputs[0].data))
+            assertEquals(42, contract2.inspectState(copiedWireTransaction.outputs[0].data))
+        }
     }
 
     @Test
@@ -312,22 +319,22 @@ class AttachmentClassLoaderTests {
         val storage = MockAttachmentStorage()
 
         // todo - think about better way to push attachmentStorage down to serializer
-        kryo.attachmentStorage = storage
-
         val attachmentRef = importJar(storage)
+        val bytes = kryo.withAttachmentStorage(storage) {
 
-        tx.addAttachment(storage.openAttachment(attachmentRef)!!.id)
+            tx.addAttachment(storage.openAttachment(attachmentRef)!!.id)
 
-        val wireTransaction = tx.toWireTransaction()
+            val wireTransaction = tx.toWireTransaction()
 
-        val bytes = wireTransaction.serialize(kryo)
+            wireTransaction.serialize(kryo)
-
+        }
-        // use empty attachmentStorage
+        // use empty attachmentStorage
-        kryo2.attachmentStorage = MockAttachmentStorage()
+        kryo2.withAttachmentStorage(MockAttachmentStorage()) {
-
+
-        val e = assertFailsWith(MissingAttachmentsException::class) {
+            val e = assertFailsWith(MissingAttachmentsException::class) {
-            bytes.deserialize(kryo2)
+                bytes.deserialize(kryo2)
+            }
+            assertEquals(attachmentRef, e.ids.single())
         }
-        assertEquals(attachmentRef, e.ids.single())
     }
 }

core/src/test/kotlin/net/corda/core/serialization/KryoTests.kt

@@ -10,6 +10,7 @@ import org.bouncycastle.jce.provider.BouncyCastleProvider
 import org.bouncycastle.pqc.jcajce.provider.BouncyCastlePQCProvider
 import org.junit.After
 import org.junit.Before
+import org.junit.Ignore
 import org.junit.Test
 import org.slf4j.LoggerFactory
 import java.io.InputStream
@@ -25,14 +26,10 @@ class KryoTests {
 
     @Before
     fun setup() {
+        // We deliberately do not return this, since we do some unorthodox registering below and do not want to pollute the pool.
         kryo = p2PKryo().borrow()
     }
 
-    @After
-    fun teardown() {
-        p2PKryo().release(kryo)
-    }
-
     @Test
     fun ok() {
         val birthday = Instant.parse("1984-04-17T00:30:00.00Z")

node/src/main/kotlin/net/corda/node/serialization/DefaultWhitelist.kt

@@ -6,6 +6,7 @@ import net.corda.core.node.CordaPluginRegistry
 import net.corda.core.serialization.SerializationCustomization
 import org.apache.activemq.artemis.api.core.SimpleString
 import rx.Notification
+import rx.exceptions.OnErrorNotImplementedException
 import java.math.BigDecimal
 import java.time.LocalDate
 import java.time.Period
@@ -49,6 +50,7 @@ class DefaultWhitelist : CordaPluginRegistry() {
             addToWhitelist(LocalDate::class.java)
             addToWhitelist(Period::class.java)
             addToWhitelist(BitSet::class.java)
+            addToWhitelist(OnErrorNotImplementedException::class.java)
         }
         return true
     }