Refactor KryoAMQPSerializer to go through generic APIs to access AMQP serialization (#1225)

client/rpc/src/main/kotlin/net/corda/client/rpc/CordaRPCClient.kt

@@ -71,7 +71,7 @@ class CordaRPCClient(
         fun initialiseSerialization() {
             try {
                 SerializationDefaults.SERIALIZATION_FACTORY = SerializationFactoryImpl().apply {
-                    registerScheme(KryoClientSerializationScheme())
+                    registerScheme(KryoClientSerializationScheme(this))
                     registerScheme(AMQPClientSerializationScheme())
                 }
                 SerializationDefaults.P2P_CONTEXT = KRYO_P2P_CONTEXT

client/rpc/src/main/kotlin/net/corda/client/rpc/serialization/SerializationScheme.kt

@@ -3,20 +3,21 @@ package net.corda.client.rpc.serialization
 import com.esotericsoftware.kryo.pool.KryoPool
 import net.corda.client.rpc.internal.RpcClientObservableSerializer
 import net.corda.core.serialization.SerializationContext
+import net.corda.core.serialization.SerializationFactory
 import net.corda.core.utilities.ByteSequence
 import net.corda.nodeapi.RPCKryo
 import net.corda.nodeapi.internal.serialization.AbstractKryoSerializationScheme
 import net.corda.nodeapi.internal.serialization.DefaultKryoCustomizer
 import net.corda.nodeapi.internal.serialization.KryoHeaderV0_1
 
-class KryoClientSerializationScheme : AbstractKryoSerializationScheme() {
+class KryoClientSerializationScheme(serializationFactory: SerializationFactory) : AbstractKryoSerializationScheme(serializationFactory) {
     override fun canDeserializeVersion(byteSequence: ByteSequence, target: SerializationContext.UseCase): Boolean {
         return byteSequence == KryoHeaderV0_1 && (target == SerializationContext.UseCase.RPCClient || target == SerializationContext.UseCase.P2P)
     }
 
     override fun rpcClientKryoPool(context: SerializationContext): KryoPool {
         return KryoPool.Builder {
-            DefaultKryoCustomizer.customize(RPCKryo(RpcClientObservableSerializer, context.whitelist)).apply { classLoader = context.deserializationClassLoader }
+            DefaultKryoCustomizer.customize(RPCKryo(RpcClientObservableSerializer, serializationFactory, context)).apply { classLoader = context.deserializationClassLoader }
         }.build()
     }
 

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

@@ -81,6 +81,11 @@ interface SerializationContext {
      */
     fun withWhitelisted(clazz: Class<*>): SerializationContext
 
+    /**
+     * Helper method to return a new context based on this context but with serialization using the format this header sequence represents.
+     */
+    fun withPreferredSerializationVersion(versionHeader: ByteSequence): SerializationContext
+
     /**
      * The use case that we are serializing for, since it influences the implementations chosen.
      */

node-api/src/main/kotlin/net/corda/nodeapi/RPCStructures.kt

@@ -8,6 +8,8 @@ import net.corda.core.concurrent.CordaFuture
 import net.corda.core.CordaRuntimeException
 import net.corda.core.serialization.ClassWhitelist
 import net.corda.core.serialization.CordaSerializable
+import net.corda.core.serialization.SerializationContext
+import net.corda.core.serialization.SerializationFactory
 import net.corda.core.toFuture
 import net.corda.core.toObservable
 import net.corda.nodeapi.config.OldConfig
@@ -47,7 +49,7 @@ class PermissionException(msg: String) : RuntimeException(msg)
 // The Kryo used for the RPC wire protocol. Every type in the wire protocol is listed here explicitly.
 // This is annoying to write out, but will make it easier to formalise the wire protocol when the time comes,
 // because we can see everything we're using in one place.
-class RPCKryo(observableSerializer: Serializer<Observable<*>>, whitelist: ClassWhitelist) : CordaKryo(CordaClassResolver(whitelist)) {
+class RPCKryo(observableSerializer: Serializer<Observable<*>>, val serializationFactory: SerializationFactory, val serializationContext: SerializationContext) : CordaKryo(CordaClassResolver(serializationFactory, serializationContext)) {
     init {
         DefaultKryoCustomizer.customize(this)
 

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

@@ -11,9 +11,22 @@ import net.corda.nodeapi.internal.serialization.amqp.SerializationOutput
 import net.corda.nodeapi.internal.serialization.amqp.SerializerFactory
 import java.util.concurrent.ConcurrentHashMap
 
-private const val AMQP_ENABLED = false
+internal val AMQP_ENABLED get() = SerializationDefaults.P2P_CONTEXT.preferedSerializationVersion == AmqpHeaderV1_0
 
 abstract class AbstractAMQPSerializationScheme : SerializationScheme {
+    internal companion object {
+        fun registerCustomSerializers(factory: SerializerFactory) {
+            factory.apply {
+                register(net.corda.nodeapi.internal.serialization.amqp.custom.PublicKeySerializer)
+                register(net.corda.nodeapi.internal.serialization.amqp.custom.ThrowableSerializer(this))
+                register(net.corda.nodeapi.internal.serialization.amqp.custom.X500NameSerializer)
+                register(net.corda.nodeapi.internal.serialization.amqp.custom.BigDecimalSerializer)
+                register(net.corda.nodeapi.internal.serialization.amqp.custom.CurrencySerializer)
+                register(net.corda.nodeapi.internal.serialization.amqp.custom.InstantSerializer(this))
+            }
+        }
+    }
+
     private val serializerFactoriesForContexts = ConcurrentHashMap<Pair<ClassWhitelist, ClassLoader>, SerializerFactory>()
 
     protected abstract fun rpcClientSerializerFactory(context: SerializationContext): SerializerFactory
@@ -30,7 +43,7 @@ abstract class AbstractAMQPSerializationScheme : SerializationScheme {
                     rpcServerSerializerFactory(context)
                 else -> SerializerFactory(context.whitelist) // TODO pass class loader also
             }
-        }
+        }.also { registerCustomSerializers(it) }
     }
 
     override fun <T : Any> deserialize(byteSequence: ByteSequence, clazz: Class<T>, context: SerializationContext): T {

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

@@ -6,10 +6,9 @@ import com.esotericsoftware.kryo.io.Output
 import com.esotericsoftware.kryo.serializers.FieldSerializer
 import com.esotericsoftware.kryo.util.DefaultClassResolver
 import com.esotericsoftware.kryo.util.Util
-import net.corda.core.serialization.AttachmentsClassLoader
+import net.corda.core.serialization.*
-import net.corda.core.serialization.ClassWhitelist
-import net.corda.core.serialization.CordaSerializable
 import net.corda.core.utilities.loggerFor
+import net.corda.nodeapi.internal.serialization.amqp.AmqpHeaderV1_0
 import java.io.PrintWriter
 import java.lang.reflect.Modifier.isAbstract
 import java.nio.charset.StandardCharsets
@@ -22,23 +21,13 @@ fun Kryo.addToWhitelist(type: Class<*>) {
     ((classResolver as? CordaClassResolver)?.whitelist as? MutableClassWhitelist)?.add(type)
 }
 
-fun makeStandardClassResolver(): ClassResolver {
-    return CordaClassResolver(GlobalTransientClassWhiteList(BuiltInExceptionsWhitelist()))
-}
-
-fun makeNoWhitelistClassResolver(): ClassResolver {
-    return CordaClassResolver(AllWhitelist)
-}
-
-fun makeAllButBlacklistedClassResolver(): ClassResolver {
-    return CordaClassResolver(AllButBlacklisted)
-}
-
 /**
  * @param amqpEnabled Setting this to true turns on experimental AMQP serialization for any class annotated with
  * [CordaSerializable].
  */
-class CordaClassResolver(val whitelist: ClassWhitelist, val amqpEnabled: Boolean = false) : DefaultClassResolver() {
+class CordaClassResolver(val serializationFactory: SerializationFactory, val serializationContext: SerializationContext) : DefaultClassResolver() {
+    val whitelist: ClassWhitelist = TransientClassWhiteList(serializationContext.whitelist)
+
     /** Returns the registration for the specified class, or null if the class is not registered.  */
     override fun getRegistration(type: Class<*>): Registration? {
         return super.getRegistration(type) ?: checkClass(type)
@@ -78,9 +67,9 @@ class CordaClassResolver(val whitelist: ClassWhitelist, val amqpEnabled: Boolean
         // If something is not annotated, or AMQP is disabled, we stay serializing with Kryo.  This will typically be the
         // case for flow checkpoints (ignoring all cases where AMQP is disabled) since our top level messaging data structures
         // are annotated and once we enter AMQP serialisation we stay with it for the entire object subgraph.
-        if (checkForAnnotation(type) && amqpEnabled) {
+        if (checkForAnnotation(type) && AMQP_ENABLED) {
             // Build AMQP serializer
-            return register(Registration(type, KryoAMQPSerializer, NAME.toInt()))
+            return register(Registration(type, KryoAMQPSerializer(serializationFactory, serializationContext), NAME.toInt()))
         }
 
         val objectInstance = try {
@@ -179,6 +168,21 @@ class GlobalTransientClassWhiteList(val delegate: ClassWhitelist) : MutableClass
     }
 }
 
+/**
+ * A whitelist that can be customised via the [CordaPluginRegistry], since implements [MutableClassWhitelist].
+ */
+class TransientClassWhiteList(val delegate: ClassWhitelist) : MutableClassWhitelist, ClassWhitelist by delegate {
+    val whitelist: MutableSet<String> = Collections.synchronizedSet(mutableSetOf())
+
+    override fun hasListed(type: Class<*>): Boolean {
+        return (type.name in whitelist) || delegate.hasListed(type)
+    }
+
+    override fun add(entry: Class<*>) {
+        whitelist += entry.name
+    }
+}
+
 
 /**
  * This class is not currently used, but can be installed to log a large number of missing entries from the whitelist

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

@@ -46,10 +46,7 @@ import kotlin.collections.ArrayList
 
 object DefaultKryoCustomizer {
     private val pluginRegistries: List<CordaPluginRegistry> by lazy {
-        // No ClassResolver only constructor.  MapReferenceResolver is the default as used by Kryo in other constructors.
+        ServiceLoader.load(CordaPluginRegistry::class.java, this.javaClass.classLoader).toList()
-        val unusedKryo = Kryo(makeStandardClassResolver(), MapReferenceResolver())
-        val customization = KryoSerializationCustomization(unusedKryo)
-        ServiceLoader.load(CordaPluginRegistry::class.java, this.javaClass.classLoader).toList().filter { it.customizeSerialization(customization) }
     }
 
     fun customize(kryo: Kryo): Kryo {

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

@@ -4,7 +4,11 @@ import com.esotericsoftware.kryo.Kryo
 import com.esotericsoftware.kryo.Serializer
 import com.esotericsoftware.kryo.io.Input
 import com.esotericsoftware.kryo.io.Output
+import net.corda.core.serialization.SerializationContext
+import net.corda.core.serialization.SerializationFactory
 import net.corda.core.serialization.SerializedBytes
+import net.corda.core.utilities.sequence
+import net.corda.nodeapi.internal.serialization.amqp.AmqpHeaderV1_0
 import net.corda.nodeapi.internal.serialization.amqp.DeserializationInput
 import net.corda.nodeapi.internal.serialization.amqp.SerializationOutput
 import net.corda.nodeapi.internal.serialization.amqp.SerializerFactory
@@ -15,38 +19,19 @@ import net.corda.nodeapi.internal.serialization.amqp.SerializerFactory
  *
  * There is no need to write out the length, since this can be peeked out of the first few bytes of the stream.
  */
-object KryoAMQPSerializer : Serializer<Any>() {
+class KryoAMQPSerializer(val serializationFactory: SerializationFactory, val serializationContext: SerializationContext) : Serializer<Any>() {
-    internal fun registerCustomSerializers(factory: SerializerFactory) {
-        factory.apply {
-            register(net.corda.nodeapi.internal.serialization.amqp.custom.PublicKeySerializer)
-            register(net.corda.nodeapi.internal.serialization.amqp.custom.ThrowableSerializer(this))
-            register(net.corda.nodeapi.internal.serialization.amqp.custom.X500NameSerializer)
-            register(net.corda.nodeapi.internal.serialization.amqp.custom.BigDecimalSerializer)
-            register(net.corda.nodeapi.internal.serialization.amqp.custom.CurrencySerializer)
-            register(net.corda.nodeapi.internal.serialization.amqp.custom.InstantSerializer(this))
-        }
-    }
-
-    // TODO: need to sort out the whitelist... we currently do not apply the whitelist attached to the [Kryo]
-    // instance to the factory.  We need to do this before turning on AMQP serialization.
-    private val serializerFactory = SerializerFactory().apply {
-        registerCustomSerializers(this)
-    }
-
     override fun write(kryo: Kryo, output: Output, obj: Any) {
-        val amqpOutput = SerializationOutput(serializerFactory)
+        val bytes = serializationFactory.serialize(obj, serializationContext.withPreferredSerializationVersion(AmqpHeaderV1_0)).bytes
-        val bytes = amqpOutput.serialize(obj).bytes
         // No need to write out the size since it's encoded within the AMQP.
         output.write(bytes)
     }
 
     override fun read(kryo: Kryo, input: Input, type: Class<Any>): Any {
-        val amqpInput = DeserializationInput(serializerFactory)
         // Use our helper functions to peek the size of the serialized object out of the AMQP byte stream.
         val peekedBytes = input.readBytes(DeserializationInput.BYTES_NEEDED_TO_PEEK)
         val size = DeserializationInput.peekSize(peekedBytes)
         val allBytes = peekedBytes.copyOf(size)
         input.readBytes(allBytes, peekedBytes.size, size - peekedBytes.size)
-        return amqpInput.deserialize(SerializedBytes<Any>(allBytes), type)
+        return serializationFactory.deserialize(allBytes.sequence(), type, serializationContext)
     }
 }

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

@@ -34,7 +34,6 @@ data class SerializationContextImpl(override val preferedSerializationVersion: B
                                     override val properties: Map<Any, Any>,
                                     override val objectReferencesEnabled: Boolean,
                                     override val useCase: SerializationContext.UseCase) : SerializationContext {
-
     override fun withProperty(property: Any, value: Any): SerializationContext {
         return copy(properties = properties + (property to value))
     }
@@ -52,6 +51,8 @@ data class SerializationContextImpl(override val preferedSerializationVersion: B
             override fun hasListed(type: Class<*>): Boolean = whitelist.hasListed(type) || type.name == clazz.name
         })
     }
+
+    override fun withPreferredSerializationVersion(versionHeader: ByteSequence) = copy(preferedSerializationVersion = versionHeader)
 }
 
 private const val HEADER_SIZE: Int = 8
@@ -118,7 +119,7 @@ private object AutoCloseableSerialisationDetector : Serializer<AutoCloseable>()
     override fun read(kryo: Kryo, input: Input, type: Class<AutoCloseable>) = throw IllegalStateException("Should not reach here!")
 }
 
-abstract class AbstractKryoSerializationScheme : SerializationScheme {
+abstract class AbstractKryoSerializationScheme(val serializationFactory: SerializationFactory) : SerializationScheme {
     private val kryoPoolsForContexts = ConcurrentHashMap<Pair<ClassWhitelist, ClassLoader>, KryoPool>()
 
     protected abstract fun rpcClientKryoPool(context: SerializationContext): KryoPool
@@ -130,7 +131,7 @@ abstract class AbstractKryoSerializationScheme : SerializationScheme {
                 SerializationContext.UseCase.Checkpoint ->
                     KryoPool.Builder {
                         val serializer = Fiber.getFiberSerializer(false) as KryoSerializer
-                        val classResolver = makeNoWhitelistClassResolver().apply { setKryo(serializer.kryo) }
+                        val classResolver = CordaClassResolver(serializationFactory, context).apply { setKryo(serializer.kryo) }
                         // TODO The ClassResolver can only be set in the Kryo constructor and Quasar doesn't provide us with a way of doing that
                         val field = Kryo::class.java.getDeclaredField("classResolver").apply { isAccessible = true }
                         serializer.kryo.apply {
@@ -146,7 +147,7 @@ abstract class AbstractKryoSerializationScheme : SerializationScheme {
                     rpcServerKryoPool(context)
                 else ->
                     KryoPool.Builder {
-                        DefaultKryoCustomizer.customize(CordaKryo(CordaClassResolver(context.whitelist))).apply { classLoader = it.second }
+                        DefaultKryoCustomizer.customize(CordaKryo(CordaClassResolver(serializationFactory, context))).apply { classLoader = it.second }
                     }.build()
             }
         }

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

@@ -5,8 +5,8 @@ import com.esotericsoftware.kryo.io.Input
 import com.esotericsoftware.kryo.io.Output
 import com.esotericsoftware.kryo.util.MapReferenceResolver
 import net.corda.core.node.services.AttachmentStorage
-import net.corda.core.serialization.AttachmentsClassLoader
+import net.corda.core.serialization.*
-import net.corda.core.serialization.CordaSerializable
+import net.corda.core.utilities.ByteSequence
 import net.corda.nodeapi.AttachmentClassLoaderTests
 import net.corda.testing.node.MockAttachmentStorage
 import org.junit.Rule
@@ -76,71 +76,84 @@ class DefaultSerializableSerializer : Serializer<DefaultSerializable>() {
 }
 
 class CordaClassResolverTests {
+    val factory: SerializationFactory = object : SerializationFactory {
+        override fun <T : Any> deserialize(byteSequence: ByteSequence, clazz: Class<T>, context: SerializationContext): T {
+            TODO("not implemented") //To change body of created functions use File | Settings | File Templates.
+        }
+
+        override fun <T : Any> serialize(obj: T, context: SerializationContext): SerializedBytes<T> {
+            TODO("not implemented") //To change body of created functions use File | Settings | File Templates.
+        }
+
+    }
+
+    val emptyWhitelistContext: SerializationContext = SerializationContextImpl(KryoHeaderV0_1, this.javaClass.classLoader, EmptyWhitelist, emptyMap(), true, SerializationContext.UseCase.P2P)
+    val allButBlacklistedContext: SerializationContext = SerializationContextImpl(KryoHeaderV0_1, this.javaClass.classLoader, AllButBlacklisted, emptyMap(), true, SerializationContext.UseCase.P2P)
+
     @Test
     fun `Annotation on enum works for specialised entries`() {
         // TODO: Remove this suppress when we upgrade to kotlin 1.1 or when JetBrain fixes the bug.
         @Suppress("UNSUPPORTED_FEATURE")
-        CordaClassResolver(EmptyWhitelist).getRegistration(Foo.Bar::class.java)
+        CordaClassResolver(factory, emptyWhitelistContext).getRegistration(Foo.Bar::class.java)
     }
 
     @Test
     fun `Annotation on array element works`() {
         val values = arrayOf(Element())
-        CordaClassResolver(EmptyWhitelist).getRegistration(values.javaClass)
+        CordaClassResolver(factory, emptyWhitelistContext).getRegistration(values.javaClass)
     }
 
     @Test
     fun `Annotation not needed on abstract class`() {
-        CordaClassResolver(EmptyWhitelist).getRegistration(AbstractClass::class.java)
+        CordaClassResolver(factory, emptyWhitelistContext).getRegistration(AbstractClass::class.java)
     }
 
     @Test
     fun `Annotation not needed on interface`() {
-        CordaClassResolver(EmptyWhitelist).getRegistration(Interface::class.java)
+        CordaClassResolver(factory, emptyWhitelistContext).getRegistration(Interface::class.java)
     }
 
     @Test
     fun `Calling register method on modified Kryo does not consult the whitelist`() {
-        val kryo = CordaKryo(CordaClassResolver(EmptyWhitelist))
+        val kryo = CordaKryo(CordaClassResolver(factory, emptyWhitelistContext))
         kryo.register(NotSerializable::class.java)
     }
 
     @Test(expected = KryoException::class)
     fun `Calling register method on unmodified Kryo does consult the whitelist`() {
-        val kryo = Kryo(CordaClassResolver(EmptyWhitelist), MapReferenceResolver())
+        val kryo = Kryo(CordaClassResolver(factory, emptyWhitelistContext), MapReferenceResolver())
         kryo.register(NotSerializable::class.java)
     }
 
     @Test(expected = KryoException::class)
     fun `Annotation is needed without whitelisting`() {
-        CordaClassResolver(EmptyWhitelist).getRegistration(NotSerializable::class.java)
+        CordaClassResolver(factory, emptyWhitelistContext).getRegistration(NotSerializable::class.java)
     }
 
     @Test
     fun `Annotation is not needed with whitelisting`() {
-        val resolver = CordaClassResolver(GlobalTransientClassWhiteList(EmptyWhitelist))
+        val resolver = CordaClassResolver(factory, emptyWhitelistContext.withWhitelisted(NotSerializable::class.java))
-        (resolver.whitelist as MutableClassWhitelist).add(NotSerializable::class.java)
         resolver.getRegistration(NotSerializable::class.java)
     }
 
     @Test
     fun `Annotation not needed on Object`() {
-        CordaClassResolver(EmptyWhitelist).getRegistration(Object::class.java)
+        CordaClassResolver(factory, emptyWhitelistContext).getRegistration(Object::class.java)
     }
 
     @Test
     fun `Annotation not needed on primitive`() {
-        CordaClassResolver(EmptyWhitelist).getRegistration(Integer.TYPE)
+        CordaClassResolver(factory, emptyWhitelistContext).getRegistration(Integer.TYPE)
     }
 
     @Test(expected = KryoException::class)
     fun `Annotation does not work for custom serializable`() {
-        CordaClassResolver(EmptyWhitelist).getRegistration(CustomSerializable::class.java)
+        CordaClassResolver(factory, emptyWhitelistContext).getRegistration(CustomSerializable::class.java)
     }
 
     @Test(expected = KryoException::class)
     fun `Annotation does not work in conjunction with Kryo annotation`() {
-        CordaClassResolver(EmptyWhitelist).getRegistration(DefaultSerializable::class.java)
+        CordaClassResolver(factory, emptyWhitelistContext).getRegistration(DefaultSerializable::class.java)
     }
 
     private fun importJar(storage: AttachmentStorage) = AttachmentClassLoaderTests.ISOLATED_CONTRACTS_JAR_PATH.openStream().use { storage.importAttachment(it) }
@@ -151,20 +164,20 @@ class CordaClassResolverTests {
         val attachmentHash = importJar(storage)
         val classLoader = AttachmentsClassLoader(arrayOf(attachmentHash).map { storage.openAttachment(it)!! })
         val attachedClass = Class.forName("net.corda.contracts.isolated.AnotherDummyContract", true, classLoader)
-        CordaClassResolver(EmptyWhitelist).getRegistration(attachedClass)
+        CordaClassResolver(factory, emptyWhitelistContext).getRegistration(attachedClass)
     }
 
     @Test
     fun `Annotation is inherited from interfaces`() {
-        CordaClassResolver(EmptyWhitelist).getRegistration(SerializableViaInterface::class.java)
+        CordaClassResolver(factory, emptyWhitelistContext).getRegistration(SerializableViaInterface::class.java)
-        CordaClassResolver(EmptyWhitelist).getRegistration(SerializableViaSubInterface::class.java)
+        CordaClassResolver(factory, emptyWhitelistContext).getRegistration(SerializableViaSubInterface::class.java)
     }
 
     @Test
     fun `Annotation is inherited from superclass`() {
-        CordaClassResolver(EmptyWhitelist).getRegistration(SubElement::class.java)
+        CordaClassResolver(factory, emptyWhitelistContext).getRegistration(SubElement::class.java)
-        CordaClassResolver(EmptyWhitelist).getRegistration(SubSubElement::class.java)
+        CordaClassResolver(factory, emptyWhitelistContext).getRegistration(SubSubElement::class.java)
-        CordaClassResolver(EmptyWhitelist).getRegistration(SerializableViaSuperSubInterface::class.java)
+        CordaClassResolver(factory, emptyWhitelistContext).getRegistration(SerializableViaSuperSubInterface::class.java)
     }
 
     // Blacklist tests.
@@ -175,7 +188,7 @@ class CordaClassResolverTests {
     fun `Check blacklisted class`() {
         expectedEx.expect(IllegalStateException::class.java)
         expectedEx.expectMessage("Class java.util.HashSet is blacklisted, so it cannot be used in serialization.")
-        val resolver = CordaClassResolver(AllButBlacklisted)
+        val resolver = CordaClassResolver(factory, allButBlacklistedContext)
         // HashSet is blacklisted.
         resolver.getRegistration(HashSet::class.java)
     }
@@ -185,7 +198,7 @@ class CordaClassResolverTests {
     fun `Check blacklisted subclass`() {
         expectedEx.expect(IllegalStateException::class.java)
         expectedEx.expectMessage("The superclass java.util.HashSet of net.corda.nodeapi.internal.serialization.CordaClassResolverTests\$SubHashSet is blacklisted, so it cannot be used in serialization.")
-        val resolver = CordaClassResolver(AllButBlacklisted)
+        val resolver = CordaClassResolver(factory, allButBlacklistedContext)
         // SubHashSet extends the blacklisted HashSet.
         resolver.getRegistration(SubHashSet::class.java)
     }
@@ -195,7 +208,7 @@ class CordaClassResolverTests {
     fun `Check blacklisted subsubclass`() {
         expectedEx.expect(IllegalStateException::class.java)
         expectedEx.expectMessage("The superclass java.util.HashSet of net.corda.nodeapi.internal.serialization.CordaClassResolverTests\$SubSubHashSet is blacklisted, so it cannot be used in serialization.")
-        val resolver = CordaClassResolver(AllButBlacklisted)
+        val resolver = CordaClassResolver(factory, allButBlacklistedContext)
         // SubSubHashSet extends SubHashSet, which extends the blacklisted HashSet.
         resolver.getRegistration(SubSubHashSet::class.java)
     }
@@ -205,7 +218,7 @@ class CordaClassResolverTests {
     fun `Check blacklisted interface impl`() {
         expectedEx.expect(IllegalStateException::class.java)
         expectedEx.expectMessage("The superinterface java.sql.Connection of net.corda.nodeapi.internal.serialization.CordaClassResolverTests\$ConnectionImpl is blacklisted, so it cannot be used in serialization.")
-        val resolver = CordaClassResolver(AllButBlacklisted)
+        val resolver = CordaClassResolver(factory, allButBlacklistedContext)
         // ConnectionImpl implements blacklisted Connection.
         resolver.getRegistration(ConnectionImpl::class.java)
     }
@@ -216,14 +229,14 @@ class CordaClassResolverTests {
     fun `Check blacklisted super-interface impl`() {
         expectedEx.expect(IllegalStateException::class.java)
         expectedEx.expectMessage("The superinterface java.sql.Connection of net.corda.nodeapi.internal.serialization.CordaClassResolverTests\$SubConnectionImpl is blacklisted, so it cannot be used in serialization.")
-        val resolver = CordaClassResolver(AllButBlacklisted)
+        val resolver = CordaClassResolver(factory, allButBlacklistedContext)
         // SubConnectionImpl implements SubConnection, which extends the blacklisted Connection.
         resolver.getRegistration(SubConnectionImpl::class.java)
     }
 
     @Test
     fun `Check forcibly allowed`() {
-        val resolver = CordaClassResolver(AllButBlacklisted)
+        val resolver = CordaClassResolver(factory, allButBlacklistedContext)
         // LinkedHashSet is allowed for serialization.
         resolver.getRegistration(LinkedHashSet::class.java)
     }
@@ -234,7 +247,7 @@ class CordaClassResolverTests {
     fun `Check blacklist precedes CordaSerializable`() {
         expectedEx.expect(IllegalStateException::class.java)
         expectedEx.expectMessage("The superclass java.util.HashSet of net.corda.nodeapi.internal.serialization.CordaClassResolverTests\$CordaSerializableHashSet is blacklisted, so it cannot be used in serialization.")
-        val resolver = CordaClassResolver(AllButBlacklisted)
+        val resolver = CordaClassResolver(factory, allButBlacklistedContext)
         // CordaSerializableHashSet is @CordaSerializable, but extends the blacklisted HashSet.
         resolver.getRegistration(CordaSerializableHashSet::class.java)
     }

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

@@ -12,6 +12,7 @@ import net.corda.core.utilities.sequence
 import net.corda.node.serialization.KryoServerSerializationScheme
 import net.corda.node.services.persistence.NodeAttachmentService
 import net.corda.testing.ALICE_PUBKEY
+import net.corda.testing.TestDependencyInjectionBase
 import org.assertj.core.api.Assertions.assertThat
 import org.assertj.core.api.Assertions.assertThatThrownBy
 import org.junit.Before
@@ -23,13 +24,13 @@ import java.time.Instant
 import kotlin.test.assertEquals
 import kotlin.test.assertTrue
 
-class KryoTests {
+class KryoTests : TestDependencyInjectionBase() {
     private lateinit var factory: SerializationFactory
     private lateinit var context: SerializationContext
 
     @Before
     fun setup() {
-        factory = SerializationFactoryImpl().apply { registerScheme(KryoServerSerializationScheme()) }
+        factory = SerializationFactoryImpl().apply { registerScheme(KryoServerSerializationScheme(this)) }
         context = SerializationContextImpl(KryoHeaderV0_1,
                 javaClass.classLoader,
                 AllWhitelist,
@@ -199,7 +200,7 @@ class KryoTests {
             }
         }
         Tmp()
-        val factory = SerializationFactoryImpl().apply { registerScheme(KryoServerSerializationScheme()) }
+        val factory = SerializationFactoryImpl().apply { registerScheme(KryoServerSerializationScheme(this)) }
         val context = SerializationContextImpl(KryoHeaderV0_1,
                 javaClass.classLoader,
                 AllWhitelist,

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

@@ -8,19 +8,20 @@ import net.corda.core.node.ServiceHub
 import net.corda.core.serialization.*
 import net.corda.core.utilities.OpaqueBytes
 import net.corda.node.serialization.KryoServerSerializationScheme
+import net.corda.testing.TestDependencyInjectionBase
 import org.assertj.core.api.Assertions.assertThat
 import org.junit.Before
 import org.junit.Test
 import java.io.ByteArrayOutputStream
 
-class SerializationTokenTest {
+class SerializationTokenTest  : TestDependencyInjectionBase() {
 
     lateinit var factory: SerializationFactory
     lateinit var context: SerializationContext
 
     @Before
     fun setup() {
-        factory = SerializationFactoryImpl().apply { registerScheme(KryoServerSerializationScheme()) }
+        factory = SerializationFactoryImpl().apply { registerScheme(KryoServerSerializationScheme(this)) }
         context = SerializationContextImpl(KryoHeaderV0_1,
                 javaClass.classLoader,
                 AllWhitelist,
@@ -96,7 +97,7 @@ class SerializationTokenTest {
         val context = serializeAsTokenContext(tokenizableBefore)
         val testContext = this.context.withTokenContext(context)
 
-        val kryo: Kryo = DefaultKryoCustomizer.customize(CordaKryo(makeNoWhitelistClassResolver()))
+        val kryo: Kryo = DefaultKryoCustomizer.customize(CordaKryo(CordaClassResolver(factory, this.context)))
         val stream = ByteArrayOutputStream()
             Output(stream).use {
                 it.write(KryoHeaderV0_1.bytes)

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

@@ -11,8 +11,8 @@ import net.corda.core.identity.AbstractParty
 import net.corda.core.serialization.CordaSerializable
 import net.corda.core.transactions.LedgerTransaction
 import net.corda.nodeapi.RPCException
+import net.corda.nodeapi.internal.serialization.AbstractAMQPSerializationScheme
 import net.corda.nodeapi.internal.serialization.EmptyWhitelist
-import net.corda.nodeapi.internal.serialization.KryoAMQPSerializer
 import net.corda.nodeapi.internal.serialization.amqp.SerializerFactory.Companion.isPrimitive
 import net.corda.nodeapi.internal.serialization.amqp.custom.*
 import net.corda.testing.MEGA_CORP
@@ -528,10 +528,10 @@ class SerializationOutputTests {
         val state = TransactionState<FooState>(FooState(), MEGA_CORP)
 
         val factory = SerializerFactory()
-        KryoAMQPSerializer.registerCustomSerializers(factory)
+        AbstractAMQPSerializationScheme.registerCustomSerializers(factory)
 
         val factory2 = SerializerFactory()
-        KryoAMQPSerializer.registerCustomSerializers(factory2)
+        AbstractAMQPSerializationScheme.registerCustomSerializers(factory2)
 
         val desState = serdes(state, factory, factory2, expectedEqual = false, expectDeserializedEqual = false)
         assertTrue(desState is TransactionState<*>)

node/src/main/kotlin/net/corda/node/internal/Node.kt

@@ -331,7 +331,7 @@ open class Node(override val configuration: FullNodeConfiguration,
 
     private fun initialiseSerialization() {
         SerializationDefaults.SERIALIZATION_FACTORY = SerializationFactoryImpl().apply {
-            registerScheme(KryoServerSerializationScheme())
+            registerScheme(KryoServerSerializationScheme(this))
             registerScheme(AMQPServerSerializationScheme())
         }
         SerializationDefaults.P2P_CONTEXT = KRYO_P2P_CONTEXT

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

@@ -2,6 +2,7 @@ package net.corda.node.serialization
 
 import com.esotericsoftware.kryo.pool.KryoPool
 import net.corda.core.serialization.SerializationContext
+import net.corda.core.serialization.SerializationFactory
 import net.corda.core.utilities.ByteSequence
 import net.corda.node.services.messaging.RpcServerObservableSerializer
 import net.corda.nodeapi.RPCKryo
@@ -9,7 +10,7 @@ import net.corda.nodeapi.internal.serialization.AbstractKryoSerializationScheme
 import net.corda.nodeapi.internal.serialization.DefaultKryoCustomizer
 import net.corda.nodeapi.internal.serialization.KryoHeaderV0_1
 
-class KryoServerSerializationScheme : AbstractKryoSerializationScheme() {
+class KryoServerSerializationScheme(serializationFactory: SerializationFactory) : AbstractKryoSerializationScheme(serializationFactory) {
     override fun canDeserializeVersion(byteSequence: ByteSequence, target: SerializationContext.UseCase): Boolean {
         return byteSequence == KryoHeaderV0_1 && target != SerializationContext.UseCase.RPCClient
     }
@@ -20,7 +21,7 @@ class KryoServerSerializationScheme : AbstractKryoSerializationScheme() {
 
     override fun rpcServerKryoPool(context: SerializationContext): KryoPool {
         return KryoPool.Builder {
-            DefaultKryoCustomizer.customize(RPCKryo(RpcServerObservableSerializer, context.whitelist)).apply { classLoader = context.deserializationClassLoader }
+            DefaultKryoCustomizer.customize(RPCKryo(RpcServerObservableSerializer, serializationFactory, context)).apply { classLoader = context.deserializationClassLoader }
         }.build()
     }
 }

node/src/test/kotlin/net/corda/node/utilities/X509UtilitiesTest.kt

@@ -398,7 +398,7 @@ class X509UtilitiesTest {
 
     @Test
     fun `serialize - deserialize X509CertififcateHolder`() {
-        val factory = SerializationFactoryImpl().apply { registerScheme(KryoServerSerializationScheme()) }
+        val factory = SerializationFactoryImpl().apply { registerScheme(KryoServerSerializationScheme(this)) }
         val context = SerializationContextImpl(KryoHeaderV0_1,
                 javaClass.classLoader,
                 AllWhitelist,
@@ -413,7 +413,7 @@ class X509UtilitiesTest {
 
     @Test
     fun `serialize - deserialize X509CertPath`() {
-        val factory = SerializationFactoryImpl().apply { registerScheme(KryoServerSerializationScheme()) }
+        val factory = SerializationFactoryImpl().apply { registerScheme(KryoServerSerializationScheme(this)) }
         val context = SerializationContextImpl(KryoHeaderV0_1,
                 javaClass.classLoader,
                 AllWhitelist,

test-utils/src/main/kotlin/net/corda/testing/SerializationTestHelpers.kt

@@ -61,8 +61,8 @@ fun initialiseTestSerialization() {
 
     // Now configure all the testing related delegates.
     (SerializationDefaults.SERIALIZATION_FACTORY as TestSerializationFactory).delegate = SerializationFactoryImpl().apply {
-        registerScheme(KryoClientSerializationScheme())
+        registerScheme(KryoClientSerializationScheme(this))
-        registerScheme(KryoServerSerializationScheme())
+        registerScheme(KryoServerSerializationScheme(this))
         registerScheme(AMQPClientSerializationScheme())
         registerScheme(AMQPServerSerializationScheme())
     }
@@ -139,4 +139,8 @@ class TestSerializationContext : SerializationContext {
     override fun withWhitelisted(clazz: Class<*>): SerializationContext {
         return TestSerializationContext().apply { delegate = this@TestSerializationContext.delegate!!.withWhitelisted(clazz) }
     }
+
+    override fun withPreferredSerializationVersion(versionHeader: ByteSequence): SerializationContext {
+        return TestSerializationContext().apply { delegate = this@TestSerializationContext.delegate!!.withPreferredSerializationVersion(versionHeader) }
+    }
 }

verifier/src/main/kotlin/net/corda/verifier/Verifier.kt

@@ -7,6 +7,7 @@ import com.typesafe.config.ConfigParseOptions
 import net.corda.core.internal.div
 import net.corda.core.serialization.SerializationContext
 import net.corda.core.serialization.SerializationDefaults
+import net.corda.core.serialization.SerializationFactory
 import net.corda.core.utilities.ByteSequence
 import net.corda.core.utilities.NetworkHostAndPort
 import net.corda.core.utilities.debug
@@ -89,13 +90,13 @@ class Verifier {
 
         private fun initialiseSerialization() {
             SerializationDefaults.SERIALIZATION_FACTORY = SerializationFactoryImpl().apply {
-                registerScheme(KryoVerifierSerializationScheme)
+                registerScheme(KryoVerifierSerializationScheme(this))
             }
             SerializationDefaults.P2P_CONTEXT = KRYO_P2P_CONTEXT
         }
     }
 
-    object KryoVerifierSerializationScheme : AbstractKryoSerializationScheme() {
+    class KryoVerifierSerializationScheme(serializationFactory: SerializationFactory) : AbstractKryoSerializationScheme(serializationFactory) {
         override fun canDeserializeVersion(byteSequence: ByteSequence, target: SerializationContext.UseCase): Boolean {
             return byteSequence.equals(KryoHeaderV0_1) && target == SerializationContext.UseCase.P2P
         }