1 files changed, 29 insertions, 21 deletions

diff --git a/src/Room.cpp b/src/Room.cpp
index 7a51760..924f2da 100644
--- a/src/Room.cpp
+++ b/src/Room.cpp
@@ -200,11 +200,14 @@ namespace dchat
         auto group = room->addGroupLocally(request.groupId);
         auto user = room->addUserLocally(*request.creatorPublicKey, group);
         bool isLocalUser = false;
-        auto localUserIt = roomLocalUser.find(*request.nodeHash);
-        if(localUserIt != roomLocalUser.end())
+        if(!room->localUser)
         {
-            room->setLocalUser(user, localUserIt->second);
-            isLocalUser = true;
+            auto localUserIt = roomLocalUser.find(*request.nodeHash);
+            if(localUserIt != roomLocalUser.end() && user->publicKey == localUserIt->second->getPublicKey())
+            {
+                room->setLocalUser(user, localUserIt->second);
+                isLocalUser = true;
+            }
         }
 
         if(callbackFuncs.addUserCallbackFunc)
@@ -355,23 +358,26 @@ namespace dchat
         }
 
         auto user = room->addUserLocally(*request.userToAddPublicKey, group);
-        std::lock_guard<std::recursive_mutex> waitingToJoinRoomLock(waitingToJoinRoomMutex);
         bool isLocalUser = false;
-        auto waitingToJoinRoomIt = waitingToJoinRoom.find(*request.nodeHash);
-        if(waitingToJoinRoomIt != waitingToJoinRoom.end())
-        {
-            room->setLocalUser(user, waitingToJoinRoomIt->second);
-            waitingToJoinRoom.erase(waitingToJoinRoomIt);
-            isLocalUser = true;
-        }
-        else
+        if(!room->localUser)
         {
-            auto localUserIt = roomLocalUser.find(*request.nodeHash);
-            if(localUserIt != roomLocalUser.end())
+            std::lock_guard<std::recursive_mutex> waitingToJoinRoomLock(waitingToJoinRoomMutex);
+            auto waitingToJoinRoomIt = waitingToJoinRoom.find(*request.nodeHash);
+            if(waitingToJoinRoomIt != waitingToJoinRoom.end() && user->publicKey == waitingToJoinRoomIt->second->getPublicKey())
             {
-                room->setLocalUser(user, localUserIt->second);
+                room->setLocalUser(user, waitingToJoinRoomIt->second);
+                waitingToJoinRoom.erase(waitingToJoinRoomIt);
                 isLocalUser = true;
             }
+            else
+            {
+                auto localUserIt = roomLocalUser.find(*request.nodeHash);
+                if(localUserIt != roomLocalUser.end() && user->publicKey == localUserIt->second->getPublicKey())
+                {
+                    room->setLocalUser(user, localUserIt->second);
+                    isLocalUser = true;
+                }
+            }
         }
 
         if(callbackFuncs.addUserCallbackFunc)
@@ -491,8 +497,8 @@ namespace dchat
         }
         std::string roomEncryptionKeyHex = inviteKey.substr(64, 64);
         std::string roomEncryptionKeyStr = odhtdb::hex2bin(roomEncryptionKeyHex.c_str(), roomEncryptionKeyHex.size());
-        auto roomEncryptionKey = std::make_shared<odhtdb::OwnedByteArray>(new u8[odhtdb::ENCRYPTION_KEY_BYTE_SIZE], odhtdb::ENCRYPTION_KEY_BYTE_SIZE);
-        memcpy(roomEncryptionKey->data, roomEncryptionKeyStr.data(), roomEncryptionKeyStr.size());
+        auto newEncryptionKey = std::make_shared<odhtdb::OwnedByteArray>(new u8[odhtdb::ENCRYPTION_KEY_BYTE_SIZE], odhtdb::ENCRYPTION_KEY_BYTE_SIZE);
+        memcpy(newEncryptionKey->data, roomEncryptionKeyStr.data(), roomEncryptionKeyStr.size());
         std::string inviteKeyChecksum = inviteKey.substr(128, 2);
         if(inviteKeyChecksum != getInviteKeyChecksum(inviteKey.substr(0, 128)))
         {
@@ -505,16 +511,18 @@ namespace dchat
         serializer.add(ourNewUser->getPublicKey().getData(), ourNewUser->getPublicKey().getSize());
         serializer.add(message.data(), message.size());
         odhtdb::DataView encryptionData { serializer.getBuffer().data(), serializer.getSize() };
-        odhtdb::DataView encryptionKey { (void*)roomEncryptionKey->data, roomEncryptionKey->size };
+        odhtdb::DataView encryptionKey { (void*)newEncryptionKey->data, newEncryptionKey->size };
         odhtdb::Encryption encryption(encryptionData, encryptionKey);
 
         sibs::SafeSerializer serializerFinish;
         serializerFinish.add((const u8*)encryption.getNonce().data, encryption.getNonce().size);
         serializerFinish.add((const u8*)encryption.getCipherText().data, encryption.getCipherText().size);
         odhtdb::InfoHash inviteInfoHash = odhtdb::InfoHash::generateHash((const u8*)inviteKey.data(), inviteKey.size());
-        database->sendCustomMessage(inviteInfoHash, serializerFinish.getBuffer().data(), serializerFinish.getSize());
+        roomLocalUser[*roomId] = ourNewUser;
+        roomEncryptionKey[*roomId] = newEncryptionKey;
         waitingToJoinRoom[*roomId] = ourNewUser;
-        odhtdb::DatabaseNode roomNode(roomEncryptionKey, roomId);
+        database->sendCustomMessage(inviteInfoHash, serializerFinish.getBuffer().data(), serializerFinish.getSize());
+        odhtdb::DatabaseNode roomNode(newEncryptionKey, roomId);
         database->seed(roomNode);
         database->storeNodeInfoForUserEncrypted(roomNode, currentUsername, currentUserPassword, *ourNewUser);
     }