Fix crash when posting message after creating room

1 files changed, 25 insertions, 4 deletions

diff --git a/src/Room.cpp b/src/Room.cpp
index 30375a3..95e5818 100644
--- a/src/Room.cpp
+++ b/src/Room.cpp
@@ -9,10 +9,9 @@
 // TODO: Remove error checks when odhtdb has been improved to take care of such errors
 namespace dchat
 {
-    Room::Room(Rooms *_rooms, std::shared_ptr<odhtdb::Hash> _id, std::shared_ptr<odhtdb::OwnedByteArray> _encryptionKey) : 
+    Room::Room(Rooms *_rooms, std::shared_ptr<odhtdb::Hash> _id) : 
         rooms(_rooms),
         id(_id),
-        encryptionKey(_encryptionKey),
         userdata(nullptr)
     {
 
@@ -43,6 +42,9 @@ namespace dchat
 
     void Room::publishMessage(const std::string &msg)
     {
+        assert(localUser);
+        assert(publicKeyToKeyPairMap.find(localUser->publicKey) != publicKeyToKeyPairMap.end());
+        assert(encryptionKey);
         sibs::SafeSerializer serializer;
         serializer.add(RoomDataType::ADD_MESSAGE);
         serializer.add((const u8*)msg.data(), msg.size());
@@ -64,6 +66,7 @@ namespace dchat
 
     void Rooms::createNodeCallbackFunc(const odhtdb::DatabaseCreateNodeRequest &request)
     {
+        std::lock_guard<std::recursive_mutex> lock(roomModifyMutex);
         auto roomIt = roomById.find(*request.nodeHash);
         if(roomIt != roomById.end())
         {
@@ -72,9 +75,11 @@ namespace dchat
             return;
         }
 
-        auto encryptionKey = roomEncryptionKey[*request.nodeHash];
         auto roomId = std::make_shared<odhtdb::Hash>(*request.nodeHash);
-        auto room = std::make_shared<Room>(this, roomId, encryptionKey);
+        auto room = std::make_shared<Room>(this, roomId);
+        auto encryptionKeyIt = roomEncryptionKey.find(*request.nodeHash);
+        if(encryptionKeyIt != roomEncryptionKey.end())
+            room->encryptionKey = encryptionKeyIt->second;
         roomById[*request.nodeHash] = room;
         if(callbackFuncs.createRoomCallbackFunc)
             callbackFuncs.createRoomCallbackFunc(room);
@@ -92,6 +97,7 @@ namespace dchat
 
     void Rooms::addNodeCallbackFunc(const odhtdb::DatabaseAddNodeRequest &request)
     {
+        std::lock_guard<std::recursive_mutex> lock(roomModifyMutex);
         auto roomIt = roomById.find(*request.nodeHash);
         if(roomIt == roomById.end())
         {
@@ -197,6 +203,7 @@ namespace dchat
 
     void Rooms::addUserCallbackFunc(const odhtdb::DatabaseAddUserRequest &request)
     {
+        std::lock_guard<std::recursive_mutex> lock(roomModifyMutex);
         auto roomIt = roomById.find(*request.nodeHash);
         if(roomIt == roomById.end())
         {
@@ -234,6 +241,7 @@ namespace dchat
 
     void Rooms::loginUser(const std::string &username, const std::string &password)
     {
+        std::lock_guard<std::recursive_mutex> lock(roomModifyMutex);
         if(loggedIn)
             throw std::runtime_error(std::string("You are already logged in as ") + username);
 
@@ -255,6 +263,7 @@ namespace dchat
 
     void Rooms::registerUser(const std::string &username, const std::string &password)
     {
+        std::lock_guard<std::recursive_mutex> lock(roomModifyMutex);
         if(loggedIn)
             throw std::runtime_error(std::string("You are already logged in as ") + username);
 
@@ -266,12 +275,24 @@ namespace dchat
 
     void Rooms::createRoom(const std::string &name)
     {
+        std::lock_guard<std::recursive_mutex> lock(roomModifyMutex);
         if(!loggedIn)
             throw std::runtime_error("You need to be logged in to create a room ");
 
         auto newNode = database->create();
         roomLocalUser[*newNode->getRequestHash()] = newNode->getNodeAdminKeyPair();
         roomEncryptionKey[*newNode->getRequestHash()] = newNode->getNodeEncryptionKey();
+        auto roomIt = roomById.find(*newNode->getRequestHash());
+        if(roomIt != roomById.end())
+        {
+            roomIt->second->encryptionKey = newNode->getNodeEncryptionKey();
+            auto user = roomIt->second->getUserByPublicKey(newNode->getNodeAdminKeyPair()->getPublicKey());
+            if(user)
+            {
+                roomIt->second->localUser = user;
+                roomIt->second->publicKeyToKeyPairMap[user->publicKey] = newNode->getNodeAdminKeyPair();
+            }
+        }
 
         odhtdb::DatabaseNode nodeInfo(newNode->getNodeEncryptionKey(), newNode->getRequestHash());
         database->storeNodeInfoForUserEncrypted(nodeInfo, currentUsername, currentUserPassword, *newNode->getNodeAdminKeyPair());