From ead4d0901182c1b75df66acdb14592f5abe424ac Mon Sep 17 00:00:00 2001
From: Rick Parker <rick.parker@r3.com>
Date: Tue, 5 Sep 2017 08:35:26 +0100
Subject: [PATCH] Plumb in plugin whitelist additions to AMQP implementation.
 (#1415)

---
 .../serialization/AMQPSerializationScheme.kt  | 25 ++++++++++++++++---
 .../serialization/DefaultWhitelist.kt         |  3 ++-
 2 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/node-api/src/main/kotlin/net/corda/nodeapi/internal/serialization/AMQPSerializationScheme.kt b/node-api/src/main/kotlin/net/corda/nodeapi/internal/serialization/AMQPSerializationScheme.kt
index 965755d253..21566a0b88 100644
--- a/node-api/src/main/kotlin/net/corda/nodeapi/internal/serialization/AMQPSerializationScheme.kt
+++ b/node-api/src/main/kotlin/net/corda/nodeapi/internal/serialization/AMQPSerializationScheme.kt
@@ -1,20 +1,35 @@
 package net.corda.nodeapi.internal.serialization
 
-import net.corda.core.serialization.ClassWhitelist
-import net.corda.core.serialization.SerializationContext
-import net.corda.core.serialization.SerializationDefaults
-import net.corda.core.serialization.SerializedBytes
+import net.corda.core.node.CordaPluginRegistry
+import net.corda.core.serialization.*
 import net.corda.core.utilities.ByteSequence
 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
+import java.util.*
 import java.util.concurrent.ConcurrentHashMap
 
 val AMQP_ENABLED get() = SerializationDefaults.P2P_CONTEXT.preferredSerializationVersion == AmqpHeaderV1_0
 
+class AMQPSerializationCustomization(val factory: SerializerFactory) : SerializationCustomization {
+    override fun addToWhitelist(vararg types: Class<*>) {
+        factory.addToWhitelist(*types)
+    }
+}
+
+fun SerializerFactory.addToWhitelist(vararg types: Class<*>) {
+    for (type in types) {
+        (this.whitelist as? MutableClassWhitelist)?.add(type)
+    }
+}
+
 abstract class AbstractAMQPSerializationScheme : SerializationScheme {
     internal companion object {
+        private val pluginRegistries: List<CordaPluginRegistry> by lazy {
+            ServiceLoader.load(CordaPluginRegistry::class.java, this.javaClass.classLoader).toList()
+        }
+
         fun registerCustomSerializers(factory: SerializerFactory) {
             factory.apply {
                 register(net.corda.nodeapi.internal.serialization.amqp.custom.PublicKeySerializer)
@@ -40,6 +55,8 @@ abstract class AbstractAMQPSerializationScheme : SerializationScheme {
                 register(net.corda.nodeapi.internal.serialization.amqp.custom.X509CertificateHolderSerializer)
                 register(net.corda.nodeapi.internal.serialization.amqp.custom.PartyAndCertificateSerializer(factory))
             }
+            val customizer = AMQPSerializationCustomization(factory)
+            pluginRegistries.forEach { it.customizeSerialization(customizer) }
         }
     }
 
diff --git a/node-api/src/main/kotlin/net/corda/nodeapi/internal/serialization/DefaultWhitelist.kt b/node-api/src/main/kotlin/net/corda/nodeapi/internal/serialization/DefaultWhitelist.kt
index d44c063c43..f2e3f3835e 100644
--- a/node-api/src/main/kotlin/net/corda/nodeapi/internal/serialization/DefaultWhitelist.kt
+++ b/node-api/src/main/kotlin/net/corda/nodeapi/internal/serialization/DefaultWhitelist.kt
@@ -64,7 +64,8 @@ class DefaultWhitelist : CordaPluginRegistry() {
                     LocalDate::class.java,
                     Period::class.java,
                     BitSet::class.java,
-                    OnErrorNotImplementedException::class.java)
+                    OnErrorNotImplementedException::class.java,
+                    StackTraceElement::class.java)
         }
         return true
     }