From 9561013b65a638da2a9437d2c66650056ecaf62a Mon Sep 17 00:00:00 2001
From: sofusmortensen <sofusmortensen@gmail.com>
Date: Tue, 5 Apr 2016 01:02:54 +0200
Subject: [PATCH] Replaced class Kryo2 with extension method on Kryo

---
 .../main/kotlin/core/serialization/Kryo.kt    | 19 +++++++++----------
 src/test/kotlin/core/node/ClassLoaderTests.kt |  9 +++++----
 2 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/core/src/main/kotlin/core/serialization/Kryo.kt b/core/src/main/kotlin/core/serialization/Kryo.kt
index a2eed8c1f4..296fcbf928 100644
--- a/core/src/main/kotlin/core/serialization/Kryo.kt
+++ b/core/src/main/kotlin/core/serialization/Kryo.kt
@@ -190,7 +190,7 @@ inline fun <T> Kryo.useClassLoader(cl: ClassLoader, body: () -> T) : T {
     }
 }
 
-fun createKryo(k: Kryo = core.serialization.Kryo2()): Kryo {
+fun createKryo(k: Kryo = Kryo()): Kryo {
     return k.apply {
         // Allow any class to be deserialized (this is insecure but for prototyping we don't care)
         isRegistrationRequired = false
@@ -228,7 +228,7 @@ fun createKryo(k: Kryo = core.serialization.Kryo2()): Kryo {
                 var inputs = kryo.readClassAndObject( input ) as List<StateRef>
                 var attachments = kryo.readClassAndObject( input ) as List<SecureHash>
 
-                val attachmentStorage = (kryo as? core.serialization.Kryo2)?.attachmentStorage
+                val attachmentStorage = kryo.attachmentStorage
 
                 // .filterNotNull in order for TwoPartyTradeProtocolTests.checkDependenciesOfSaleAssetAreResolved test to run
                 val classLoader = core.node.AttachmentsClassLoader.create( attachments.map { attachmentStorage?.openAttachment(it) }.filterNotNull() )
@@ -264,11 +264,10 @@ fun createKryo(k: Kryo = core.serialization.Kryo2()): Kryo {
     }
 }
 
-/**
- * Extends Kryo with a field for passing attachmentStorage to serializer for WireTransaction
- *
- * TODO: Think of better solution, or at least better name
- */
-class Kryo2() : Kryo() {
-    var attachmentStorage: AttachmentStorage? = null
-}
+val ATTACHMENT_STORAGE = "ATTACHMENT_STORAGE"
+
+var Kryo.attachmentStorage: AttachmentStorage?
+    get() = this.context.get(ATTACHMENT_STORAGE, null) as AttachmentStorage?
+    set(value) {
+        this.context.put(ATTACHMENT_STORAGE, value)
+    }
diff --git a/src/test/kotlin/core/node/ClassLoaderTests.kt b/src/test/kotlin/core/node/ClassLoaderTests.kt
index c66c39fe5b..676ff88540 100644
--- a/src/test/kotlin/core/node/ClassLoaderTests.kt
+++ b/src/test/kotlin/core/node/ClassLoaderTests.kt
@@ -5,6 +5,7 @@ import contracts.DUMMY_PROGRAM_ID
 import contracts.DummyContract
 import core.*
 import core.crypto.SecureHash
+import core.serialization.attachmentStorage
 import core.serialization.createKryo
 import core.serialization.deserialize
 import core.serialization.serialize
@@ -215,7 +216,7 @@ class ClassLoaderTests {
 
         var storage = MockAttachmentStorage()
 
-        var kryo = createKryo() as core.serialization.Kryo2
+        var kryo = createKryo()
 
         // todo - think about better way to push attachmentStorage down to serializer
         kryo.attachmentStorage = storage
@@ -228,7 +229,7 @@ class ClassLoaderTests {
 
         val bytes = wireTransaction.serialize(kryo)
 
-        kryo = createKryo() as core.serialization.Kryo2
+        kryo = createKryo()
 
         // use empty attachmentStorage
         kryo.attachmentStorage = storage
@@ -252,7 +253,7 @@ class ClassLoaderTests {
 
         var storage = MockAttachmentStorage()
 
-        var kryo = createKryo() as core.serialization.Kryo2
+        var kryo = createKryo()
 
         // todo - think about better way to push attachmentStorage down to serializer
         kryo.attachmentStorage = storage
@@ -265,7 +266,7 @@ class ClassLoaderTests {
 
         val bytes = wireTransaction.serialize(kryo)
 
-        kryo = createKryo() as core.serialization.Kryo2
+        kryo = createKryo()
         // use empty attachmentStorage
         kryo.attachmentStorage = MockAttachmentStorage()