Split out payload field JSON construction from encryption/decryption.… (

#10)

* Split out payload field JSON construction from encryption/decryption. Also introduce MessageLogger interface that logs messages with unencrypted payloads.

lib/src/main/kotlin/org/walletconnect/Session.kt

@@ -73,6 +73,10 @@ interface Session {
         }
     }
 
+    interface MessageLogger {
+        fun log(message: Session.Transport.Message, isOwnMessage: Boolean)
+    }
+
     interface Callback {
         fun onStatus(status: Status)
         fun onMethodCall(call: MethodCall)
@@ -90,9 +94,21 @@ interface Session {
     data class TransportError(override val cause: Throwable) : RuntimeException("Transport exception caused by $cause", cause)
 
     interface PayloadAdapter {
-        fun parse(payload: String, key: String): MethodCall
-        fun prepare(data: MethodCall, key: String): String
+        /**
+         * Takes in the decrypted payload JSON and returns the parsed [MethodCall].
+         */
+        fun parse(decryptedPayload: String): MethodCall
+
+        /**
+         * Takes in a [MethodCall] and returns the unencrypted payload JSON.
+         */
+        fun prepare(data: MethodCall): String
+
+    }
 
+    interface PayloadEncryption {
+        fun encrypt(unencryptedPayloadJson: String, key: String): String
+        fun decrypt(encryptedPayloadJson: String, key: String): String
     }
 
     interface Transport {
@@ -121,7 +137,7 @@ interface Session {
             fun build(
                     url: String,
                     statusHandler: (Status) -> Unit,
-                    messageHandler: (Message) -> Unit
+                    messageHandler: (Message) -> Unit,
             ): Transport
         }
 

lib/src/main/kotlin/org/walletconnect/impls/MoshiPayloadAdapter.kt

@@ -4,23 +4,11 @@ import com.squareup.moshi.Json
 import com.squareup.moshi.JsonClass
 import com.squareup.moshi.Moshi
 import com.squareup.moshi.Types
-import org.bouncycastle.crypto.digests.SHA256Digest
-import org.bouncycastle.crypto.engines.AESEngine
-import org.bouncycastle.crypto.macs.HMac
-import org.bouncycastle.crypto.modes.CBCBlockCipher
-import org.bouncycastle.crypto.paddings.PKCS7Padding
-import org.bouncycastle.crypto.paddings.PaddedBufferedBlockCipher
-import org.bouncycastle.crypto.params.KeyParameter
-import org.bouncycastle.crypto.params.ParametersWithIV
-import org.komputing.khex.decode
-import org.komputing.khex.extensions.toNoPrefixHexString
 import org.walletconnect.Session
 import org.walletconnect.types.*
-import java.security.SecureRandom
 
 class MoshiPayloadAdapter(moshi: Moshi) : Session.PayloadAdapter {
 
-    private val payloadAdapter = moshi.adapter(EncryptedPayload::class.java)
     private val mapAdapter = moshi.adapter<Map<String, Any?>>(
         Types.newParameterizedType(
             Map::class.java,
@@ -29,74 +17,19 @@ class MoshiPayloadAdapter(moshi: Moshi) : Session.PayloadAdapter {
         )
     )
 
-    private fun createRandomBytes(i: Int) = ByteArray(i).also { SecureRandom().nextBytes(it) }
-
-    override fun parse(payload: String, key: String): Session.MethodCall {
-        val encryptedPayload = payloadAdapter.fromJson(payload) ?: throw IllegalArgumentException("Invalid json payload!")
-
-        // TODO verify hmac
-
-        val padding = PKCS7Padding()
-        val aes = PaddedBufferedBlockCipher(
-            CBCBlockCipher(AESEngine()),
-            padding
-        )
-        val ivAndKey = ParametersWithIV(
-            KeyParameter(decode(key)),
-            decode(encryptedPayload.iv)
-        )
-        aes.init(false, ivAndKey)
-
-        val encryptedData = decode(encryptedPayload.data)
-        val minSize = aes.getOutputSize(encryptedData.size)
-        val outBuf = ByteArray(minSize)
-        var len = aes.processBytes(encryptedData, 0, encryptedData.size, outBuf, 0)
-        len += aes.doFinal(outBuf, len)
-
-        return outBuf.copyOf(len).toMethodCall()
+    override fun parse(decryptedPayloadJson: String): Session.MethodCall {
+        return decryptedPayloadJson.toMethodCall()
     }
 
-    override fun prepare(data: Session.MethodCall, key: String): String {
-        val bytesData = data.toBytes()
-        val hexKey = decode(key)
-        val iv = createRandomBytes(16)
-
-        val padding = PKCS7Padding()
-        val aes = PaddedBufferedBlockCipher(
-            CBCBlockCipher(AESEngine()),
-            padding
-        )
-        aes.init(true, ParametersWithIV(KeyParameter(hexKey), iv))
-
-        val minSize = aes.getOutputSize(bytesData.size)
-        val outBuf = ByteArray(minSize)
-        val length1 = aes.processBytes(bytesData, 0, bytesData.size, outBuf, 0)
-        aes.doFinal(outBuf, length1)
-
-
-        val hmac = HMac(SHA256Digest())
-        hmac.init(KeyParameter(hexKey))
-
-        val hmacResult = ByteArray(hmac.macSize)
-        hmac.update(outBuf, 0, outBuf.size)
-        hmac.update(iv, 0, iv.size)
-        hmac.doFinal(hmacResult, 0)
-
-        return payloadAdapter.toJson(
-            EncryptedPayload(
-                outBuf.toNoPrefixHexString(),
-                hmac = hmacResult.toNoPrefixHexString(),
-                iv = iv.toNoPrefixHexString()
-            )
-        )
+    override fun prepare(data: Session.MethodCall): String {
+        return data.toJson()
     }
 
     /**
      * Convert FROM request bytes
      */
-    private fun ByteArray.toMethodCall(): Session.MethodCall =
-        String(this).let { json ->
-            mapAdapter.fromJson(json)?.let {
+    private fun String.toMethodCall(): Session.MethodCall =
+            mapAdapter.fromJson(this)?.let {
                 try {
                     val method = it["method"]
                     when (method) {
@@ -108,10 +41,9 @@ class MoshiPayloadAdapter(moshi: Moshi) : Session.PayloadAdapter {
                         else -> it.toCustom()
                     }
                 } catch (e: Exception) {
-                    throw Session.MethodCallException.InvalidRequest(it.getId(), "$json (${e.message ?: "Unknown error"})")
+                    throw Session.MethodCallException.InvalidRequest(it.getId(), "$this (${e.message ?: "Unknown error"})")
                 }
             } ?: throw IllegalArgumentException("Invalid json")
-        }
 
     private fun Map<String, *>.toSessionUpdate(): Session.MethodCall.SessionUpdate {
         val params = this["params"] as? List<*> ?: throw IllegalArgumentException("params missing")
@@ -163,7 +95,7 @@ class MoshiPayloadAdapter(moshi: Moshi) : Session.PayloadAdapter {
     /**
      * Convert INTO request bytes
      */
-    private fun Session.MethodCall.toBytes() =
+    private fun Session.MethodCall.toJson() =
         mapAdapter.toJson(
             when (this) {
                 is Session.MethodCall.SessionRequest -> this.toMap()
@@ -173,7 +105,7 @@ class MoshiPayloadAdapter(moshi: Moshi) : Session.PayloadAdapter {
                 is Session.MethodCall.SignMessage -> this.toMap()
                 is Session.MethodCall.Custom -> this.toMap()
             }
-        ).toByteArray()
+        )
 
     private fun Session.MethodCall.SessionRequest.toMap() =
         jsonRpc(id, "wc_sessionRequest", peer.intoMap())

lib/src/main/kotlin/org/walletconnect/impls/MoshiPayloadEncryption.kt

@@ -0,0 +1,82 @@
+package org.walletconnect.impls
+
+import com.squareup.moshi.Moshi
+import org.bouncycastle.crypto.digests.SHA256Digest
+import org.bouncycastle.crypto.engines.AESEngine
+import org.bouncycastle.crypto.macs.HMac
+import org.bouncycastle.crypto.modes.CBCBlockCipher
+import org.bouncycastle.crypto.paddings.PKCS7Padding
+import org.bouncycastle.crypto.paddings.PaddedBufferedBlockCipher
+import org.bouncycastle.crypto.params.KeyParameter
+import org.bouncycastle.crypto.params.ParametersWithIV
+import org.komputing.khex.decode
+import org.komputing.khex.extensions.toNoPrefixHexString
+import org.walletconnect.Session
+import java.security.SecureRandom
+
+class MoshiPayloadEncryption(moshi: Moshi) : Session.PayloadEncryption {
+
+    private val encryptedPayloadAdapter = moshi.adapter(MoshiPayloadAdapter.EncryptedPayload::class.java)
+
+    override fun encrypt(unencryptedPayloadJson: String, key: String): String {
+        val bytesData = unencryptedPayloadJson.toByteArray()
+        val hexKey = decode(key)
+        val iv = createRandomBytes(16)
+
+        val padding = PKCS7Padding()
+        val aes = PaddedBufferedBlockCipher(
+            CBCBlockCipher(AESEngine()),
+            padding
+        )
+        aes.init(true, ParametersWithIV(KeyParameter(hexKey), iv))
+
+        val minSize = aes.getOutputSize(bytesData.size)
+        val outBuf = ByteArray(minSize)
+        val length1 = aes.processBytes(bytesData, 0, bytesData.size, outBuf, 0)
+        aes.doFinal(outBuf, length1)
+
+
+        val hmac = HMac(SHA256Digest())
+        hmac.init(KeyParameter(hexKey))
+
+        val hmacResult = ByteArray(hmac.macSize)
+        hmac.update(outBuf, 0, outBuf.size)
+        hmac.update(iv, 0, iv.size)
+        hmac.doFinal(hmacResult, 0)
+
+        return encryptedPayloadAdapter.toJson(
+            MoshiPayloadAdapter.EncryptedPayload(
+                outBuf.toNoPrefixHexString(),
+                hmac = hmacResult.toNoPrefixHexString(),
+                iv = iv.toNoPrefixHexString()
+            )
+        )
+    }
+
+    override fun decrypt(encryptedPayloadJson: String, key: String): String {
+        val encryptedPayload = encryptedPayloadAdapter.fromJson(encryptedPayloadJson) ?: throw IllegalArgumentException("Invalid json payload!")
+
+        // TODO verify hmac
+
+        val padding = PKCS7Padding()
+        val aes = PaddedBufferedBlockCipher(
+            CBCBlockCipher(AESEngine()),
+            padding
+        )
+        val ivAndKey = ParametersWithIV(
+            KeyParameter(decode(key)),
+            decode(encryptedPayload.iv)
+        )
+        aes.init(false, ivAndKey)
+
+        val encryptedData = decode(encryptedPayload.data)
+        val minSize = aes.getOutputSize(encryptedData.size)
+        val outBuf = ByteArray(minSize)
+        var len = aes.processBytes(encryptedData, 0, encryptedData.size, outBuf, 0)
+        len += aes.doFinal(outBuf, len)
+
+        return String(outBuf.copyOf(len))
+    }
+
+    private fun createRandomBytes(i: Int) = ByteArray(i).also { SecureRandom().nextBytes(it) }
+}

lib/src/main/kotlin/org/walletconnect/impls/OkHttpTransport.kt

@@ -121,12 +121,12 @@ class OkHttpTransport(
         }
     }
 
-    class Builder(val client: OkHttpClient, val moshi: Moshi) :
+    class Builder(private val client: OkHttpClient, private val moshi: Moshi) :
         Session.Transport.Builder {
         override fun build(
             url: String,
             statusHandler: (Session.Transport.Status) -> Unit,
-            messageHandler: (Session.Transport.Message) -> Unit
+            messageHandler: (Session.Transport.Message) -> Unit,
         ): Session.Transport =
             OkHttpTransport(client, url, statusHandler, messageHandler, moshi)
 

lib/src/main/kotlin/org/walletconnect/impls/WCSession.kt

@@ -9,12 +9,14 @@ import java.util.*
 import java.util.concurrent.ConcurrentHashMap
 
 class WCSession(
-        private val config: Session.FullyQualifiedConfig,
-        private val payloadAdapter: Session.PayloadAdapter,
-        private val sessionStore: WCSessionStore,
-        transportBuilder: Session.Transport.Builder,
-        clientMeta: Session.PeerMeta,
-        clientId: String? = null
+    private val config: Session.FullyQualifiedConfig,
+    private val payloadAdapter: Session.PayloadAdapter,
+    private val payloadEncryption: Session.PayloadEncryption,
+    private val sessionStore: WCSessionStore,
+    private val messageLogger: Session.MessageLogger,
+    transportBuilder: Session.Transport.Builder,
+    clientMeta: Session.PeerMeta,
+    clientId: String? = null
 ) : Session {
 
     private val keyLock = Any()
@@ -91,11 +93,11 @@ class WCSession(
     override fun init() {
         if (transport.connect()) {
             // Register for all messages for this client
-            transport.send(
-                    Session.Transport.Message(
-                            config.handshakeTopic, "sub", ""
-                    )
+            val message = Session.Transport.Message(
+                config.handshakeTopic, "sub", ""
             )
+            transport.send(message)
+            messageLogger.log(message, isOwnMessage = true)
         }
     }
 
@@ -167,11 +169,11 @@ class WCSession(
         when (status) {
             Session.Transport.Status.Connected -> {
                 // Register for all messages for this client
-                transport.send(
-                    Session.Transport.Message(
-                        clientData.id, "sub", ""
-                    )
+                val message = Session.Transport.Message(
+                    clientData.id, "sub", ""
                 )
+                transport.send(message)
+                messageLogger.log(message, isOwnMessage = true)
             }
             Session.Transport.Status.Disconnected -> {
                 // no-op
@@ -194,7 +196,9 @@ class WCSession(
         val data: Session.MethodCall
         synchronized(keyLock) {
             try {
-                data = payloadAdapter.parse(message.payload, decryptionKey)
+                val decryptedPayload = payloadEncryption.decrypt(message.payload, decryptionKey)
+                data = payloadAdapter.parse(decryptedPayload)
+                messageLogger.log(message.copy(payload = decryptedPayload), isOwnMessage = false)
             } catch (e: Exception) {
                 handlePayloadError(e)
                 return
@@ -285,13 +289,17 @@ class WCSession(
         topic ?: return false
 
         val payload: String
+        val unencryptedPayload: String
         synchronized(keyLock) {
-            payload = payloadAdapter.prepare(msg, encryptionKey)
+            unencryptedPayload = payloadAdapter.prepare(msg)
+            payload = payloadEncryption.encrypt(unencryptedPayload, encryptionKey)
         }
         callback?.let {
             requests[msg.id()] = callback
         }
-        transport.send(Session.Transport.Message(topic, "pub", payload))
+        val message = Session.Transport.Message(topic, "pub", payload)
+        transport.send(message)
+        messageLogger.log(message.copy(payload = unencryptedPayload), isOwnMessage = true)
         return true
     }