Replace code only used in 1 test with existing general mechanism. (#600)

core/src/main/kotlin/net/corda/core/node/services/Services.kt

@@ -328,6 +328,7 @@ interface FileUploader {
 interface AttachmentsStorageService {
     /** Provides access to storage of arbitrary JAR files (which may contain only data, no code). */
     val attachments: AttachmentStorage
+    val attachmentsClassLoaderEnabled: Boolean
 }
 
 /**

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

@@ -4,20 +4,17 @@ import com.esotericsoftware.kryo.*
 import com.esotericsoftware.kryo.io.Input
 import com.esotericsoftware.kryo.io.Output
 import com.esotericsoftware.kryo.pool.KryoPool
-import com.esotericsoftware.kryo.serializers.JavaSerializer
 import com.esotericsoftware.kryo.util.MapReferenceResolver
 import com.google.common.annotations.VisibleForTesting
 import net.corda.core.contracts.*
 import net.corda.core.crypto.*
 import net.corda.core.node.AttachmentsClassLoader
-import net.corda.core.node.services.AttachmentStorage
 import net.corda.core.transactions.WireTransaction
 import net.i2p.crypto.eddsa.EdDSAPrivateKey
 import net.i2p.crypto.eddsa.EdDSAPublicKey
 import net.i2p.crypto.eddsa.spec.EdDSAPrivateKeySpec
 import net.i2p.crypto.eddsa.spec.EdDSAPublicKeySpec
 import org.bouncycastle.asn1.ASN1InputStream
-import org.bouncycastle.asn1.ASN1Sequence
 import org.bouncycastle.asn1.x500.X500Name
 import org.slf4j.Logger
 import org.slf4j.LoggerFactory
@@ -80,7 +77,7 @@ fun storageKryo(): KryoPool = internalKryoPool
  * A type safe wrapper around a byte array that contains a serialised object. You can call [SerializedBytes.deserialize]
  * to get the original object back.
  */
-@Suppress("unused")  // Type parameter is just for documentation purposes.
+@Suppress("unused") // Type parameter is just for documentation purposes.
 class SerializedBytes<T : Any>(bytes: ByteArray, val internalOnly: Boolean = false) : OpaqueBytes(bytes) {
     // It's OK to use lazy here because SerializedBytes is configured to use the ImmutableClassSerializer.
     val hash: SecureHash by lazy { bytes.sha256() }
@@ -308,6 +305,18 @@ object WireTransactionSerializer : Serializer<WireTransaction>() {
         kryo.writeClassAndObject(output, obj.timestamp)
     }
 
+    private fun attachmentsClassLoader(kryo: Kryo, attachmentHashes: List<SecureHash>): ClassLoader? {
+        val serializationContext = kryo.serializationContext() ?: return null // Some tests don't set one.
+        serializationContext.serviceHub.storageService.attachmentsClassLoaderEnabled || return null
+        val missing = ArrayList<SecureHash>()
+        val attachments = ArrayList<Attachment>()
+        attachmentHashes.forEach { id ->
+            serializationContext.serviceHub.storageService.attachments.openAttachment(id)?.let { attachments += it } ?: run { missing += id }
+        }
+        missing.isNotEmpty() && throw MissingAttachmentsException(missing)
+        return AttachmentsClassLoader(attachments)
+    }
+
     @Suppress("UNCHECKED_CAST")
     override fun read(kryo: Kryo, input: Input, type: Class<WireTransaction>): WireTransaction {
         val inputs = kryo.readClassAndObject(input) as List<StateRef>
@@ -315,30 +324,13 @@ object WireTransactionSerializer : Serializer<WireTransaction>() {
 
         // If we're deserialising in the sandbox context, we use our special attachments classloader.
         // Otherwise we just assume the code we need is on the classpath already.
-        val attachmentStorage = kryo.attachmentStorage
+        kryo.useClassLoader(attachmentsClassLoader(kryo, attachmentHashes) ?: javaClass.classLoader) {
-        val classLoader = if (attachmentStorage != null) {
-            val missing = ArrayList<SecureHash>()
-            val attachments = ArrayList<Attachment>()
-            for (id in attachmentHashes) {
-                val attachment = attachmentStorage.openAttachment(id)
-                if (attachment == null)
-                    missing += id
-                else
-                    attachments += attachment
-            }
-            if (missing.isNotEmpty())
-                throw MissingAttachmentsException(missing)
-            AttachmentsClassLoader(attachments)
-        } else javaClass.classLoader
-
-        kryo.useClassLoader(classLoader) {
             val outputs = kryo.readClassAndObject(input) as List<TransactionState<ContractState>>
             val commands = kryo.readClassAndObject(input) as List<Command>
             val notary = kryo.readClassAndObject(input) as Party?
             val signers = kryo.readClassAndObject(input) as List<PublicKey>
             val transactionType = kryo.readClassAndObject(input) as TransactionType
             val timestamp = kryo.readClassAndObject(input) as Timestamp?
-
             return WireTransaction(inputs, attachmentHashes, outputs, commands, notary, signers, transactionType, timestamp)
         }
     }
@@ -385,7 +377,7 @@ object CompositeKeySerializer : Serializer<CompositeKey>() {
         val threshold = input.readInt()
         val children = readListOfLength<CompositeKey.NodeAndWeight>(kryo, input, minLen = 2)
         val builder = CompositeKey.Builder()
-        children.forEach { builder.addKey(it.node, it.weight)  }
+        children.forEach { builder.addKey(it.node, it.weight) }
         return builder.build(threshold) as CompositeKey
     }
 }
@@ -394,7 +386,7 @@ object CompositeKeySerializer : Serializer<CompositeKey>() {
  * Helper function for reading lists with number of elements at the beginning.
  * @param minLen minimum number of elements we expect for list to include, defaults to 1
  * @param expectedLen expected length of the list, defaults to null if arbitrary length list read
-  */
+ */
 inline fun <reified T> readListOfLength(kryo: Kryo, input: Input, minLen: Int = 1, expectedLen: Int? = null): List<T> {
     val elemCount = input.readInt()
     if (elemCount < minLen) throw KryoException("Cannot deserialize list, too little elements. Minimum required: $minLen, got: $elemCount")
@@ -509,21 +501,6 @@ fun <T> Kryo.withoutReferences(block: () -> T): T {
     }
 }
 
-val ATTACHMENT_STORAGE = "ATTACHMENT_STORAGE"
-
-val Kryo.attachmentStorage: AttachmentStorage?
-    get() = this.context.get(ATTACHMENT_STORAGE, null) as AttachmentStorage?
-
-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)
-    }
-}
-
 /** For serialising a MetaData object. */
 @ThreadSafe
 object MetaDataSerializer : Serializer<MetaData>() {

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

@@ -41,29 +41,31 @@ interface SerializationToken {
  */
 class SerializeAsTokenSerializer<T : SerializeAsToken> : Serializer<T>() {
     override fun write(kryo: Kryo, output: Output, obj: T) {
-        kryo.writeClassAndObject(output, obj.toToken(getContext(kryo) ?: throw KryoException("Attempt to write a ${SerializeAsToken::class.simpleName} instance of ${obj.javaClass.name} without initialising a context")))
+        kryo.writeClassAndObject(output, obj.toToken(kryo.serializationContext() ?: throw KryoException("Attempt to write a ${SerializeAsToken::class.simpleName} instance of ${obj.javaClass.name} without initialising a context")))
     }
 
     override fun read(kryo: Kryo, input: Input, type: Class<T>): T {
         val token = (kryo.readClassAndObject(input) as? SerializationToken) ?: throw KryoException("Non-token read for tokenized type: ${type.name}")
-        val fromToken = token.fromToken(getContext(kryo) ?: throw KryoException("Attempt to read a token for a ${SerializeAsToken::class.simpleName} instance of ${type.name} without initialising a context"))
+        val fromToken = token.fromToken(kryo.serializationContext() ?: throw KryoException("Attempt to read a token for a ${SerializeAsToken::class.simpleName} instance of ${type.name} without initialising a context"))
         if (type.isAssignableFrom(fromToken.javaClass)) {
             return type.cast(fromToken)
         } else {
             throw KryoException("Token read ($token) did not return expected tokenized type: ${type.name}")
         }
     }
+}
 
-    companion object {
+private val serializationContextKey = SerializeAsTokenContext::class.java
-        private fun getContext(kryo: Kryo): SerializeAsTokenContext? = kryo.context.get(SerializeAsTokenContext::class.java) as? SerializeAsTokenContext
 
-        fun setContext(kryo: Kryo, context: SerializeAsTokenContext) {
+fun Kryo.serializationContext() = context.get(serializationContextKey) as? SerializeAsTokenContext
-            kryo.context.put(SerializeAsTokenContext::class.java, context)
-        }
 
-        fun clearContext(kryo: Kryo) {
+fun <T> Kryo.withSerializationContext(serializationContext: SerializeAsTokenContext, block: () -> T) = run {
-            kryo.context.remove(SerializeAsTokenContext::class.java)
+    context.containsKey(serializationContextKey) && throw IllegalStateException("There is already a serialization context.")
-        }
+    context.put(serializationContextKey, serializationContext)
+    try {
+        block()
+    } finally {
+        context.remove(serializationContextKey)
     }
 }
 
@@ -76,7 +78,15 @@ class SerializeAsTokenSerializer<T : SerializeAsToken> : Serializer<T>() {
  * Then it is a case of using the companion object methods on [SerializeAsTokenSerializer] to set and clear context as necessary
  * on the Kryo instance when serializing to enable/disable tokenization.
  */
-class SerializeAsTokenContext(toBeTokenized: Any, kryoPool: KryoPool, val serviceHub: ServiceHub) {
+class SerializeAsTokenContext internal constructor(val serviceHub: ServiceHub, init: SerializeAsTokenContext.() -> Unit) {
+    constructor(toBeTokenized: Any, kryoPool: KryoPool, serviceHub: ServiceHub) : this(serviceHub, {
+        kryoPool.run { kryo ->
+            kryo.withSerializationContext(this) {
+                toBeTokenized.serialize(kryo)
+            }
+        }
+    })
+
     private val classNameToSingleton = mutableMapOf<String, SerializeAsToken>()
     private var readOnly = false
 
@@ -90,11 +100,7 @@ class SerializeAsTokenContext(toBeTokenized: Any, kryoPool: KryoPool, val servic
          * accidental registrations from occuring as these could not be deserialized in a deserialization-first
          * scenario if they are not part of this iniital context construction serialization.
          */
-        kryoPool.run { kryo ->
+        init(this)
-            SerializeAsTokenSerializer.setContext(kryo, this)
-            toBeTokenized.serialize(kryo)
-            SerializeAsTokenSerializer.clearContext(kryo)
-        }
         readOnly = true
     }
 

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

@@ -1,10 +1,13 @@
 package net.corda.core.node
 
 import com.esotericsoftware.kryo.Kryo
+import com.nhaarman.mockito_kotlin.mock
+import com.nhaarman.mockito_kotlin.whenever
 import net.corda.core.contracts.*
 import net.corda.core.crypto.Party
 import net.corda.core.crypto.SecureHash
 import net.corda.core.node.services.AttachmentStorage
+import net.corda.core.node.services.StorageService
 import net.corda.core.serialization.*
 import net.corda.core.transactions.TransactionBuilder
 import net.corda.core.utilities.DUMMY_NOTARY
@@ -35,6 +38,15 @@ val ATTACHMENT_TEST_PROGRAM_ID = AttachmentClassLoaderTests.AttachmentDummyContr
 class AttachmentClassLoaderTests {
     companion object {
         val ISOLATED_CONTRACTS_JAR_PATH: URL = AttachmentClassLoaderTests::class.java.getResource("isolated.jar")
+
+        private fun <T> Kryo.withAttachmentStorage(attachmentStorage: AttachmentStorage, block: () -> T) = run {
+            val serviceHub = mock<ServiceHub>()
+            val storageService = mock<StorageService>()
+            whenever(serviceHub.storageService).thenReturn(storageService)
+            whenever(storageService.attachmentsClassLoaderEnabled).thenReturn(true)
+            whenever(storageService.attachments).thenReturn(attachmentStorage)
+            withSerializationContext(SerializeAsTokenContext(serviceHub) {}, block)
+        }
     }
 
     class AttachmentDummyContract : Contract {

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

@@ -22,7 +22,6 @@ class SerializationTokenTest {
 
     @After
     fun cleanup() {
-        SerializeAsTokenSerializer.clearContext(kryo)
         storageKryo().release(kryo)
     }
 
@@ -46,11 +45,12 @@ class SerializationTokenTest {
     fun `write token and read tokenizable`() {
         val tokenizableBefore = LargeTokenizable()
         val context = serializeAsTokenContext(tokenizableBefore)
-        SerializeAsTokenSerializer.setContext(kryo, context)
+        kryo.withSerializationContext(context) {
-        val serializedBytes = tokenizableBefore.serialize(kryo)
+            val serializedBytes = tokenizableBefore.serialize(kryo)
-        assertThat(serializedBytes.size).isLessThan(tokenizableBefore.numBytes)
+            assertThat(serializedBytes.size).isLessThan(tokenizableBefore.numBytes)
-        val tokenizableAfter = serializedBytes.deserialize(kryo)
+            val tokenizableAfter = serializedBytes.deserialize(kryo)
-        assertThat(tokenizableAfter).isSameAs(tokenizableBefore)
+            assertThat(tokenizableAfter).isSameAs(tokenizableBefore)
+        }
     }
 
     private class UnitSerializeAsToken : SingletonSerializeAsToken()
@@ -59,27 +59,30 @@ class SerializationTokenTest {
     fun `write and read singleton`() {
         val tokenizableBefore = UnitSerializeAsToken()
         val context = serializeAsTokenContext(tokenizableBefore)
-        SerializeAsTokenSerializer.setContext(kryo, context)
+        kryo.withSerializationContext(context) {
-        val serializedBytes = tokenizableBefore.serialize(kryo)
+            val serializedBytes = tokenizableBefore.serialize(kryo)
-        val tokenizableAfter = serializedBytes.deserialize(kryo)
+            val tokenizableAfter = serializedBytes.deserialize(kryo)
-        assertThat(tokenizableAfter).isSameAs(tokenizableBefore)
+            assertThat(tokenizableAfter).isSameAs(tokenizableBefore)
+        }
     }
 
     @Test(expected = UnsupportedOperationException::class)
     fun `new token encountered after context init`() {
         val tokenizableBefore = UnitSerializeAsToken()
         val context = serializeAsTokenContext(emptyList<Any>())
-        SerializeAsTokenSerializer.setContext(kryo, context)
+        kryo.withSerializationContext(context) {
-        tokenizableBefore.serialize(kryo)
+            tokenizableBefore.serialize(kryo)
+        }
     }
 
     @Test(expected = UnsupportedOperationException::class)
     fun `deserialize unregistered token`() {
         val tokenizableBefore = UnitSerializeAsToken()
         val context = serializeAsTokenContext(emptyList<Any>())
-        SerializeAsTokenSerializer.setContext(kryo, context)
+        kryo.withSerializationContext(context) {
-        val serializedBytes = tokenizableBefore.toToken(serializeAsTokenContext(emptyList<Any>())).serialize(kryo)
+            val serializedBytes = tokenizableBefore.toToken(serializeAsTokenContext(emptyList<Any>())).serialize(kryo)
-        serializedBytes.deserialize(kryo)
+            serializedBytes.deserialize(kryo)
+        }
     }
 
     @Test(expected = KryoException::class)
@@ -92,14 +95,15 @@ class SerializationTokenTest {
     fun `deserialize non-token`() {
         val tokenizableBefore = UnitSerializeAsToken()
         val context = serializeAsTokenContext(tokenizableBefore)
-        SerializeAsTokenSerializer.setContext(kryo, context)
+        kryo.withSerializationContext(context) {
-        val stream = ByteArrayOutputStream()
+            val stream = ByteArrayOutputStream()
-        Output(stream).use {
+            Output(stream).use {
-            kryo.writeClass(it, SingletonSerializeAsToken::class.java)
+                kryo.writeClass(it, SingletonSerializeAsToken::class.java)
-            kryo.writeObject(it, emptyList<Any>())
+                kryo.writeObject(it, emptyList<Any>())
+            }
+            val serializedBytes = SerializedBytes<Any>(stream.toByteArray())
+            serializedBytes.deserialize(kryo)
         }
-        val serializedBytes = SerializedBytes<Any>(stream.toByteArray())
-        serializedBytes.deserialize(kryo)
     }
 
     private class WrongTypeSerializeAsToken : SerializeAsToken {
@@ -114,8 +118,9 @@ class SerializationTokenTest {
     fun `token returns unexpected type`() {
         val tokenizableBefore = WrongTypeSerializeAsToken()
         val context = serializeAsTokenContext(tokenizableBefore)
-        SerializeAsTokenSerializer.setContext(kryo, context)
+        kryo.withSerializationContext(context) {
-        val serializedBytes = tokenizableBefore.serialize(kryo)
+            val serializedBytes = tokenizableBefore.serialize(kryo)
-        serializedBytes.deserialize(kryo)
+            serializedBytes.deserialize(kryo)
+        }
     }
 }

node/src/main/kotlin/net/corda/node/services/persistence/StorageServiceImpl.kt

@@ -7,6 +7,8 @@ open class StorageServiceImpl(override val attachments: AttachmentStorage,
                               override val validatedTransactions: TransactionStorage,
                               override val stateMachineRecordedTransactionMapping: StateMachineRecordedTransactionMappingStorage)
     : SingletonSerializeAsToken(), TxWritableStorageService {
+    override val attachmentsClassLoaderEnabled = false
+
     lateinit override var uploaders: List<FileUploader>
 
     fun initUploaders(uploadersList: List<FileUploader>) {

node/src/main/kotlin/net/corda/node/services/statemachine/StateMachineManager.kt

@@ -374,16 +374,18 @@ class StateMachineManager(val serviceHub: ServiceHubInternal,
     private fun serializeFiber(fiber: FlowStateMachineImpl<*>): SerializedBytes<FlowStateMachineImpl<*>> {
         return quasarKryo().run { kryo ->
             // add the map of tokens -> tokenizedServices to the kyro context
-            SerializeAsTokenSerializer.setContext(kryo, serializationContext)
+            kryo.withSerializationContext(serializationContext) {
-            fiber.serialize(kryo)
+                fiber.serialize(kryo)
+            }
         }
     }
 
     private fun deserializeFiber(checkpoint: Checkpoint): FlowStateMachineImpl<*> {
         return quasarKryo().run { kryo ->
             // put the map of token -> tokenized into the kryo context
-            SerializeAsTokenSerializer.setContext(kryo, serializationContext)
+            kryo.withSerializationContext(serializationContext) {
-            checkpoint.serializedFiber.deserialize(kryo).apply { fromCheckpoint = true }
+                checkpoint.serializedFiber.deserialize(kryo)
+            }.apply { fromCheckpoint = true }
         }
     }
 

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

@@ -173,7 +173,9 @@ class MockStorageService(override val attachments: AttachmentStorage = MockAttac
                          override val validatedTransactions: TransactionStorage = MockTransactionStorage(),
                          override val uploaders: List<FileUploader> = listOf<FileUploader>(),
                          override val stateMachineRecordedTransactionMapping: StateMachineRecordedTransactionMappingStorage = MockStateMachineRecordedTransactionMappingStorage())
-    : SingletonSerializeAsToken(), TxWritableStorageService
+    : SingletonSerializeAsToken(), TxWritableStorageService {
+    override val attachmentsClassLoaderEnabled = false
+}
 
 /**
  * Make properties appropriate for creating a DataSource for unit tests.