2 files changed, 29 insertions, 23 deletions

diff --git a/src/client/websocketpp_client.cpp b/src/client/websocketpp_client.cpp
index 5199343..981d4ec 100644
--- a/src/client/websocketpp_client.cpp
+++ b/src/client/websocketpp_client.cpp
@@ -1,10 +1,8 @@
-#include <nlohmann/json.hpp>
 #include <websocketpp/client.hpp>
 #include <websocketpp/config/asio_client.hpp>
 
 #include "web_socket_client.hpp"
 
-using nlohmann::json;
 using std::error_code;
 using std::function;
 using std::lock_guard;
@@ -86,26 +84,17 @@ public:
         return make_tuple(uri, true);
     };
 
-    void receive(string uri, function<void(const string&, const string&)> messageHandler) override
+    void receive(string uri, function<void(const string&)> messageHandler) override
     {
-        this->_client.set_message_handler([this, messageHandler](
+        lock_guard<mutex> lock(this->_propertyMutex);
+        auto connectionHandle = this->_connectionHandles[uri];
+        auto connection = this->_client.get_con_from_hdl(connectionHandle);
+
+        connection->set_message_handler([messageHandler](
             websocketpp::connection_hdl connectionHandle,
             websocketpp_client::message_ptr message)
         {
-            json jarr = json::array();
-            string payload = message->get_payload();
-
-            jarr.parse(payload);
-            string messageType = jarr[0];
-            
-            if (messageType == "EVENT")
-            {
-                string subscriptionId = jarr[1];
-                string messageContents = jarr[2].dump();
-                messageHandler(subscriptionId, messageContents);
-            };
-
-            // TODO: Add support for other message types.
+            messageHandler(message->get_payload());
         });
     };
 
diff --git a/src/nostr_service.cpp b/src/nostr_service.cpp
index 7efc11e..ac63f23 100644
--- a/src/nostr_service.cpp
+++ b/src/nostr_service.cpp
@@ -164,7 +164,7 @@ string NostrService::queryRelays(Filters filters)
     });
 };
 
-string NostrService::queryRelays(Filters filters, function<void(string, Event)> responseHandler)
+string NostrService::queryRelays(Filters filters, function<void(const string&, Event)> responseHandler)
 {
     RelayList successfulRelays;
     RelayList failedRelays;
@@ -182,10 +182,8 @@ string NostrService::queryRelays(Filters filters, function<void(string, Event)>
         });
         requestFutures.push_back(move(requestFuture));
 
-        this->_client->receive(relay, [responseHandler](string subscriptionId, string message) {
-            Event event;
-            event.deserialize(message);
-            responseHandler(subscriptionId, event);
+        this->_client->receive(relay, [this, responseHandler](string payload) {
+            this->onMessage(payload, responseHandler);
         });
     }
 
@@ -460,6 +458,25 @@ bool NostrService::hasSubscription(string relay, string subscriptionId)
     return false;
 };
 
+void NostrService::onMessage(string message, function<void(const string&, Event)> eventHandler)
+{
+    json jarr = json::array();
+    jarr = json::parse(message);
+
+    string messageType = jarr[0];
+
+    if (messageType == "EVENT")
+    {
+        string subscriptionId = jarr[1];
+        string serializedEvent = jarr[2].dump();
+        Event event;
+        event.deserialize(message);
+        eventHandler(subscriptionId, event);
+    }
+
+    // Support other message types here, if necessary.
+};
+
 void NostrService::onEvent(string subscriptionId, Event event)
 {
     lock_guard<mutex> lock(this->_propertyMutex);