diff --git a/Sources/tss-client-swift/Helpers.swift b/Sources/tss-client-swift/Helpers.swift
index 5f08b17..85b07de 100644
--- a/Sources/tss-client-swift/Helpers.swift
+++ b/Sources/tss-client-swift/Helpers.swift
@@ -118,23 +118,27 @@ public class TSSHelpers {
     ///
     /// - Throws: `TSSClientError`
     public static func hexUncompressedPublicKey(pubKey: Data, return64Bytes: Bool) throws -> String {
-        if pubKey.bytes.count == 65 && return64Bytes {
-            if pubKey.bytes.first == 04 {
-                return Data(pubKey.bytes.dropFirst()).hexString
+        if pubKey.bytes.count == 65 {
+            if return64Bytes {
+                if pubKey.bytes.first == 04 {
+                    return Data(pubKey.bytes.dropFirst()).hexString
+                } else {
+                    throw TSSClientError("Invalid public key bytes")
+                }
             } else {
-                throw TSSClientError("Invalid public key bytes")
+                return Data(pubKey.bytes).hexString
             }
-        } else if !return64Bytes {
-            return Data(pubKey.bytes).hexString
         }
 
-        if pubKey.bytes.count == 65 && !return64Bytes {
-            return Data(pubKey.bytes).hexString
-        } else if return64Bytes { // first byte should be 04 prefix
-            let prefix: UInt8 = 4
-            var pk = Data(pubKey)
-            pk.insert(prefix, at: 0)
-            return pk.hexString
+        if pubKey.bytes.count == 64 {
+            if return64Bytes {
+                return Data(pubKey.bytes).hexString
+            } else { // first byte should be 04 prefix
+                let prefix: UInt8 = 4
+                var pk = Data(pubKey)
+                pk.insert(prefix, at: 0)
+                return pk.hexString
+            }
         }
 
         throw TSSClientError("Invalid public key bytes")
diff --git a/Sources/tss-client-swift/TSSSocket.swift b/Sources/tss-client-swift/TSSSocket.swift
index af853ad..345b103 100644
--- a/Sources/tss-client-swift/TSSSocket.swift
+++ b/Sources/tss-client-swift/TSSSocket.swift
@@ -53,13 +53,12 @@ internal final class TSSSocket {
                     print("disconnected, party:" + String(party))
                 })
                 socket.on("precompute_complete", callback: { data, ack in
-                    if session != self.session {
-                        print("ignoring message for a different session...")
-                        return
-                    }
-                    
                     if let json = try? JSONSerialization.data(withJSONObject: data[0]) {
                         if let msg = try? JSONDecoder().decode(TssPrecomputeUpdate.self, from: json) {
+                            if msg.session != self.session {
+                                print("ignoring message for a different session...")
+                                return
+                            }
                             EventQueue.shared.addEvent(event: Event(message: String(msg.party), session: msg.session, party: Int32(msg.party), occurred: Date(), type: EventType.PrecomputeComplete))
                         } else {
                             EventQueue.shared.addEvent(event: Event(message: "Received json was not decodable", session: self.session, party: Int32(-1), occurred: Date(), type: EventType.SocketDataError))
@@ -72,14 +71,13 @@ internal final class TSSSocket {
                     }
                 })
                 socket.on("precompute_failed", callback: { data, ack in
-                    if session != self.session {
-                        print("ignoring message for a different session...")
-                        return
-                    }
-                    
                     if let json = try? JSONSerialization.data(withJSONObject: data[0])
                     {
                         if let msg = try? JSONDecoder().decode(TssPrecomputeUpdate.self, from: json) {
+                            if msg.session != self.session {
+                                print("ignoring message for a different session...")
+                                return
+                            }
                             EventQueue.shared.addEvent(event: Event(message: String(msg.party), session: msg.session, party: Int32(msg.party), occurred: Date(), type: EventType.PrecomputeError))
                         } else {
                             EventQueue.shared.addEvent(event: Event(message: "Received json was not decodable", session: self.session, party: Int32(-1), occurred: Date(), type: EventType.SocketDataError))
@@ -92,13 +90,12 @@ internal final class TSSSocket {
                     }
                 })
                 socket.on("send", callback: { data, ack in
-                    if session != self.session {
-                        print("ignoring message for a different session...")
-                        return
-                    }
-                    
                     if let json = try? JSONSerialization.data(withJSONObject: data[0]) {
                         if let msg = try? JSONDecoder().decode(TssRecvMsg.self, from: json) {
+                            if msg.session != self.session {
+                                print("ignoring message for a different session...")
+                                return
+                            }
                             MessageQueue.shared.addMessage(msg: Message(session: msg.session, sender: UInt64(Int64(msg.sender)), recipient: UInt64(Int64(msg.recipient)), msgType: msg.msg_type, msgData: msg.msg_data))
                             let tag = msg.msg_type.split(separator: "~")[1]
                             print("dkls: Received message \(tag), sender: `\(msg.sender)`, receiver: `\(msg.recipient)`")