Fix crash when posting message after creating room

3 files changed, 29 insertions, 5 deletions

diff --git a/include/dchat/Gif.hpp b/include/dchat/Gif.hpp
index 9e39473..c870d1c 100644
--- a/include/dchat/Gif.hpp
+++ b/include/dchat/Gif.hpp
@@ -19,6 +19,7 @@ namespace dchat
         GifLoadException(const std::string &errMsg) : std::runtime_error(errMsg) {}
     };
     
+    // Gif is in image format ARGB (32-bits in total)
     class Gif
     {
     public:
diff --git a/include/dchat/Room.hpp b/include/dchat/Room.hpp
index 4938aea..5dd0e3f 100644
--- a/include/dchat/Room.hpp
+++ b/include/dchat/Room.hpp
@@ -8,6 +8,7 @@
 #include <memory>
 #include <vector>
 #include <functional>
+#include <mutex>
 #include <odhtdb/Database.hpp>
 
 namespace dchat
@@ -26,7 +27,7 @@ namespace dchat
     {
         DISABLE_COPY(Room)
     public:
-        Room(Rooms *rooms, std::shared_ptr<odhtdb::Hash> id, std::shared_ptr<odhtdb::OwnedByteArray> encryptionKey);
+        Room(Rooms *rooms, std::shared_ptr<odhtdb::Hash> id);
         std::shared_ptr<User> addUser(const odhtdb::Signature::PublicKey &userPublicKey, const odhtdb::DataView groupId);
         // Returns null if user doesn't exist in room
         std::shared_ptr<User> getUserByPublicKey(const odhtdb::Signature::PublicKey &userPublicKey);
@@ -114,5 +115,6 @@ namespace dchat
 
         std::string currentUsername;
         std::string currentUserPassword;
+        std::recursive_mutex roomModifyMutex;
     };
 }
\ No newline at end of file
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());