Prepare for MSC 3700

8 files changed, 109 insertions, 50 deletions

diff --git a/lib/connection.cpp b/lib/connection.cpp
index d7460f08..0bdc0489 100644
--- a/lib/connection.cpp
+++ b/lib/connection.cpp
@@ -220,7 +220,7 @@ public:
         }
         q->database()->saveOlmSession(senderKey, session->sessionId(), std::get<QByteArray>(pickleResult), QDateTime::currentDateTime());
     }
-    QString sessionDecryptPrekey(const QOlmMessage& message, const QString &senderKey, std::unique_ptr<QOlmAccount>& olmAccount)
+    std::pair<QString, QString> sessionDecryptPrekey(const QOlmMessage& message, const QString &senderKey, std::unique_ptr<QOlmAccount>& olmAccount)
     {
         Q_ASSERT(message.type() == QOlmMessage::PreKey);
         for(auto& session : olmSessions[senderKey]) {
@@ -230,7 +230,7 @@ public:
                 const auto result = session->decrypt(message);
                 if(std::holds_alternative<QString>(result)) {
                     q->database()->setOlmSessionLastReceived(QString(session->sessionId()), QDateTime::currentDateTime());
-                    return std::get<QString>(result);
+                    return { std::get<QString>(result), session->sessionId() };
                 } else {
                     qCDebug(E2EE) << "Failed to decrypt prekey message";
                     return {};
@@ -249,47 +249,53 @@ public:
             qWarning(E2EE) << "Failed to remove one time key for session" << newSession->sessionId();
         }
         const auto result = newSession->decrypt(message);
+        QString sessionId = newSession->sessionId();
         saveSession(newSession, senderKey);
         olmSessions[senderKey].push_back(std::move(newSession));
         if(std::holds_alternative<QString>(result)) {
-            return std::get<QString>(result);
+            return { std::get<QString>(result), sessionId };
         } else {
             qCDebug(E2EE) << "Failed to decrypt prekey message with new session";
             return {};
         }
     }
-    QString sessionDecryptGeneral(const QOlmMessage& message, const QString &senderKey)
+    std::pair<QString, QString> sessionDecryptGeneral(const QOlmMessage& message, const QString &senderKey)
     {
         Q_ASSERT(message.type() == QOlmMessage::General);
         for(auto& session : olmSessions[senderKey]) {
             const auto result = session->decrypt(message);
             if(std::holds_alternative<QString>(result)) {
                 q->database()->setOlmSessionLastReceived(QString(session->sessionId()), QDateTime::currentDateTime());
-                return std::get<QString>(result);
+                return { std::get<QString>(result), session->sessionId() };
             }
         }
         qCWarning(E2EE) << "Failed to decrypt message";
         return {};
     }
 
-    QString sessionDecryptMessage(
+    std::pair<QString, QString> sessionDecryptMessage(
         const QJsonObject& personalCipherObject, const QByteArray& senderKey, std::unique_ptr<QOlmAccount>& account)
     {
         QString decrypted;
+        QString olmSessionId;
         int type = personalCipherObject.value(TypeKeyL).toInt(-1);
         QByteArray body = personalCipherObject.value(BodyKeyL).toString().toLatin1();
         if (type == QOlmMessage::PreKey) {
             QOlmMessage preKeyMessage(body, QOlmMessage::PreKey);
-            decrypted = sessionDecryptPrekey(preKeyMessage, senderKey, account);
+            auto result = sessionDecryptPrekey(preKeyMessage, senderKey, account);
+            decrypted = result.first;
+            olmSessionId = result.second;
         } else if (type == QOlmMessage::General) {
             QOlmMessage message(body, QOlmMessage::General);
-            decrypted = sessionDecryptGeneral(message, senderKey);
+            auto result = sessionDecryptGeneral(message, senderKey);
+            decrypted = result.first;
+            olmSessionId = result.second;
         }
-        return decrypted;
+        return { decrypted, olmSessionId };
     }
 #endif
 
-    EventPtr sessionDecryptMessage(const EncryptedEvent& encryptedEvent)
+    std::pair<EventPtr, QString> sessionDecryptMessage(const EncryptedEvent& encryptedEvent)
     {
 #ifndef Quotient_E2EE_ENABLED
         qCWarning(E2EE) << "End-to-end encryption (E2EE) support is turned off.";
@@ -305,7 +311,7 @@ public:
             qCDebug(E2EE) << "Encrypted event is not for the current device";
             return {};
         }
-        const auto decrypted = sessionDecryptMessage(
+        const auto [decrypted, olmSessionId] = sessionDecryptMessage(
             personalCipherObject, encryptedEvent.senderKey().toLatin1(), olmAccount);
         if (decrypted.isEmpty()) {
             qCDebug(E2EE) << "Problem with new session from senderKey:"
@@ -356,7 +362,7 @@ public:
             return {};
         }
 
-        return std::move(decryptedEvent);
+        return { std::move(decryptedEvent), olmSessionId };
 #endif // Quotient_E2EE_ENABLED
     }
 #ifdef Quotient_E2EE_ENABLED
@@ -956,16 +962,16 @@ void Connection::Private::consumeToDeviceEvents(Events&& toDeviceEvents)
 
 void Connection::Private::handleEncryptedToDeviceEvent(const EncryptedEvent& event)
 {
-    const auto decryptedEvent = sessionDecryptMessage(event);
+    const auto [decryptedEvent, olmSessionId] = sessionDecryptMessage(event);
         if(!decryptedEvent) {
             qCWarning(E2EE) << "Failed to decrypt event" << event.id();
             return;
         }
 
         switchOnType(*decryptedEvent,
-            [this, senderKey = event.senderKey()](const RoomKeyEvent& roomKeyEvent) {
+            [this, senderKey = event.senderKey(), &event, olmSessionId](const RoomKeyEvent& roomKeyEvent) {
                 if (auto* detectedRoom = q->room(roomKeyEvent.roomId())) {
-                    detectedRoom->handleRoomKeyEvent(roomKeyEvent, senderKey);
+                    detectedRoom->handleRoomKeyEvent(roomKeyEvent, event.senderId(), olmSessionId);
                 } else {
                     qCDebug(E2EE) << "Encrypted event room id" << roomKeyEvent.roomId()
                         << "is not found at the connection" << q->objectName();
@@ -2192,14 +2198,14 @@ Database* Connection::database()
     return d->database;
 }
 
-UnorderedMap<std::pair<QString, QString>, QOlmInboundGroupSessionPtr> Connection::loadRoomMegolmSessions(Room* room)
+UnorderedMap<QString, QOlmInboundGroupSessionPtr> Connection::loadRoomMegolmSessions(Room* room)
 {
     return database()->loadMegolmSessions(room->id(), picklingMode());
 }
 
-void Connection::saveMegolmSession(Room* room, const QString& senderKey, QOlmInboundGroupSession* session, const QString& ed25519Key)
+void Connection::saveMegolmSession(Room* room, QOlmInboundGroupSession* session)
 {
-    database()->saveMegolmSession(room->id(), senderKey, session->sessionId(), ed25519Key, session->pickle(picklingMode()));
+    database()->saveMegolmSession(room->id(), session->sessionId(), session->pickle(picklingMode()), session->senderId(), session->olmSessionId());
 }
 
 QStringList Connection::devicesForUser(User* user) const
diff --git a/lib/connection.h b/lib/connection.h
index 9f23902b..29731593 100644
--- a/lib/connection.h
+++ b/lib/connection.h
@@ -317,8 +317,8 @@ public:
 #ifdef Quotient_E2EE_ENABLED
     QOlmAccount* olmAccount() const;
     Database* database();
-    UnorderedMap<std::pair<QString, QString>, QOlmInboundGroupSessionPtr> loadRoomMegolmSessions(Room* room);
-    void saveMegolmSession(Room* room, const QString& senderKey, QOlmInboundGroupSession* session, const QString& ed25519Key);
+    UnorderedMap<QString, QOlmInboundGroupSessionPtr> loadRoomMegolmSessions(Room* room);
+    void saveMegolmSession(Room* room, QOlmInboundGroupSession* session);
 #endif // Quotient_E2EE_ENABLED
     Q_INVOKABLE Quotient::SyncJob* syncJob() const;
     Q_INVOKABLE int millisToReconnect() const;
diff --git a/lib/database.cpp b/lib/database.cpp
index d719d027..3189b3f1 100644
--- a/lib/database.cpp
+++ b/lib/database.cpp
@@ -29,6 +29,7 @@ Database::Database(const QString& matrixId, const QString& deviceId, QObject* pa
     switch(version()) {
         case 0: migrateTo1();
         case 1: migrateTo2();
+        case 2: migrateTo3();
     }
 }
 
@@ -112,6 +113,20 @@ void Database::migrateTo2()
     commit();
 }
 
+void Database::migrateTo3()
+{
+    qCDebug(DATABASE) << "Migrating database to version 3";
+    transaction();
+
+    execute(QStringLiteral("CREATE TABLE inbound_megolm_sessions_temp AS SELECT roomId, sessionId, pickle FROM inbound_megolm_sessions;"));
+    execute(QStringLiteral("DROP TABLE inbound_megolm_sessions;"));
+    execute(QStringLiteral("ALTER TABLE inbound_megolm_sessions_temp RENAME TO inbound_megolm_sessions;"));
+    execute(QStringLiteral("ALTER TABLE inbound_megolm_sessions ADD olmSessionId TEXT;"));
+    execute(QStringLiteral("ALTER TABLE inbound_megolm_sessions ADD senderId TEXT;"));
+    execute(QStringLiteral("PRAGMA user_version = 3;"));
+    commit();
+}
+
 QByteArray Database::accountPickle()
 {
     auto query = prepareQuery(QStringLiteral("SELECT pickle FROM accounts;"));
@@ -179,33 +194,36 @@ UnorderedMap<QString, std::vector<QOlmSessionPtr>> Database::loadOlmSessions(con
     return sessions;
 }
 
-UnorderedMap<std::pair<QString, QString>, QOlmInboundGroupSessionPtr> Database::loadMegolmSessions(const QString& roomId, const PicklingMode& picklingMode)
+UnorderedMap<QString, QOlmInboundGroupSessionPtr> Database::loadMegolmSessions(const QString& roomId, const PicklingMode& picklingMode)
 {
     auto query = prepareQuery(QStringLiteral("SELECT * FROM inbound_megolm_sessions WHERE roomId=:roomId;"));
     query.bindValue(":roomId", roomId);
     transaction();
     execute(query);
     commit();
-    UnorderedMap<std::pair<QString, QString>, QOlmInboundGroupSessionPtr> sessions;
+    UnorderedMap<QString, QOlmInboundGroupSessionPtr> sessions;
     while (query.next()) {
         auto session = QOlmInboundGroupSession::unpickle(query.value("pickle").toByteArray(), picklingMode);
         if (std::holds_alternative<QOlmError>(session)) {
             qCWarning(E2EE) << "Failed to unpickle megolm session";
             continue;
         }
-        sessions[{query.value("senderKey").toString(), query.value("sessionId").toString()}] = std::move(std::get<QOlmInboundGroupSessionPtr>(session));
+
+        sessions[query.value("sessionId").toString()] = std::move(std::get<QOlmInboundGroupSessionPtr>(session));
+        sessions[query.value("sessionId").toString()]->setOlmSessionId(query.value("olmSessionId").toString());
+        sessions[query.value("sessionId").toString()]->setSenderId(query.value("senderId").toString());
     }
     return sessions;
 }
 
-void Database::saveMegolmSession(const QString& roomId, const QString& senderKey, const QString& sessionId, const QString& ed25519Key, const QByteArray& pickle)
+void Database::saveMegolmSession(const QString& roomId, const QString& sessionId, const QByteArray& pickle, const QString& senderId, const QString& olmSessionId)
 {
-    auto query = prepareQuery(QStringLiteral("INSERT INTO inbound_megolm_sessions(roomId, senderKey, sessionId, ed25519Key, pickle) VALUES(:roomId, :senderKey, :sessionId, :ed25519Key, :pickle);"));
+    auto query = prepareQuery(QStringLiteral("INSERT INTO inbound_megolm_sessions(roomId, sessionId, pickle, senderId, olmSessionId) VALUES(:roomId, :sessionId, :pickle, :senderId, :olmSessionId);"));
     query.bindValue(":roomId", roomId);
-    query.bindValue(":senderKey", senderKey);
     query.bindValue(":sessionId", sessionId);
-    query.bindValue(":ed25519Key", ed25519Key);
     query.bindValue(":pickle", pickle);
+    query.bindValue(":senderId", senderId);
+    query.bindValue(":olmSessionId", olmSessionId);
     transaction();
     execute(query);
     commit();
diff --git a/lib/database.h b/lib/database.h
index cf241dbc..08fe49f3 100644
--- a/lib/database.h
+++ b/lib/database.h
@@ -8,7 +8,6 @@
 #include <QtCore/QVector>
 
 #include "e2ee/e2ee.h"
-
 namespace Quotient {
 class QUOTIENT_API Database : public QObject
 {
@@ -29,8 +28,8 @@ public:
     void clear();
     void saveOlmSession(const QString& senderKey, const QString& sessionId, const QByteArray &pickle, const QDateTime& timestamp);
     UnorderedMap<QString, std::vector<QOlmSessionPtr>> loadOlmSessions(const PicklingMode& picklingMode);
-    UnorderedMap<std::pair<QString, QString>, QOlmInboundGroupSessionPtr> loadMegolmSessions(const QString& roomId, const PicklingMode& picklingMode);
-    void saveMegolmSession(const QString& roomId, const QString& senderKey, const QString& sessionKey, const QString& ed25519Key, const QByteArray& pickle);
+    UnorderedMap<QString, QOlmInboundGroupSessionPtr> loadMegolmSessions(const QString& roomId, const PicklingMode& picklingMode);
+    void saveMegolmSession(const QString& roomId, const QString& sessionId, const QByteArray& pickle, const QString& senderId, const QString& olmSessionId);
     void addGroupSessionIndexRecord(const QString& roomId, const QString& sessionId, uint32_t index, const QString& eventId, qint64 ts);
     std::pair<QString, qint64> groupSessionIndexRecord(const QString& roomId, const QString& sessionId, qint64 index);
     void clearRoomData(const QString& roomId);
@@ -39,6 +38,8 @@ public:
 private:
     void migrateTo1();
     void migrateTo2();
+    void migrateTo3();
+
     QString m_matrixId;
 };
 }
diff --git a/lib/e2ee/qolminboundsession.cpp b/lib/e2ee/qolminboundsession.cpp
index 2e9cc716..60d871ef 100644
--- a/lib/e2ee/qolminboundsession.cpp
+++ b/lib/e2ee/qolminboundsession.cpp
@@ -149,3 +149,21 @@ bool QOlmInboundGroupSession::isVerified() const
 {
     return olm_inbound_group_session_is_verified(m_groupSession) != 0;
 }
+
+QString QOlmInboundGroupSession::olmSessionId() const
+{
+    return m_olmSessionId;
+}
+void QOlmInboundGroupSession::setOlmSessionId(const QString& olmSessionId)
+{
+    m_olmSessionId = olmSessionId;
+}
+
+QString QOlmInboundGroupSession::senderId() const
+{
+    return m_senderId;
+}
+void QOlmInboundGroupSession::setSenderId(const QString& senderId)
+{
+    m_senderId = senderId;
+}
diff --git a/lib/e2ee/qolminboundsession.h b/lib/e2ee/qolminboundsession.h
index 437f753d..32112b97 100644
--- a/lib/e2ee/qolminboundsession.h
+++ b/lib/e2ee/qolminboundsession.h
@@ -41,9 +41,20 @@ public:
     QByteArray sessionId() const;
     bool isVerified() const;
 
+    //! The olm session that this session was received from.
+    //! Required to get the device this session is from.
+    QString olmSessionId() const;
+    void setOlmSessionId(const QString& setOlmSessionId);
+
+    //! The sender of this session.
+    QString senderId() const;
+    void setSenderId(const QString& senderId);
+
     QOlmInboundGroupSession(OlmInboundGroupSession* session);
 private:
     OlmInboundGroupSession* m_groupSession;
+    QString m_olmSessionId;
+    QString m_senderId;
 };
 
 using QOlmInboundGroupSessionPtr = std::unique_ptr<QOlmInboundGroupSession>;
diff --git a/lib/room.cpp b/lib/room.cpp
index 3ba6890a..041d88c8 100644
--- a/lib/room.cpp
+++ b/lib/room.cpp
@@ -337,27 +337,25 @@ public:
     bool isLocalUser(const User* u) const { return u == q->localUser(); }
 
 #ifdef Quotient_E2EE_ENABLED
-    // A map from (senderKey, sessionId) to InboundGroupSession
-    UnorderedMap<std::pair<QString, QString>, QOlmInboundGroupSessionPtr> groupSessions;
+    UnorderedMap<QString, QOlmInboundGroupSessionPtr> groupSessions;
 
-    bool addInboundGroupSession(QString senderKey, QString sessionId,
-                                QString sessionKey, QString ed25519Key)
+    bool addInboundGroupSession(QString sessionId, QString sessionKey, const QString& senderId, const QString& olmSessionId)
     {
-        if (groupSessions.find({senderKey, sessionId}) != groupSessions.end()) {
-            qCWarning(E2EE) << "Inbound Megolm session" << sessionId
-                          << "with senderKey" << senderKey << "already exists";
+        if (groupSessions.find(sessionId) != groupSessions.end()) {
+            qCWarning(E2EE) << "Inbound Megolm session" << sessionId << "already exists";
             return false;
         }
 
         auto megolmSession = QOlmInboundGroupSession::create(sessionKey.toLatin1());
         if (megolmSession->sessionId() != sessionId) {
-            qCWarning(E2EE) << "Session ID mismatch in m.room_key event sent "
-                             "from sender with key" << senderKey;
+            qCWarning(E2EE) << "Session ID mismatch in m.room_key event";
             return false;
         }
+        megolmSession->setSenderId(senderId);
+        megolmSession->setOlmSessionId(olmSessionId);
         qCWarning(E2EE) << "Adding inbound session";
-        connection->saveMegolmSession(q, senderKey, megolmSession.get(), ed25519Key);
-        groupSessions[{senderKey, sessionId}] = std::move(megolmSession);
+        connection->saveMegolmSession(q, megolmSession.get());
+        groupSessions[sessionId] = std::move(megolmSession);
         return true;
     }
 
@@ -365,9 +363,10 @@ public:
                                        const QString& senderKey,
                                        const QString& sessionId,
                                        const QString& eventId,
-                                       QDateTime timestamp)
+                                       QDateTime timestamp,
+                                       const QString& senderId)
     {
-        auto groupSessionIt = groupSessions.find({ senderKey, sessionId });
+        auto groupSessionIt = groupSessions.find(sessionId);
         if (groupSessionIt == groupSessions.end()) {
             // qCWarning(E2EE) << "Unable to decrypt event" << eventId
             //               << "The sender's device has not sent us the keys for "
@@ -375,6 +374,10 @@ public:
             return QString();
         }
         auto& senderSession = groupSessionIt->second;
+        if (senderSession->senderId() != senderId) {
+            qCWarning(E2EE) << "Sender from event does not match sender from session";
+            return {};
+        }
         auto decryptResult = senderSession->decrypt(cipher);
         if(std::holds_alternative<QOlmError>(decryptResult)) {
             qCWarning(E2EE) << "Unable to decrypt event" << eventId
@@ -1482,9 +1485,9 @@ RoomEventPtr Room::decryptMessage(const EncryptedEvent& encryptedEvent)
     QString decrypted = d->groupSessionDecryptMessage(
         encryptedEvent.ciphertext(), encryptedEvent.senderKey(),
         encryptedEvent.sessionId(), encryptedEvent.id(),
-        encryptedEvent.originTimestamp());
+        encryptedEvent.originTimestamp(), encryptedEvent.senderId());
     if (decrypted.isEmpty()) {
-        // qCWarning(E2EE) << "Encrypted message is empty";
+        qCWarning(E2EE) << "Encrypted message is empty";
         return {};
     }
     auto decryptedEvent = encryptedEvent.createDecrypted(decrypted);
@@ -1497,19 +1500,21 @@ RoomEventPtr Room::decryptMessage(const EncryptedEvent& encryptedEvent)
 }
 
 void Room::handleRoomKeyEvent(const RoomKeyEvent& roomKeyEvent,
-                              const QString& senderKey)
+                              const QString& senderId,
+                              const QString& olmSessionId)
 {
 #ifndef Quotient_E2EE_ENABLED
     Q_UNUSED(roomKeyEvent)
-    Q_UNUSED(senderKey)
+    Q_UNUSED(senderId)
+    Q_UNUSED(olmSessionId)
     qCWarning(E2EE) << "End-to-end encryption (E2EE) support is turned off.";
 #else // Quotient_E2EE_ENABLED
     if (roomKeyEvent.algorithm() != MegolmV1AesSha2AlgoKey) {
         qCWarning(E2EE) << "Ignoring unsupported algorithm"
                         << roomKeyEvent.algorithm() << "in m.room_key event";
     }
-    if (d->addInboundGroupSession(senderKey, roomKeyEvent.sessionId(),
-                                  roomKeyEvent.sessionKey(), roomKeyEvent.fullJson()["keys"]["ed25519"].toString())) {
+    if (d->addInboundGroupSession(roomKeyEvent.sessionId(),
+                                  roomKeyEvent.sessionKey(), senderId, olmSessionId)) {
         qCWarning(E2EE) << "added new inboundGroupSession:"
                       << d->groupSessions.size();
         auto undecryptedEvents = d->undecryptedEvents[roomKeyEvent.sessionId()];
diff --git a/lib/room.h b/lib/room.h
index 9f70d77a..6ba7feac 100644
--- a/lib/room.h
+++ b/lib/room.h
@@ -277,7 +277,7 @@ public:
     int timelineSize() const;
     bool usesEncryption() const;
     RoomEventPtr decryptMessage(const EncryptedEvent& encryptedEvent);
-    void handleRoomKeyEvent(const RoomKeyEvent& roomKeyEvent, const QString& senderKey);
+    void handleRoomKeyEvent(const RoomKeyEvent& roomKeyEvent, const QString& senderId, const QString& olmSessionId);
     int joinedCount() const;
     int invitedCount() const;
     int totalMemberCount() const;