Merge pull request #2024 from corda/christians_public-key-caching-hooks

Public key customization hooks

.idea/compiler.xml

@@ -40,6 +40,9 @@
       <module name="docs_source_example-code_main" target="1.8" />
       <module name="docs_source_example-code_test" target="1.8" />
       <module name="docs_test" target="1.8" />
+      <module name="example-code_integrationTest" target="1.8" />
+      <module name="example-code_main" target="1.8" />
+      <module name="example-code_test" target="1.8" />
       <module name="experimental_main" target="1.8" />
       <module name="experimental_test" target="1.8" />
       <module name="explorer-capsule_main" target="1.6" />

client/rpc/src/main/kotlin/net/corda/client/rpc/internal/KryoClientSerializationScheme.kt

@@ -19,7 +19,7 @@ class KryoClientSerializationScheme : AbstractKryoSerializationScheme() {
 
     override fun rpcClientKryoPool(context: SerializationContext): KryoPool {
         return KryoPool.Builder {
-            DefaultKryoCustomizer.customize(RPCKryo(RpcClientObservableSerializer, context)).apply {
+            DefaultKryoCustomizer.customize(RPCKryo(RpcClientObservableSerializer, context), publicKeySerializer).apply {
                 classLoader = context.deserializationClassLoader
             }
         }.build()

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

@@ -7,6 +7,7 @@ import net.corda.core.utilities.ByteSequence
 import net.corda.nodeapi.internal.serialization.DefaultWhitelist
 import net.corda.nodeapi.internal.serialization.MutableClassWhitelist
 import net.corda.nodeapi.internal.serialization.SerializationScheme
+import java.security.PublicKey
 import java.util.*
 import java.util.concurrent.ConcurrentHashMap
 
@@ -28,10 +29,11 @@ abstract class AbstractAMQPSerializationScheme : SerializationScheme {
         private val serializationWhitelists: List<SerializationWhitelist> by lazy {
             ServiceLoader.load(SerializationWhitelist::class.java, this::class.java.classLoader).toList() + DefaultWhitelist
         }
+    }
 
-        fun registerCustomSerializers(factory: SerializerFactory) {
+    private fun registerCustomSerializers(factory: SerializerFactory) {
         with(factory) {
-                register(net.corda.nodeapi.internal.serialization.amqp.custom.PublicKeySerializer)
+            register(publicKeySerializer)
             register(net.corda.nodeapi.internal.serialization.amqp.custom.PrivateKeySerializer)
             register(net.corda.nodeapi.internal.serialization.amqp.custom.ThrowableSerializer(this))
             register(net.corda.nodeapi.internal.serialization.amqp.custom.X500NameSerializer)
@@ -63,12 +65,13 @@ abstract class AbstractAMQPSerializationScheme : SerializationScheme {
         for (whitelistProvider in serializationWhitelists)
             factory.addToWhitelist(*whitelistProvider.whitelist.toTypedArray())
     }
-    }
 
     private val serializerFactoriesForContexts = ConcurrentHashMap<Pair<ClassWhitelist, ClassLoader>, SerializerFactory>()
 
     protected abstract fun rpcClientSerializerFactory(context: SerializationContext): SerializerFactory
     protected abstract fun rpcServerSerializerFactory(context: SerializationContext): SerializerFactory
+    open protected val publicKeySerializer: CustomSerializer.Implements<PublicKey>
+            = net.corda.nodeapi.internal.serialization.amqp.custom.PublicKeySerializer
 
     private fun getSerializerFactory(context: SerializationContext): SerializerFactory {
         return serializerFactoriesForContexts.computeIfAbsent(Pair(context.whitelist, context.deserializationClassLoader)) {

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

@@ -51,6 +51,7 @@ import java.io.ByteArrayOutputStream
 import java.io.FileInputStream
 import java.io.InputStream
 import java.lang.reflect.Modifier.isPublic
+import java.security.PublicKey
 import java.security.cert.CertPath
 import java.util.*
 import kotlin.collections.ArrayList
@@ -60,7 +61,7 @@ object DefaultKryoCustomizer {
         ServiceLoader.load(SerializationWhitelist::class.java, this.javaClass.classLoader).toList() + DefaultWhitelist
     }
 
-    fun customize(kryo: Kryo): Kryo {
+    fun customize(kryo: Kryo, publicKeySerializer: Serializer<PublicKey> = PublicKeySerializer): Kryo {
         return kryo.apply {
             // Store a little schema of field names in the stream the first time a class is used which increases tolerance
             // for change to a class.
@@ -95,10 +96,10 @@ object DefaultKryoCustomizer {
             register(BufferedInputStream::class.java, InputStreamSerializer)
             register(Class.forName("sun.net.www.protocol.jar.JarURLConnection\$JarURLInputStream"), InputStreamSerializer)
             noReferencesWithin<WireTransaction>()
-            register(ECPublicKeyImpl::class.java, PublicKeySerializer)
+            register(ECPublicKeyImpl::class.java, publicKeySerializer)
-            register(EdDSAPublicKey::class.java, PublicKeySerializer)
+            register(EdDSAPublicKey::class.java, publicKeySerializer)
             register(EdDSAPrivateKey::class.java, PrivateKeySerializer)
-            register(CompositeKey::class.java, PublicKeySerializer)  // Using a custom serializer for compactness
+            register(CompositeKey::class.java, publicKeySerializer)  // Using a custom serializer for compactness
             // Exceptions. We don't bother sending the stack traces as the client will fill in its own anyway.
             register(Array<StackTraceElement>::class, read = { _, _ -> emptyArray() }, write = { _, _, _ -> })
             // This ensures a NonEmptySetSerializer is constructed with an initial value.
@@ -111,11 +112,11 @@ object DefaultKryoCustomizer {
             register(X500Name::class.java, X500NameSerializer)
             register(X509CertificateHolder::class.java, X509CertificateSerializer)
             register(BCECPrivateKey::class.java, PrivateKeySerializer)
-            register(BCECPublicKey::class.java, PublicKeySerializer)
+            register(BCECPublicKey::class.java, publicKeySerializer)
             register(BCRSAPrivateCrtKey::class.java, PrivateKeySerializer)
-            register(BCRSAPublicKey::class.java, PublicKeySerializer)
+            register(BCRSAPublicKey::class.java, publicKeySerializer)
             register(BCSphincs256PrivateKey::class.java, PrivateKeySerializer)
-            register(BCSphincs256PublicKey::class.java, PublicKeySerializer)
+            register(BCSphincs256PublicKey::class.java, publicKeySerializer)
             register(NotaryChangeWireTransaction::class.java, NotaryChangeWireTransactionSerializer)
             register(PartyAndCertificate::class.java, PartyAndCertificateSerializer)
 

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

@@ -18,6 +18,7 @@ import net.corda.core.serialization.*
 import net.corda.core.internal.LazyPool
 import net.corda.nodeapi.internal.serialization.CordaClassResolver
 import net.corda.nodeapi.internal.serialization.SerializationScheme
+import java.security.PublicKey
 
 // "corda" + majorVersionByte + minorVersionMSB + minorVersionLSB
 val KryoHeaderV0_1: OpaqueBytes = OpaqueBytes("corda\u0000\u0000\u0001".toByteArray(Charsets.UTF_8))
@@ -39,6 +40,9 @@ abstract class AbstractKryoSerializationScheme : SerializationScheme {
     protected abstract fun rpcClientKryoPool(context: SerializationContext): KryoPool
     protected abstract fun rpcServerKryoPool(context: SerializationContext): KryoPool
 
+    // this can be overriden in derived serialization schemes
+    open protected val publicKeySerializer: Serializer<PublicKey> = PublicKeySerializer
+
     private fun getPool(context: SerializationContext): KryoPool {
         return kryoPoolsForContexts.computeIfAbsent(Pair(context.whitelist, context.deserializationClassLoader)) {
             when (context.useCase) {
@@ -50,6 +54,7 @@ abstract class AbstractKryoSerializationScheme : SerializationScheme {
                         val field = Kryo::class.java.getDeclaredField("classResolver").apply { isAccessible = true }
                         serializer.kryo.apply {
                             field.set(this, classResolver)
+                            // don't allow overriding the public key serializer for checkpointing
                             DefaultKryoCustomizer.customize(this)
                             addDefaultSerializer(AutoCloseable::class.java, AutoCloseableSerialisationDetector)
                             register(ClosureSerializer.Closure::class.java, CordaClosureSerializer)
@@ -62,7 +67,7 @@ abstract class AbstractKryoSerializationScheme : SerializationScheme {
                     rpcServerKryoPool(context)
                 else ->
                     KryoPool.Builder {
-                        DefaultKryoCustomizer.customize(CordaKryo(CordaClassResolver(context))).apply { classLoader = it.second }
+                        DefaultKryoCustomizer.customize(CordaKryo(CordaClassResolver(context)), publicKeySerializer).apply { classLoader = it.second }
                     }.build()
             }
         }

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

@@ -0,0 +1,65 @@
+package net.corda.nodeapi.internal.serialization.amqp
+
+import net.corda.core.serialization.SerializationContext
+import net.corda.core.utilities.ByteSequence
+import net.corda.nodeapi.internal.serialization.AMQP_P2P_CONTEXT
+import org.apache.qpid.proton.codec.Data
+import org.assertj.core.api.Assertions
+import org.junit.Test
+import java.lang.reflect.Type
+import java.security.PublicKey
+
+class OverridePKSerializerTest {
+    class SerializerTestException(message: String) : Exception(message)
+
+    class TestPublicKeySerializer : CustomSerializer.Implements<PublicKey>(PublicKey::class.java) {
+        override fun writeDescribedObject(obj: PublicKey, data: Data, type: Type, output: SerializationOutput) {
+            throw SerializerTestException("Custom write call")
+        }
+
+        override fun readObject(obj: Any, schema: Schema, input: DeserializationInput): PublicKey {
+            throw SerializerTestException("Custom read call")
+        }
+
+        override val schemaForDocumentation: Schema
+            get() = TODO("not implemented") //To change initializer of created properties use File | Settings | File Templates.
+    }
+
+    class AMQPTestSerializationScheme : AbstractAMQPSerializationScheme() {
+        override fun rpcServerSerializerFactory(context: SerializationContext): SerializerFactory {
+            TODO("not implemented") //To change body of created functions use File | Settings | File Templates.
+        }
+
+        override fun canDeserializeVersion(byteSequence: ByteSequence, target: SerializationContext.UseCase): Boolean = true
+
+        override fun rpcClientSerializerFactory(context: SerializationContext): SerializerFactory {
+            TODO("not implemented") //To change body of created functions use File | Settings | File Templates.
+        }
+
+        override val publicKeySerializer = TestPublicKeySerializer()
+    }
+
+    class TestPublicKey : PublicKey {
+        override fun getAlgorithm(): String {
+            TODO("not implemented") //To change body of created functions use File | Settings | File Templates.
+        }
+
+        override fun getEncoded(): ByteArray {
+            TODO("not implemented") //To change body of created functions use File | Settings | File Templates.
+        }
+
+        override fun getFormat(): String {
+            TODO("not implemented") //To change body of created functions use File | Settings | File Templates.
+        }
+    }
+
+    @Test
+    fun `test publicKeySerializer is overridden`() {
+        val scheme = AMQPTestSerializationScheme()
+        val key = TestPublicKey()
+
+        Assertions
+                .assertThatThrownBy { scheme.serialize(key, AMQP_P2P_CONTEXT) }
+                .hasMessageMatching("Custom write call")
+    }
+}

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

@@ -36,6 +36,10 @@ import java.nio.ByteBuffer
 import java.time.*
 import java.time.temporal.ChronoUnit
 import java.util.*
+import kotlin.reflect.full.declaredFunctions
+import kotlin.reflect.full.declaredMemberFunctions
+import kotlin.reflect.full.superclasses
+import kotlin.reflect.jvm.javaMethod
 import kotlin.test.assertEquals
 import kotlin.test.assertNotNull
 import kotlin.test.assertTrue
@@ -557,11 +561,16 @@ class SerializationOutputTests {
     fun `test transaction state`() {
         val state = TransactionState(FooState(), FOO_PROGRAM_ID, MEGA_CORP)
 
+        val scheme = AMQPServerSerializationScheme()
+        val func = scheme::class.superclasses.single { it.simpleName == "AbstractAMQPSerializationScheme" }
+                .java.getDeclaredMethod("registerCustomSerializers", SerializerFactory::class.java)
+        func.isAccessible = true
+
         val factory = SerializerFactory(AllWhitelist, ClassLoader.getSystemClassLoader())
-        AbstractAMQPSerializationScheme.registerCustomSerializers(factory)
+        func.invoke(scheme, factory)
 
         val factory2 = SerializerFactory(AllWhitelist, ClassLoader.getSystemClassLoader())
-        AbstractAMQPSerializationScheme.registerCustomSerializers(factory2)
+        func.invoke(scheme, factory2)
 
         val desState = serdes(state, factory, factory2, expectedEqual = false, expectDeserializedEqual = false)
         assertTrue((desState as TransactionState<*>).data is FooState)

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

@@ -18,7 +18,7 @@ class KryoServerSerializationScheme : AbstractKryoSerializationScheme() {
 
     override fun rpcServerKryoPool(context: SerializationContext): KryoPool {
         return KryoPool.Builder {
-            DefaultKryoCustomizer.customize(RPCKryo(RpcServerObservableSerializer, context)).apply {
+            DefaultKryoCustomizer.customize(RPCKryo(RpcServerObservableSerializer, context), publicKeySerializer).apply {
                 classLoader = context.deserializationClassLoader
             }
         }.build()