Properly share and rotate sessions on member and device changes

src/Cache.cpp

@@ -362,6 +362,7 @@ Cache::inboundMegolmSessionExists(const MegolmSessionIndex &index)
 
 void
 Cache::updateOutboundMegolmSession(const std::string &room_id,
+                                   const OutboundGroupSessionData &data_,
                                    mtx::crypto::OutboundGroupSessionPtr &ptr)
 {
         using namespace mtx::crypto;
@@ -369,7 +370,7 @@ Cache::updateOutboundMegolmSession(const std::string &room_id,
         if (!outboundMegolmSessionExists(room_id))
                 return;
 
-        OutboundGroupSessionData data;
+        OutboundGroupSessionData data = data_;
         data.message_index            = olm_outbound_group_session_message_index(ptr.get());
         data.session_id               = mtx::crypto::session_id(ptr.get());
         data.session_key              = mtx::crypto::session_key(ptr.get());
@@ -402,7 +403,7 @@ Cache::dropOutboundMegolmSession(const std::string &room_id)
 void
 Cache::saveOutboundMegolmSession(const std::string &room_id,
                                  const OutboundGroupSessionData &data,
-                                 mtx::crypto::OutboundGroupSessionPtr session)
+                                 mtx::crypto::OutboundGroupSessionPtr &session)
 {
         using namespace mtx::crypto;
         const auto pickled = pickle<OutboundSessionObject>(session.get(), SECRET);
@@ -3095,6 +3096,39 @@ Cache::roomMembers(const std::string &room_id)
         return members;
 }
 
+std::map<std::string, std::optional<UserKeyCache>>
+Cache::getMembersWithKeys(const std::string &room_id)
+{
+        lmdb::val keys;
+
+        try {
+                auto txn = lmdb::txn::begin(env_, nullptr, MDB_RDONLY);
+                std::map<std::string, std::optional<UserKeyCache>> members;
+
+                auto db     = getMembersDb(txn, room_id);
+                auto keysDb = getUserKeysDb(txn);
+
+                std::string user_id, unused;
+                auto cursor = lmdb::cursor::open(txn, db);
+                while (cursor.get(user_id, unused, MDB_NEXT)) {
+                        auto res = lmdb::dbi_get(txn, keysDb, lmdb::val(user_id), keys);
+
+                        if (res) {
+                                members[user_id] =
+                                  json::parse(std::string_view(keys.data(), keys.size()))
+                                    .get<UserKeyCache>();
+                        } else {
+                                members[user_id] = {};
+                        }
+                }
+                cursor.close();
+
+                return members;
+        } catch (std::exception &) {
+                return {};
+        }
+}
+
 QString
 Cache::displayName(const QString &room_id, const QString &user_id)
 {
@@ -3235,6 +3269,8 @@ Cache::updateUserKeys(const std::string &sync_token, const mtx::responses::Query
                 updates[user].self_signing_keys = keys;
 
         for (auto &[user, update] : updates) {
+                nhlog::db()->debug("Updated user keys: {}", user);
+
                 lmdb::val oldKeys;
                 auto res = lmdb::dbi_get(txn, db, lmdb::val(user), oldKeys);
 
@@ -3297,6 +3333,8 @@ Cache::markUserKeysOutOfDate(lmdb::txn &txn,
         query.token = sync_token;
 
         for (const auto &user : user_ids) {
+                nhlog::db()->debug("Marking user keys out of date: {}", user);
+
                 lmdb::val oldKeys;
                 auto res = lmdb::dbi_get(txn, db, lmdb::val(user), oldKeys);
 
@@ -3650,12 +3688,41 @@ from_json(const json &j, MemberInfo &info)
         info.avatar_url = j.at("avatar_url");
 }
 
+void
+to_json(nlohmann::json &obj, const DeviceAndMasterKeys &msg)
+{
+        obj["devices"]     = msg.devices;
+        obj["master_keys"] = msg.master_keys;
+}
+
+void
+from_json(const nlohmann::json &obj, DeviceAndMasterKeys &msg)
+{
+        msg.devices     = obj.at("devices").get<decltype(msg.devices)>();
+        msg.master_keys = obj.at("master_keys").get<decltype(msg.master_keys)>();
+}
+
+void
+to_json(nlohmann::json &obj, const SharedWithUsers &msg)
+{
+        obj["keys"] = msg.keys;
+}
+
+void
+from_json(const nlohmann::json &obj, SharedWithUsers &msg)
+{
+        msg.keys = obj.at("keys").get<std::map<std::string, DeviceAndMasterKeys>>();
+}
+
 void
 to_json(nlohmann::json &obj, const OutboundGroupSessionData &msg)
 {
         obj["session_id"]    = msg.session_id;
         obj["session_key"]   = msg.session_key;
         obj["message_index"] = msg.message_index;
+
+        obj["initially"] = msg.initially;
+        obj["currently"] = msg.currently;
 }
 
 void
@@ -3664,6 +3731,9 @@ from_json(const nlohmann::json &obj, OutboundGroupSessionData &msg)
         msg.session_id    = obj.at("session_id");
         msg.session_key   = obj.at("session_key");
         msg.message_index = obj.at("message_index");
+
+        msg.initially = obj.value("initially", SharedWithUsers{});
+        msg.currently = obj.value("currently", SharedWithUsers{});
 }
 
 void
@@ -4098,9 +4168,9 @@ isRoomMember(const std::string &user_id, const std::string &room_id)
 void
 saveOutboundMegolmSession(const std::string &room_id,
                           const OutboundGroupSessionData &data,
-                          mtx::crypto::OutboundGroupSessionPtr session)
+                          mtx::crypto::OutboundGroupSessionPtr &session)
 {
-        instance_->saveOutboundMegolmSession(room_id, data, std::move(session));
+        instance_->saveOutboundMegolmSession(room_id, data, session);
 }
 OutboundGroupSessionDataRef
 getOutboundMegolmSession(const std::string &room_id)
@@ -4114,9 +4184,10 @@ outboundMegolmSessionExists(const std::string &room_id) noexcept
 }
 void
 updateOutboundMegolmSession(const std::string &room_id,
+                            const OutboundGroupSessionData &data,
                             mtx::crypto::OutboundGroupSessionPtr &session)
 {
-        instance_->updateOutboundMegolmSession(room_id, session);
+        instance_->updateOutboundMegolmSession(room_id, data, session);
 }
 void
 dropOutboundMegolmSession(const std::string &room_id)

src/Cache.h

@@ -235,13 +235,14 @@ isRoomMember(const std::string &user_id, const std::string &room_id);
 void
 saveOutboundMegolmSession(const std::string &room_id,
                           const OutboundGroupSessionData &data,
-                          mtx::crypto::OutboundGroupSessionPtr session);
+                          mtx::crypto::OutboundGroupSessionPtr &session);
 OutboundGroupSessionDataRef
 getOutboundMegolmSession(const std::string &room_id);
 bool
 outboundMegolmSessionExists(const std::string &room_id) noexcept;
 void
 updateOutboundMegolmSession(const std::string &room_id,
+                            const OutboundGroupSessionData &data,
                             mtx::crypto::OutboundGroupSessionPtr &session);
 void
 dropOutboundMegolmSession(const std::string &room_id);

src/CacheCryptoStructs.h

@@ -6,12 +6,28 @@
 #include <mtx/responses/crypto.hpp>
 #include <mtxclient/crypto/objects.hpp>
 
+struct DeviceAndMasterKeys
+{
+        // map from device id or master key id to message_index
+        std::map<std::string, uint64_t> devices, master_keys;
+};
+
+struct SharedWithUsers
+{
+        // userid to keys
+        std::map<std::string, DeviceAndMasterKeys> keys;
+};
+
 // Extra information associated with an outbound megolm session.
 struct OutboundGroupSessionData
 {
         std::string session_id;
         std::string session_key;
         uint64_t message_index = 0;
+
+        // who has access to this session.
+        // Rotate, when a user leaves the room and share, when a user gets added.
+        SharedWithUsers initially, currently;
 };
 
 void

src/Cache_p.h

@@ -59,6 +59,8 @@ public:
 
         // user cache stores user keys
         std::optional<UserKeyCache> userKeys(const std::string &user_id);
+        std::map<std::string, std::optional<UserKeyCache>> getMembersWithKeys(
+          const std::string &room_id);
         void updateUserKeys(const std::string &sync_token,
                             const mtx::responses::QueryKeys &keyQuery);
         void markUserKeysOutOfDate(lmdb::txn &txn,
@@ -232,10 +234,11 @@ public:
         //
         void saveOutboundMegolmSession(const std::string &room_id,
                                        const OutboundGroupSessionData &data,
-                                       mtx::crypto::OutboundGroupSessionPtr session);
+                                       mtx::crypto::OutboundGroupSessionPtr &session);
         OutboundGroupSessionDataRef getOutboundMegolmSession(const std::string &room_id);
         bool outboundMegolmSessionExists(const std::string &room_id) noexcept;
         void updateOutboundMegolmSession(const std::string &room_id,
+                                         const OutboundGroupSessionData &data,
                                          mtx::crypto::OutboundGroupSessionPtr &session);
         void dropOutboundMegolmSession(const std::string &room_id);
 

src/Olm.cpp

@@ -278,11 +278,168 @@ mtx::events::msg::Encrypted
 encrypt_group_message(const std::string &room_id, const std::string &device_id, nlohmann::json body)
 {
         using namespace mtx::events;
+        using namespace mtx::identifiers;
+
+        auto own_user_id = http::client()->user_id().to_string();
+
+        auto members = cache::client()->getMembersWithKeys(room_id);
+
+        std::map<std::string, std::vector<std::string>> sendSessionTo;
+        mtx::crypto::OutboundGroupSessionPtr session = nullptr;
+        OutboundGroupSessionData group_session_data;
+
+        if (cache::outboundMegolmSessionExists(room_id)) {
+                auto res = cache::getOutboundMegolmSession(room_id);
+
+                auto member_it             = members.begin();
+                auto session_member_it     = res.data.currently.keys.begin();
+                auto session_member_it_end = res.data.currently.keys.end();
+
+                while (member_it != members.end() || session_member_it != session_member_it_end) {
+                        if (member_it == members.end()) {
+                                // a member left, purge session!
+                                nhlog::crypto()->debug(
+                                  "Rotating megolm session because of left member");
+                                break;
+                        }
+
+                        if (session_member_it == session_member_it_end) {
+                                // share with all remaining members
+                                while (member_it != members.end()) {
+                                        sendSessionTo[member_it->first] = {};
+
+                                        if (member_it->second)
+                                                for (const auto &dev :
+                                                     member_it->second->device_keys)
+                                                        if (member_it->first != own_user_id ||
+                                                            dev.first != device_id)
+                                                                sendSessionTo[member_it->first]
+                                                                  .push_back(dev.first);
+
+                                        ++member_it;
+                                }
+
+                                session = std::move(res.session);
+                                break;
+                        }
+
+                        if (member_it->first > session_member_it->first) {
+                                // a member left, purge session
+                                nhlog::crypto()->debug(
+                                  "Rotating megolm session because of left member");
+                                break;
+                        } else if (member_it->first < session_member_it->first) {
+                                // new member, send them the session at this index
+                                sendSessionTo[member_it->first] = {};
+
+                                for (const auto &dev : member_it->second->device_keys)
+                                        if (member_it->first != own_user_id ||
+                                            dev.first != device_id)
+                                                sendSessionTo[member_it->first].push_back(
+                                                  dev.first);
+
+                                ++member_it;
+                        } else {
+                                // compare devices
+                                bool device_removed = false;
+                                for (const auto &dev : session_member_it->second.devices) {
+                                        if (!member_it->second ||
+                                            !member_it->second->device_keys.count(dev.first)) {
+                                                device_removed = true;
+                                                break;
+                                        }
+                                }
+
+                                if (device_removed) {
+                                        // device removed, rotate session!
+                                        nhlog::crypto()->debug(
+                                          "Rotating megolm session because of removed device of {}",
+                                          member_it->first);
+                                        break;
+                                }
+
+                                // check for new devices to share with
+                                if (member_it->second)
+                                        for (const auto &dev : member_it->second->device_keys)
+                                                if (!session_member_it->second.devices.count(
+                                                      dev.first) &&
+                                                    (member_it->first != own_user_id ||
+                                                     dev.first != device_id))
+                                                        sendSessionTo[member_it->first].push_back(
+                                                          dev.first);
+
+                                ++member_it;
+                                ++session_member_it;
+                                if (member_it == members.end() &&
+                                    session_member_it == session_member_it_end) {
+                                        // all devices match or are newly added
+                                        session = std::move(res.session);
+                                }
+                        }
+                }
+
+                group_session_data = std::move(res.data);
+        }
+
+        if (!session) {
+                nhlog::ui()->debug("creating new outbound megolm session");
+
+                // Create a new outbound megolm session.
+                session                = olm::client()->init_outbound_group_session();
+                const auto session_id  = mtx::crypto::session_id(session.get());
+                const auto session_key = mtx::crypto::session_key(session.get());
+
+                // Saving the new megolm session.
+                OutboundGroupSessionData session_data{};
+                session_data.session_id    = mtx::crypto::session_id(session.get());
+                session_data.session_key   = mtx::crypto::session_key(session.get());
+                session_data.message_index = 0;
+
+                sendSessionTo.clear();
+
+                for (const auto &[user, devices] : members) {
+                        sendSessionTo[user]               = {};
+                        session_data.initially.keys[user] = {};
+                        if (devices) {
+                                for (const auto &[device_id_, key] : devices->device_keys) {
+                                        (void)key;
+                                        if (device_id != device_id_ || user != own_user_id) {
+                                                sendSessionTo[user].push_back(device_id_);
+                                                session_data.initially.keys[user]
+                                                  .devices[device_id_] = 0;
+                                        }
+                                }
+                        }
+                }
+
+                cache::saveOutboundMegolmSession(room_id, session_data, session);
+                group_session_data = std::move(session_data);
+
+                {
+                        MegolmSessionIndex index;
+                        index.room_id    = room_id;
+                        index.session_id = session_id;
+                        index.sender_key = olm::client()->identity_keys().curve25519;
+                        auto megolm_session =
+                          olm::client()->init_inbound_group_session(session_key);
+                        cache::saveInboundMegolmSession(index, std::move(megolm_session));
+                }
+        }
+
+        mtx::events::DeviceEvent<mtx::events::msg::RoomKey> megolm_payload{};
+        megolm_payload.content.algorithm   = MEGOLM_ALGO;
+        megolm_payload.content.room_id     = room_id;
+        megolm_payload.content.session_id  = mtx::crypto::session_id(session.get());
+        megolm_payload.content.session_key = mtx::crypto::session_key(session.get());
+        megolm_payload.type                = mtx::events::EventType::RoomKey;
+
+        if (!sendSessionTo.empty())
+                olm::send_encrypted_to_device_messages(sendSessionTo, megolm_payload);
 
-        // relations shouldn't be encrypted...
         mtx::common::ReplyRelatesTo relation;
         mtx::common::RelatesTo r_relation;
 
+        // relations shouldn't be encrypted...
         if (body["content"].contains("m.relates_to") &&
             body["content"]["m.relates_to"].contains("m.in_reply_to")) {
                 relation = body["content"]["m.relates_to"];
@@ -292,25 +449,35 @@ encrypt_group_message(const std::string &room_id, const std::string &device_id,
                 body["content"].erase("m.relates_to");
         }
 
-        // Always check before for existence.
-        auto res     = cache::getOutboundMegolmSession(room_id);
-        auto payload = olm::client()->encrypt_group_message(res.session.get(), body.dump());
+        auto payload = olm::client()->encrypt_group_message(session.get(), body.dump());
 
         // Prepare the m.room.encrypted event.
         msg::Encrypted data;
         data.ciphertext   = std::string((char *)payload.data(), payload.size());
         data.sender_key   = olm::client()->identity_keys().curve25519;
-        data.session_id   = mtx::crypto::session_id(res.session.get());
+        data.session_id   = mtx::crypto::session_id(session.get());
         data.device_id    = device_id;
         data.algorithm    = MEGOLM_ALGO;
         data.relates_to   = relation;
         data.r_relates_to = r_relation;
 
-        res.data.message_index = olm_outbound_group_session_message_index(res.session.get());
-        nhlog::crypto()->debug("next message_index {}", res.data.message_index);
+        group_session_data.message_index = olm_outbound_group_session_message_index(session.get());
+        nhlog::crypto()->debug("next message_index {}", group_session_data.message_index);
+
+        // update current set of members for the session with the new members and that message_index
+        for (const auto &[user, devices] : sendSessionTo) {
+                if (!group_session_data.currently.keys.count(user))
+                        group_session_data.currently.keys[user] = {};
+
+                for (const auto &device_id : devices) {
+                        if (!group_session_data.currently.keys[user].devices.count(device_id))
+                                group_session_data.currently.keys[user].devices[device_id] =
+                                  group_session_data.message_index;
+                }
+        }
 
         // We need to re-pickle the session after we send a message to save the new message_index.
-        cache::updateOutboundMegolmSession(room_id, res.session);
+        cache::updateOutboundMegolmSession(room_id, group_session_data, session);
 
         return data;
 }

src/timeline/TimelineModel.cpp

@@ -910,8 +910,6 @@ TimelineModel::sendEncryptedMessage(mtx::events::RoomEvent<T> msg, mtx::events::
                     {"room_id", room_id}};
 
         try {
-                // Check if we have already an outbound megolm session then we can use.
-                if (cache::outboundMegolmSessionExists(room_id)) {
                 mtx::events::EncryptedEvent<mtx::events::msg::Encrypted> event;
                 event.content =
                   olm::encrypt_group_message(room_id, http::client()->device_id(), doc);
@@ -922,68 +920,6 @@ TimelineModel::sendEncryptedMessage(mtx::events::RoomEvent<T> msg, mtx::events::
                 event.origin_server_ts = QDateTime::currentMSecsSinceEpoch();
 
                 emit this->addPendingMessageToStore(event);
-                        return;
-                }
-
-                nhlog::ui()->debug("creating new outbound megolm session");
-
-                // Create a new outbound megolm session.
-                auto outbound_session  = olm::client()->init_outbound_group_session();
-                const auto session_id  = mtx::crypto::session_id(outbound_session.get());
-                const auto session_key = mtx::crypto::session_key(outbound_session.get());
-
-                mtx::events::DeviceEvent<mtx::events::msg::RoomKey> megolm_payload;
-                megolm_payload.content.algorithm   = "m.megolm.v1.aes-sha2";
-                megolm_payload.content.room_id     = room_id;
-                megolm_payload.content.session_id  = session_id;
-                megolm_payload.content.session_key = session_key;
-                megolm_payload.type                = mtx::events::EventType::RoomKey;
-
-                // Saving the new megolm session.
-                // TODO: Maybe it's too early to save.
-                OutboundGroupSessionData session_data;
-                session_data.session_id    = session_id;
-                session_data.session_key   = session_key;
-                session_data.message_index = 0;
-                cache::saveOutboundMegolmSession(
-                  room_id, session_data, std::move(outbound_session));
-
-                {
-                        MegolmSessionIndex index;
-                        index.room_id    = room_id;
-                        index.session_id = session_id;
-                        index.sender_key = olm::client()->identity_keys().curve25519;
-                        auto megolm_session =
-                          olm::client()->init_inbound_group_session(session_key);
-                        cache::saveInboundMegolmSession(index, std::move(megolm_session));
-                }
-
-                const auto members = cache::roomMembers(room_id);
-                nhlog::ui()->info("retrieved {} members for {}", members.size(), room_id);
-
-                std::map<std::string, std::vector<std::string>> targets;
-                for (const auto &member : members)
-                        targets[member] = {};
-
-                olm::send_encrypted_to_device_messages(targets, megolm_payload);
-
-                try {
-                        mtx::events::EncryptedEvent<mtx::events::msg::Encrypted> event;
-                        event.content =
-                          olm::encrypt_group_message(room_id, http::client()->device_id(), doc);
-                        event.event_id         = msg.event_id;
-                        event.room_id          = room_id;
-                        event.sender           = http::client()->user_id().to_string();
-                        event.type             = mtx::events::EventType::RoomEncrypted;
-                        event.origin_server_ts = QDateTime::currentMSecsSinceEpoch();
-
-                        emit this->addPendingMessageToStore(event);
-                } catch (const lmdb::error &e) {
-                        nhlog::db()->critical("failed to save megolm outbound session: {}",
-                                              e.what());
-                        emit ChatPage::instance()->showNotification(
-                          tr("Failed to encrypt event, sending aborted!"));
-                }
 
                 // TODO: Let the user know about the errors.
         } catch (const lmdb::error &e) {