Matrix: attempt to fix threading issues

1 files changed, 195 insertions, 126 deletions

diff --git a/src/plugins/Matrix.cpp b/src/plugins/Matrix.cpp
index ec7064b..8ef7f47 100644
--- a/src/plugins/Matrix.cpp
+++ b/src/plugins/Matrix.cpp
@@ -16,7 +16,79 @@
 // When reaching top/bottom message, show older/newer messages.
 // Remove older messages (outside screen) to save memory. Reload them when the selected body item is the top/bottom one.
 
+// TODO: Verify if this class really is thread-safe (for example room data fields, user fields, message fields; etc that are updated in /sync)
+
 namespace QuickMedia {
+    std::shared_ptr<UserInfo> RoomData::get_user_by_id(const std::string &user_id) {
+        std::lock_guard<std::mutex> lock(room_mutex);
+        auto user_it = user_info_by_user_id.find(user_id);
+        if(user_it == user_info_by_user_id.end())
+            return nullptr;
+        return user_it->second;
+    }
+
+    void RoomData::add_user(std::shared_ptr<UserInfo> user) {
+        std::lock_guard<std::mutex> lock(room_mutex);
+        user_info_by_user_id.insert(std::make_pair(user->user_id, user));
+    }
+
+    void RoomData::set_user_read_marker(std::shared_ptr<UserInfo> &user, const std::string &event_id) {
+        std::lock_guard<std::mutex> lock(user_mutex);
+        user->read_marker_event_id = event_id;
+    }
+
+    std::string RoomData::get_user_read_marker(std::shared_ptr<UserInfo> &user) {
+        std::lock_guard<std::mutex> lock(user_mutex);
+        return user->read_marker_event_id;
+    }
+
+    void RoomData::prepend_messages_reverse(std::vector<std::shared_ptr<Message>> new_messages) {
+        std::lock_guard<std::mutex> lock(room_mutex);
+        for(auto &new_message : new_messages) {
+            message_by_event_id[new_message->event_id] = new_message;
+        }
+        messages.insert(messages.begin(), new_messages.rbegin(), new_messages.rend());
+    }
+
+    void RoomData::append_messages(std::vector<std::shared_ptr<Message>> new_messages) {
+        std::lock_guard<std::mutex> lock(room_mutex);
+        for(auto &new_message : new_messages) {
+            message_by_event_id[new_message->event_id] = new_message;
+        }
+        messages.insert(messages.end(), new_messages.begin(), new_messages.end());
+    }
+
+    std::shared_ptr<Message> RoomData::get_message_by_id(const std::string &id) {
+        std::lock_guard<std::mutex> lock(room_mutex);
+        auto message_it = message_by_event_id.find(id);
+        if(message_it == message_by_event_id.end())
+            return nullptr;
+        return message_it->second;
+    }
+
+    std::vector<std::shared_ptr<UserInfo>> RoomData::get_users_excluding_me(const std::string &my_user_id) {
+        std::lock_guard<std::mutex> lock(user_mutex);
+        std::vector<std::shared_ptr<UserInfo>> users_excluding_me;
+        for(auto &[user_id, user] : user_info_by_user_id) {
+            if(user->user_id != my_user_id) {
+                users_excluding_me.push_back(user);
+            }
+        }
+        return users_excluding_me;
+    }
+
+    void RoomData::acquire_room_lock() {
+        room_mutex.lock();
+    }
+
+    void RoomData::release_room_lock() {
+        room_mutex.unlock();
+    }
+
+    const std::vector<std::shared_ptr<Message>>& RoomData::get_messages_thread_unsafe() const {
+        return messages;
+    }
+
     Matrix::Matrix() : Plugin("matrix") {
         
     }
@@ -101,18 +173,18 @@ namespace QuickMedia {
             std::string room_name;
             std::string avatar_url;
 
-            auto room_it = room_data_by_id.find(room_id_str);
-            if(room_it == room_data_by_id.end()) {
-                auto room_data = std::make_unique<RoomData>();
-                room_data->id = room_id_json.asString();
-                room_data_by_id.insert(std::make_pair(room_id_str, std::move(room_data)));
+            auto room = get_room_by_id(room_id_str);
+            if(!room) {
+                room = std::make_shared<RoomData>();
+                room->id = room_id_json.asString();
+                add_room(std::move(room));
                 room_name = room_id_str;
                 fprintf(stderr, "Missing room %s from /sync, adding in joined_rooms\n", room_id_str.c_str());
             } else {
-                room_name = room_it->second->name;
+                room_name = room->name;
                 if(room_name.empty())
                     room_name = room_id_str;
-                avatar_url = room_it->second->avatar_url;
+                avatar_url = room->avatar_url;
             }
 
             auto body_item = BodyItem::create(std::move(room_name));
@@ -124,21 +196,20 @@ namespace QuickMedia {
         return PluginResult::OK;
     }
 
-    static void room_messages_to_body_items(RoomData *room_data, std::shared_ptr<Message> *messages, size_t num_messages, BodyItems &result_items) {
+    static void room_messages_to_body_items(const std::shared_ptr<Message> *messages, size_t num_messages, BodyItems &result_items) {
         for(size_t i = 0; i < num_messages; ++i) {
-            const UserInfo &user_info = room_data->user_info[messages[i]->user_id];
             auto body_item = BodyItem::create("");
-            body_item->set_author(user_info.display_name);
+            body_item->set_author(messages[i]->user->display_name);
             body_item->set_description(messages[i]->body);
             if(!messages[i]->thumbnail_url.empty())
                 body_item->thumbnail_url = messages[i]->thumbnail_url;
             else if(!messages[i]->url.empty() && messages[i]->type == MessageType::IMAGE)
                 body_item->thumbnail_url = messages[i]->url;
             else
-                body_item->thumbnail_url = user_info.avatar_url;
+                body_item->thumbnail_url = messages[i]->user->avatar_url;
             // TODO: Show image thumbnail inline instead of url to image and showing it as the thumbnail of the body item
             body_item->url = messages[i]->url;
-            body_item->author_color = user_info.display_name_color;
+            body_item->author_color = messages[i]->user->display_name_color;
             body_item->userdata = (void*)messages[i].get(); // Note: messages[i] has to be valid as long as body_item is used!
             result_items.push_back(std::move(body_item));
         }
@@ -146,65 +217,74 @@ namespace QuickMedia {
 
     // TODO: Merge common code with |get_new_room_messages|
     PluginResult Matrix::get_all_synced_room_messages(const std::string &room_id, BodyItems &result_items) {
-        auto room_it = room_data_by_id.find(room_id);
-        if(room_it == room_data_by_id.end()) {
+        auto room = get_room_by_id(room_id);
+        if(!room) {
             fprintf(stderr, "Error: no such room: %s\n", room_id.c_str());
             return PluginResult::ERR;
         }
 
-        if(!room_it->second->initial_fetch_finished) {
-            PluginResult result = get_previous_room_messages(room_id, room_it->second.get());
+        // TODO: Thread safe?
+        if(!room->initial_fetch_finished) {
+            PluginResult result = get_previous_room_messages(room);
             if(result == PluginResult::OK) {
-                room_it->second->initial_fetch_finished = true;
+                room->initial_fetch_finished = true;
             } else {
                 fprintf(stderr, "Initial sync failed for room: %s\n", room_id.c_str());
                 return result;
             }
         }
 
-        room_messages_to_body_items(room_it->second.get(), room_it->second->messages.data(), room_it->second->messages.size(), result_items);
-        room_it->second->last_read_index = room_it->second->messages.size();
+        room->acquire_room_lock();
+        room_messages_to_body_items(room->get_messages_thread_unsafe().data(), room->get_messages_thread_unsafe().size(), result_items);
+        room->last_read_index = room->get_messages_thread_unsafe().size();
+        room->release_room_lock();
         return PluginResult::OK;
     }
 
     PluginResult Matrix::get_new_room_messages(const std::string &room_id, BodyItems &result_items) {
-        auto room_it = room_data_by_id.find(room_id);
-        if(room_it == room_data_by_id.end()) {
+        auto room = get_room_by_id(room_id);
+        if(!room) {
             fprintf(stderr, "Error: no such room: %s\n", room_id.c_str());
             return PluginResult::ERR;
         }
 
-        if(!room_it->second->initial_fetch_finished) {
-            PluginResult result = get_previous_room_messages(room_id, room_it->second.get());
+        if(!room->initial_fetch_finished) {
+            PluginResult result = get_previous_room_messages(room);
             if(result == PluginResult::OK) {
-                room_it->second->initial_fetch_finished = true;
+                room->initial_fetch_finished = true;
             } else {
                 fprintf(stderr, "Initial sync failed for room: %s\n", room_id.c_str());
                 return result;
             }
         }
 
-        size_t num_new_messages = room_it->second->messages.size() - room_it->second->last_read_index;
-        room_messages_to_body_items(room_it->second.get(), room_it->second->messages.data() + room_it->second->last_read_index, num_new_messages, result_items);
-        room_it->second->last_read_index = room_it->second->messages.size();
+        room->acquire_room_lock();
+        size_t num_new_messages = room->get_messages_thread_unsafe().size() - room->last_read_index;
+        room_messages_to_body_items(room->get_messages_thread_unsafe().data() + room->last_read_index, num_new_messages, result_items);
+        room->last_read_index = room->get_messages_thread_unsafe().size();
+        room->release_room_lock();
         return PluginResult::OK;
     }
 
     PluginResult Matrix::get_previous_room_messages(const std::string &room_id, BodyItems &result_items) {
-        auto room_it = room_data_by_id.find(room_id);
-        if(room_it == room_data_by_id.end()) {
+        auto room = get_room_by_id(room_id);
+        if(!room) {
             fprintf(stderr, "Error: no such room: %s\n", room_id.c_str());
             return PluginResult::ERR;
         }
         
-        size_t num_messages_before = room_it->second->messages.size();
-        PluginResult result = get_previous_room_messages(room_id, room_it->second.get());
+        room->acquire_room_lock();
+        size_t num_messages_before = room->get_messages_thread_unsafe().size();
+        room->release_room_lock();
+        PluginResult result = get_previous_room_messages(room);
         if(result != PluginResult::OK)
             return result;
 
-        size_t num_messages_after = room_it->second->messages.size();
+        room->acquire_room_lock();
+        size_t num_messages_after = room->get_messages_thread_unsafe().size();
         size_t num_new_messages = num_messages_after - num_messages_before;
-        room_messages_to_body_items(room_it->second.get(), room_it->second->messages.data(), num_new_messages, result_items);
+        room_messages_to_body_items(room->get_messages_thread_unsafe().data(), num_new_messages, result_items);
+        room->release_room_lock();
         return PluginResult::OK;
     }
 
@@ -230,28 +310,27 @@ namespace QuickMedia {
 
             std::string room_id_str = room_id.asString();
 
-            auto room_it = room_data_by_id.find(room_id_str);
-            if(room_it == room_data_by_id.end()) {
-                auto room_data = std::make_unique<RoomData>();
-                room_data->id = room_id_str;
-                room_data_by_id.insert(std::make_pair(room_id_str, std::move(room_data)));
-                room_it = room_data_by_id.find(room_id_str); // TODO: Get iterator from above insert
+            auto room = get_room_by_id(room_id_str);
+            if(!room) {
+                room = std::make_shared<RoomData>();
+                room->id = room_id_str;
+                add_room(room);
             }
 
             const Json::Value &state_json = (*it)["state"];
             if(state_json.isObject()) {
                 const Json::Value &events_json = state_json["events"];
-                events_add_user_info(events_json, room_it->second.get());
-                events_set_room_name(events_json, room_it->second.get());
+                events_add_user_info(events_json, room.get());
+                events_set_room_name(events_json, room.get());
             }
 
             const Json::Value &timeline_json = (*it)["timeline"];
             if(timeline_json.isObject()) {
-                if(room_it->second->prev_batch.empty()) {
+                if(room->prev_batch.empty()) {
                     // This may be non-existent if this is the first event in the room
                     const Json::Value &prev_batch_json = timeline_json["prev_batch"];
                     if(prev_batch_json.isString())
-                        room_it->second->prev_batch = prev_batch_json.asString();
+                        room->prev_batch = prev_batch_json.asString();
                 }
 
                 // TODO: Is there no better way to check for notifications? this is not robust...
@@ -264,15 +343,15 @@ namespace QuickMedia {
                 }
 
                 const Json::Value &events_json = timeline_json["events"];
-                events_add_user_info(events_json, room_it->second.get());
-                events_add_messages(events_json, room_it->second.get(), MessageDirection::AFTER, &room_messages, has_unread_notifications);
-                events_set_room_name(events_json, room_it->second.get());
+                events_add_user_info(events_json, room.get());
+                events_add_messages(events_json, room, MessageDirection::AFTER, &room_messages, has_unread_notifications);
+                events_set_room_name(events_json, room.get());
             }
 
             const Json::Value &ephemeral_json = (*it)["ephemeral"];
             if(ephemeral_json.isObject()) {
                 const Json::Value &events_json = ephemeral_json["events"];
-                events_add_user_read_markers(events_json, room_it->second.get());
+                events_add_user_read_markers(events_json, room.get());
             }
         }
 
@@ -327,23 +406,18 @@ namespace QuickMedia {
 
             std::string sender_json_str = sender_json.asString();
 
-            UserInfo user_info;
-            user_info.user_id = sender_json_str;
-            user_info.avatar_url = std::move(avatar_url_str);
-            if(strncmp(user_info.avatar_url.c_str(), "mxc://", 6) == 0)
-                user_info.avatar_url.erase(user_info.avatar_url.begin(), user_info.avatar_url.begin() + 6);
+            auto user_info = std::make_shared<UserInfo>();
+            user_info->user_id = sender_json_str;
+            user_info->avatar_url = std::move(avatar_url_str);
+            if(strncmp(user_info->avatar_url.c_str(), "mxc://", 6) == 0)
+                user_info->avatar_url.erase(user_info->avatar_url.begin(), user_info->avatar_url.begin() + 6);
             // TODO: What if the user hasn't selected an avatar?
-            user_info.avatar_url = homeserver + "/_matrix/media/r0/thumbnail/" + user_info.avatar_url + "?width=32&height=32&method=crop";
-            user_info.display_name = display_name_json.asString();
-            user_info.display_name_color = user_id_to_color(sender_json_str);
+            user_info->avatar_url = homeserver + "/_matrix/media/r0/thumbnail/" + user_info->avatar_url + "?width=32&height=32&method=crop";
+            user_info->display_name = display_name_json.asString();
+            user_info->display_name_color = user_id_to_color(sender_json_str);
             
-            auto user_it = room_data->user_info_by_user_id.find(sender_json_str);
-            if(user_it != room_data->user_info_by_user_id.end()) {
-                room_data->user_info[user_it->second] = std::move(user_info);
-            } else {
-                room_data->user_info.push_back(user_info);
-                room_data->user_info_by_user_id.insert(std::make_pair(sender_json_str, room_data->user_info.size() - 1));
-            }
+            // Overwrites user data
+            room_data->add_user(std::move(user_info));
         }
     }
 
@@ -385,13 +459,13 @@ namespace QuickMedia {
                     if(!user_id_json.isString())
                         continue;
 
-                    auto user_it = room_data->user_info_by_user_id.find(user_id_json.asString());
-                    if(user_it == room_data->user_info_by_user_id.end()) {
+                    auto user = room_data->get_user_by_id(user_id_json.asString());
+                    if(!user) {
                         fprintf(stderr, "Receipt read receipt for unknown user: %s, ignoring...\n", user_id_json.asCString());
                         continue;
                     }
 
-                    room_data->user_info[user_it->second].read_marker_event_id = event_id_str;
+                    room_data->set_user_read_marker(user, event_id_str);
                 }
             }
         }
@@ -463,7 +537,7 @@ namespace QuickMedia {
         return false;
     }
 
-    void Matrix::events_add_messages(const Json::Value &events_json, RoomData *room_data, MessageDirection message_dir, RoomSyncMessages *room_messages, bool has_unread_notifications) {
+    void Matrix::events_add_messages(const Json::Value &events_json, std::shared_ptr<RoomData> &room_data, MessageDirection message_dir, RoomSyncMessages *room_messages, bool has_unread_notifications) {
         if(!events_json.isArray())
             return;
 
@@ -501,8 +575,8 @@ namespace QuickMedia {
             if(!content_type.isString())
                 continue;
 
-            auto user_it = room_data->user_info_by_user_id.find(sender_json_str);
-            if(user_it == room_data->user_info_by_user_id.end()) {
+            auto user = room_data->get_user_by_id(sender_json_str);
+            if(!user) {
                 // Note: this is important because otherwise replying and such is broken
                 fprintf(stderr, "Warning: skipping unknown user: %s\n", sender_json_str.c_str());
                 continue;
@@ -565,7 +639,7 @@ namespace QuickMedia {
                 message->type = MessageType::FILE;
             } else if(strcmp(content_type.asCString(), "m.emote") == 0) { // this is a /me message, TODO: show /me messages differently
                 message->type = MessageType::TEXT;
-                prefix = "*" + room_data->user_info[user_it->second].display_name + "* ";
+                prefix = "*" + user->display_name + "* ";
             } else if(strcmp(content_type.asCString(), "m.notice") == 0) { // TODO: show notices differently
                 message->type = MessageType::TEXT;
                 prefix = "* NOTICE * ";
@@ -580,7 +654,7 @@ namespace QuickMedia {
                 continue;
             }
 
-            message->user_id = user_it->second;
+            message->user = user;
             message->event_id = event_id_str;
             message->body = prefix + body_json.asString();
             message->replaces_event_id = std::move(replaces_event_id);
@@ -589,18 +663,18 @@ namespace QuickMedia {
                 message->mentions_me = message_contains_user_mention(message->body, username) || message_contains_user_mention(message->body, "@room");
             message->timestamp = timestamp;
             new_messages.push_back(message);
-            room_data->message_by_event_id[event_id_str] = message;
             if(room_sync_messages)
                 room_sync_messages->push_back(message);
         }
 
         // TODO: Loop and std::move instead? doesn't insert create copies?
         if(message_dir == MessageDirection::BEFORE) {
-            room_data->messages.insert(room_data->messages.begin(), new_messages.rbegin(), new_messages.rend());
+            room_data->prepend_messages_reverse(std::move(new_messages));
+            // TODO: Is this thread-safe?
             if(room_data->last_read_index != 0)
                 room_data->last_read_index += new_messages.size();
         } else if(message_dir == MessageDirection::AFTER) {
-            room_data->messages.insert(room_data->messages.end(), new_messages.begin(), new_messages.end());
+            room_data->append_messages(std::move(new_messages));
         }
     }
 
@@ -612,17 +686,7 @@ namespace QuickMedia {
         return user_id.substr(1, index - 1);
     }
 
-    static std::vector<UserInfo*> get_users_excluding_me(const std::string &my_user_id, std::vector<UserInfo> &user_info) {
-        std::vector<UserInfo*> users_excluding_me;
-        for(UserInfo &user : user_info) {
-            if(user.user_id != my_user_id) {
-                users_excluding_me.push_back(&user);
-            }
-        }
-        return users_excluding_me;
-    }
-
-    static std::string combine_user_display_names_for_room_name(const std::vector<UserInfo*> &user_info, const std::string &fallback_user_id) {
+    static std::string combine_user_display_names_for_room_name(const std::vector<std::shared_ptr<UserInfo>> &user_info, const std::string &fallback_user_id) {
         std::string result;
         if(user_info.size() == 0)
             result = extract_user_name_from_user_id(fallback_user_id);
@@ -658,9 +722,9 @@ namespace QuickMedia {
             room_data->name = name_json.asString();
         }
 
-        std::vector<UserInfo*> users_excluding_me;
+        std::vector<std::shared_ptr<UserInfo>> users_excluding_me;
         if(room_data->name.empty() || room_data->avatar_url.empty())
-            users_excluding_me = get_users_excluding_me(user_id, room_data->user_info);
+            users_excluding_me = room_data->get_users_excluding_me(user_id); // TODO: What about thread safety with user_id? its reset in /logout
 
         for(const Json::Value &event_item_json : events_json) {
             if(!event_item_json.isObject())
@@ -683,9 +747,9 @@ namespace QuickMedia {
 
             if(room_data->avatar_url.empty()) {
                 if(users_excluding_me.empty()) {
-                    auto user_it = room_data->user_info_by_user_id.find(creator_json.asString());
-                    if(user_it != room_data->user_info_by_user_id.end())
-                        room_data->avatar_url = room_data->user_info[user_it->second].avatar_url;
+                    auto user = room_data->get_user_by_id(creator_json.asString());
+                    if(user)
+                        room_data->avatar_url = user->avatar_url;
                 } else {
                     // TODO: If there are multiple users, then we want to use some other type of avatar, not the first users avatar
                     room_data->avatar_url = users_excluding_me.front()->avatar_url;
@@ -714,10 +778,10 @@ namespace QuickMedia {
         }
     }
 
-    PluginResult Matrix::get_previous_room_messages(const std::string &room_id, RoomData *room_data) {
+    PluginResult Matrix::get_previous_room_messages(std::shared_ptr<RoomData> &room_data) {
         std::string from = room_data->prev_batch;
         if(from.empty()) {
-            fprintf(stderr, "Info: missing previous batch for room: %s, using /sync next batch\n", room_id.c_str());
+            fprintf(stderr, "Info: missing previous batch for room: %s, using /sync next batch\n", room_data->id.c_str());
             from = next_batch;
             if(from.empty()) {
                 fprintf(stderr, "Error: missing next batch!\n");
@@ -739,7 +803,7 @@ namespace QuickMedia {
         std::string filter = url_param_encode(Json::writeString(builder, std::move(request_data)));
 
         char url[512];
-        snprintf(url, sizeof(url), "%s/_matrix/client/r0/rooms/%s/messages?from=%s&limit=20&dir=b&filter=%s", homeserver.c_str(), room_id.c_str(), from.c_str(), filter.c_str());
+        snprintf(url, sizeof(url), "%s/_matrix/client/r0/rooms/%s/messages?from=%s&limit=20&dir=b&filter=%s", homeserver.c_str(), room_data->id.c_str(), from.c_str(), filter.c_str());
         fprintf(stderr, "load initial room data, url: |%s|\n", url);
 
         std::string server_response;
@@ -762,8 +826,8 @@ namespace QuickMedia {
             return PluginResult::ERR;
 
         const Json::Value &state_json = json_root["state"];
-        events_add_user_info(state_json, room_data);
-        events_set_room_name(state_json, room_data);
+        events_add_user_info(state_json, room_data.get());
+        events_set_room_name(state_json, room_data.get());
 
         const Json::Value &chunk_json = json_root["chunk"];
         events_add_messages(chunk_json, room_data, MessageDirection::BEFORE, nullptr, false);
@@ -954,7 +1018,7 @@ namespace QuickMedia {
         return result;
     }
 
-    static std::string create_body_for_message_reply(const RoomData *room_data, const Message *message, const std::string &body) {
+    static std::string create_body_for_message_reply(const Message *message, const std::string &body) {
         std::string related_to_body;
         switch(message->type) {
             case MessageType::TEXT: {
@@ -977,20 +1041,21 @@ namespace QuickMedia {
                 related_to_body = "sent a file";
                 break;
         }
-        return "> <" + room_data->user_info[message->user_id].user_id + "> " + block_quote(std::move(related_to_body)) + "\n\n" + body;
+        return "> <" + message->user->user_id + "> " + block_quote(std::move(related_to_body)) + "\n\n" + body;
     }   
 
     // TODO: Add formatted_body just like element does with <mx-reply><blockquote... and also support greentext with that
     PluginResult Matrix::post_reply(const std::string &room_id, const std::string &body, void *relates_to) {
-        auto room_it = room_data_by_id.find(room_id);
-        if(room_it == room_data_by_id.end()) {
+        auto room = get_room_by_id(room_id);
+        if(!room) {
             fprintf(stderr, "Error: no such room: %s\n", room_id.c_str());
             return PluginResult::ERR;
         }
 
+        // TODO: Store shared_ptr<Message> instead of raw pointer...
         Message *relates_to_message_raw = (Message*)relates_to;
-        std::shared_ptr<Message> relates_to_message_shared = room_it->second->message_by_event_id[relates_to_message_raw->event_id];
-        std::shared_ptr<Message> relates_to_message_original = get_edited_message_original_message(room_it->second.get(), relates_to_message_shared);
+        std::shared_ptr<Message> relates_to_message_shared = room->get_message_by_id(relates_to_message_raw->event_id);
+        std::shared_ptr<Message> relates_to_message_original = get_edited_message_original_message(room.get(), relates_to_message_shared);
         if(!relates_to_message_original) {
             fprintf(stderr, "Failed to get the original message for message with event id: %s\n", relates_to_message_raw->event_id.c_str());
             return PluginResult::ERR;
@@ -1010,7 +1075,7 @@ namespace QuickMedia {
 
         Json::Value request_data(Json::objectValue);
         request_data["msgtype"] = "m.text"; // TODO: Allow image reply? element doesn't do that but we could!
-        request_data["body"] = create_body_for_message_reply(room_it->second.get(), relates_to_message_raw, body); // Yes, the reply is to the edited message but the event_id reference is to the original message...
+        request_data["body"] = create_body_for_message_reply(relates_to_message_raw, body); // Yes, the reply is to the edited message but the event_id reference is to the original message...
         request_data["m.relates_to"] = std::move(relates_to_json);
 
         Json::StreamWriterBuilder builder;
@@ -1056,15 +1121,15 @@ namespace QuickMedia {
     }
 
     PluginResult Matrix::post_edit(const std::string &room_id, const std::string &body, void *relates_to) {
-        auto room_it = room_data_by_id.find(room_id);
-        if(room_it == room_data_by_id.end()) {
+        auto room = get_room_by_id(room_id);
+        if(!room) {
             fprintf(stderr, "Error: no such room: %s\n", room_id.c_str());
             return PluginResult::ERR;
         }
 
         Message *relates_to_message_raw = (Message*)relates_to;
-        std::shared_ptr<Message> relates_to_message_shared = room_it->second->message_by_event_id[relates_to_message_raw->event_id];
-        std::shared_ptr<Message> relates_to_message_original = get_edited_message_original_message(room_it->second.get(), relates_to_message_shared);
+        std::shared_ptr<Message> relates_to_message_shared = room->get_message_by_id(relates_to_message_raw->event_id);
+        std::shared_ptr<Message> relates_to_message_original = get_edited_message_original_message(room.get(), relates_to_message_shared);
         if(!relates_to_message_original) {
             fprintf(stderr, "Failed to get the original message for message with event id: %s\n", relates_to_message_raw->event_id.c_str());
             return PluginResult::ERR;
@@ -1166,8 +1231,8 @@ namespace QuickMedia {
         if(message->replaces_event_id.empty())
             return message;
 
-        auto message_it = room_data->message_by_event_id.find(message->replaces_event_id);
-        if(message_it == room_data->message_by_event_id.end()) {
+        auto replaced_message = room_data->get_message_by_id(message->replaces_event_id);
+        if(!replaced_message) {
             Json::Value request_data(Json::objectValue);
             request_data["lazy_load_members"] = true;
 
@@ -1234,7 +1299,6 @@ namespace QuickMedia {
                 return nullptr;
 
             auto new_message = std::make_shared<Message>();
-            new_message->user_id = -1;
             new_message->event_id = event_id_json.asString();
             new_message->replaces_event_id = std::move(replaces_event_id);
             if(strcmp(content_type.asCString(), "m.text") == 0) {
@@ -1253,7 +1317,7 @@ namespace QuickMedia {
 
             return get_edited_message_original_message(room_data, std::move(new_message));
         } else {
-            return get_edited_message_original_message(room_data, message_it->second);
+            return get_edited_message_original_message(room_data, replaced_message);
         }
     }
 
@@ -1696,22 +1760,14 @@ namespace QuickMedia {
         return PluginResult::OK;
     }
 
-    bool Matrix::was_message_posted_by_me(const std::string &room_id, void *message) const {
-        auto room_it = room_data_by_id.find(room_id);
-        if(room_it == room_data_by_id.end()) {
-            fprintf(stderr, "Error: no such room: %s\n", room_id.c_str());
-            return false;
-        }
+    bool Matrix::was_message_posted_by_me(void *message) {
         Message *message_typed = (Message*)message;
-        return user_id == room_it->second->user_info[message_typed->user_id].user_id;
+        return user_id == message_typed->user->user_id;
     }
 
-    std::string Matrix::message_get_author_displayname(RoomData *room_data, Message *message) const {
-        if(message->user_id >= room_data->user_info.size()) {
-            fprintf(stderr, "Warning: no user with the index: %zu\n", message->user_id);
-            return "";
-        }
-        return room_data->user_info[message->user_id].display_name;
+    std::string Matrix::message_get_author_displayname(Message *message) const {
+        // TODO: Thread safe?
+        return message->user->display_name;
     }
 
     PluginResult Matrix::get_config(int *upload_size) {
@@ -1761,12 +1817,25 @@ namespace QuickMedia {
         return PluginResult::OK;
     }
 
-    const UserInfo* Matrix::get_me(const std::string &room_id) const {
-        auto room_it = room_data_by_id.find(room_id);
-        if(room_it == room_data_by_id.end()) {
+    std::shared_ptr<UserInfo> Matrix::get_me(const std::string &room_id) {
+        auto room = get_room_by_id(room_id);
+        if(!room) {
             fprintf(stderr, "Error: no such room: %s\n", room_id.c_str());
             return nullptr;
         }
-        return &room_it->second->user_info[room_it->second->user_info_by_user_id[user_id]];
+        return room->get_user_by_id(user_id);
+    }
+
+    std::shared_ptr<RoomData> Matrix::get_room_by_id(const std::string &id) {
+        std::lock_guard<std::mutex> lock(room_data_mutex);
+        auto room_it = room_data_by_id.find(id);
+        if(room_it == room_data_by_id.end())
+            return nullptr;
+        return room_it->second;
+    }
+
+    void Matrix::add_room(std::shared_ptr<RoomData> room) {
+        std::lock_guard<std::mutex> lock(room_data_mutex);
+        room_data_by_id.insert(std::make_pair(room->id, room));
     }
 }
\ No newline at end of file