Keep log of where we send keys and send keys to new devices and users

4 files changed, 88 insertions, 19 deletions

diff --git a/lib/connection.cpp b/lib/connection.cpp
index b11ec731..2a1b39f9 100644
--- a/lib/connection.cpp
+++ b/lib/connection.cpp
@@ -2247,6 +2247,7 @@ bool Connection::hasOlmSession(User* user, const QString& deviceId) const
 QPair<QOlmMessage::Type, QByteArray> Connection::olmEncryptMessage(User* user, const QString& device, const QByteArray& message)
 {
     //TODO be smarter about choosing a session; see e2ee impl guide
+    //TODO do we need to save the olm session after sending a message?
     const auto& curveKey = curveKeyForUserDevice(user->id(), device);
     QOlmMessage::Type type = d->olmSessions[curveKey][0]->encryptMessageType();
     auto result = d->olmSessions[curveKey][0]->encrypt(message);
diff --git a/lib/database.cpp b/lib/database.cpp
index 8cb3a9d1..4a28fd4c 100644
--- a/lib/database.cpp
+++ b/lib/database.cpp
@@ -13,6 +13,9 @@
 #include "e2ee/e2ee.h"
 #include "e2ee/qolmsession.h"
 #include "e2ee/qolminboundsession.h"
+#include "connection.h"
+#include "user.h"
+#include "room.h"
 
 using namespace Quotient;
 Database::Database(const QString& matrixId, const QString& deviceId, QObject* parent)
@@ -91,6 +94,7 @@ void Database::migrateTo1()
     execute(QStringLiteral("CREATE TABLE tracked_users (matrixId TEXT);"));
     execute(QStringLiteral("CREATE TABLE outdated_users (matrixId TEXT);"));
     execute(QStringLiteral("CREATE TABLE tracked_devices (matrixId TEXT, deviceId TEXT, curveKeyId TEXT, curveKey TEXT, edKeyId TEXT, edKey TEXT);"));
+    execute(QStringLiteral("CREATE TABLE sent_megolm_sessions (roomId TEXT, userId TEXT, deviceId TEXT, identityKey TEXT, sessionId TEXT, i INTEGER);"));
 
     execute(QStringLiteral("PRAGMA user_version = 1;"));
     commit();
@@ -331,3 +335,43 @@ QOlmOutboundGroupSessionPtr Database::loadCurrentOutboundMegolmSession(const QSt
     }
     return nullptr;
 }
+
+void Database::setDevicesReceivedKey(const QString& roomId, QHash<User *, QStringList> devices, const QString& sessionId, int index)
+{
+    //TODO this better
+    auto connection = dynamic_cast<Connection *>(parent());
+    transaction();
+    for (const auto& user : devices.keys()) {
+        for (const auto& device : devices[user]) {
+            auto query = prepareQuery(QStringLiteral("INSERT INTO sent_megolm_sessions(roomId, userId, deviceId, identityKey, sessionId, i) VALUES(:roomId, :userId, :deviceId, :identityKey, :sessionId, :i);"));
+            query.bindValue(":roomId", roomId);
+            query.bindValue(":userId", user->id());
+            query.bindValue(":deviceId", device);
+            query.bindValue(":identityKey", connection->curveKeyForUserDevice(user->id(), device));
+            query.bindValue(":sessionId", sessionId);
+            query.bindValue(":i", index);
+            execute(query);
+        }
+    }
+    commit();
+}
+
+QHash<QString, QStringList> Database::devicesWithoutKey(Room* room, const QString &sessionId)
+{
+    auto connection = dynamic_cast<Connection *>(parent());
+    QHash<QString, QStringList> devices;
+    for (const auto& user : room->users()) {//TODO does this include invited & left?
+        devices[user->id()] = connection->devicesForUser(user);
+    }
+
+    auto query = prepareQuery(QStringLiteral("SELECT userId, deviceId FROM sent_megolm_sessions WHERE roomId=:roomId AND sessionId=:sessionId"));
+    query.bindValue(":roomId", room->id());
+    query.bindValue(":sessionId", sessionId);
+    transaction();
+    execute(query);
+    commit();
+    while (query.next()) {
+        devices[query.value("userId").toString()].removeAll(query.value("deviceId").toString());
+    }
+    return devices;
+}
diff --git a/lib/database.h b/lib/database.h
index 751ebd1d..30f2f203 100644
--- a/lib/database.h
+++ b/lib/database.h
@@ -7,9 +7,14 @@
 #include <QtSql/QSqlQuery>
 #include <QtCore/QVector>
 
+#include <QtCore/QHash>
+
 #include "e2ee/e2ee.h"
 
 namespace Quotient {
+class User;
+class Room;
+
 class QUOTIENT_API Database : public QObject
 {
     Q_OBJECT
@@ -38,6 +43,10 @@ public:
     QOlmOutboundGroupSessionPtr loadCurrentOutboundMegolmSession(const QString& roomId, const PicklingMode& picklingMode);
     void saveCurrentOutboundMegolmSession(const QString& roomId, const PicklingMode& picklingMode, const QOlmOutboundGroupSessionPtr& data);
 
+    // Returns a map User -> [Device] that have not received key yet
+    QHash<QString, QStringList> devicesWithoutKey(Room* room, const QString &sessionId);
+    void setDevicesReceivedKey(const QString& roomId, QHash<User *, QStringList> devices, const QString& sessionId, int index);
+
 private:
     void migrateTo1();
     void migrateTo2();
diff --git a/lib/room.cpp b/lib/room.cpp
index 9e2bd7dd..7db9f8e9 100644
--- a/lib/room.cpp
+++ b/lib/room.cpp
@@ -447,6 +447,13 @@ public:
         qCDebug(E2EE) << "Creating new outbound megolm session for room " << q->id();
         currentOutboundMegolmSession = QOlmOutboundGroupSession::create();
         connection->saveCurrentOutboundMegolmSession(q, currentOutboundMegolmSession);
+
+        const auto sessionKey = currentOutboundMegolmSession->sessionKey();
+        if(std::holds_alternative<QOlmError>(sessionKey)) {
+            qCWarning(E2EE) << "Session error";
+            //TODO something
+        }
+        addInboundGroupSession(q->connection()->olmAccount()->identityKeys().curve25519, currentOutboundMegolmSession->sessionId(), std::get<QByteArray>(sessionKey), QString(connection->olmAccount()->identityKeys().ed25519));
     }
 
     std::unique_ptr<EncryptedEvent> payloadForUserDevice(User* user, const QString& device, const QByteArray& sessionId, const QByteArray& sessionKey)
@@ -472,13 +479,23 @@ public:
         return makeEvent<EncryptedEvent>(encrypted, connection->olmAccount()->identityKeys().curve25519);
     }
 
-    void sendRoomKeyToDevices(const QByteArray& sessionId, const QByteArray& sessionKey)
+    QHash<User*, QStringList> getDevicesWithoutKey() const
+    {
+        QHash<User*, QStringList> devices;
+        auto rawDevices = q->connection()->database()->devicesWithoutKey(q, QString(currentOutboundMegolmSession->sessionId()));
+        for (const auto& user : rawDevices.keys()) {
+            devices[q->connection()->user(user)] = rawDevices[user];
+        }
+        return devices;
+    }
+
+    void sendRoomKeyToDevices(const QByteArray& sessionId, const QByteArray& sessionKey, const QHash<User*, QStringList> devices, int index)
     {
         qCDebug(E2EE) << "Sending room key to devices" << sessionId, sessionKey.toHex();
         QHash<QString, QHash<QString, QString>> hash;
-        for (const auto& user : q->users()) {
+        for (const auto& user : devices.keys()) {
             QHash<QString, QString> u;
-            for(const auto &device : connection->devicesForUser(user)) {
+            for(const auto &device : devices[user]) {
                 if (!connection->hasOlmSession(user, device)) {
                     u[device] = "signed_curve25519"_ls;
                     qCDebug(E2EE) << "Adding" << user << device << "to keys to claim";
@@ -489,30 +506,28 @@ public:
             }
         }
         auto job = connection->callApi<ClaimKeysJob>(hash);
-        connect(job, &BaseJob::success, q, [job, this, sessionId, sessionKey](){
+        connect(job, &BaseJob::success, q, [job, this, sessionId, sessionKey, devices, index](){
             Connection::UsersToDevicesToEvents usersToDevicesToEvents;
             const auto data = job->jsonData();
-            for(const auto &user : q->users()) {
-                for(const auto &device : connection->devicesForUser(user)) {
+            for(const auto &user : devices.keys()) {
+                for(const auto &device : devices[user]) {
                     const auto recipientCurveKey = connection->curveKeyForUserDevice(user->id(), device);
                     if (!connection->hasOlmSession(user, device)) {
                         qCDebug(E2EE) << "Creating a new session for" << user << device;
-                        if(data["one_time_keys"].toObject()[user->id()].toObject()[device].toObject().isEmpty()) {
+                        if(data["one_time_keys"][user->id()][device].toObject().isEmpty()) {
                             qWarning() << "No one time key for" << user << device;
                             continue;
                         }
-                        auto keyId = data["one_time_keys"].toObject()[user->id()].toObject()[device].toObject().keys()[0];
-                        auto oneTimeKey = data["one_time_keys"].toObject()[user->id()].toObject()[device].toObject()[keyId].toObject()["key"].toString();
-                        auto signature = data["one_time_keys"].toObject()[user->id()].toObject()[device].toObject()[keyId].toObject()["signatures"].toObject()[user->id()].toObject()[QStringLiteral("ed25519:") + device].toString().toLatin1();
-                        auto signedData = data["one_time_keys"].toObject()[user->id()].toObject()[device].toObject()[keyId].toObject();
+                        const auto keyId = data["one_time_keys"][user->id()][device].toObject().keys()[0];
+                        const auto oneTimeKey = data["one_time_keys"][user->id()][device][keyId]["key"].toString();
+                        const auto signature = data["one_time_keys"][user->id()][device][keyId]["signatures"][user->id()][QStringLiteral("ed25519:") + device].toString().toLatin1();
+                        auto signedData = data["one_time_keys"][user->id()][device][keyId].toObject();
                         signedData.remove("unsigned");
                         signedData.remove("signatures");
                         auto signatureMatch = QOlmUtility().ed25519Verify(connection->edKeyForUserDevice(user->id(), device).toLatin1(), QJsonDocument(signedData).toJson(QJsonDocument::Compact), signature);
                         if (std::holds_alternative<QOlmError>(signatureMatch)) {
                             //TODO i think there are more failed signature checks than expected. Investigate
-                            qDebug() << signedData;
                             qCWarning(E2EE) << "Failed to verify one-time-key signature for" << user->id() << device << ". Skipping this device.";
-                            //Q_ASSERT(false);
                             continue;
                         } else {
                         }
@@ -522,10 +537,11 @@ public:
                 }
             }
             connection->sendToDevices("m.room.encrypted", usersToDevicesToEvents);
+            connection->database()->setDevicesReceivedKey(q->id(), devices, sessionId, index);
         });
     }
 
-    void sendMegolmSession() {
+    void sendMegolmSession(const QHash<User *, QStringList>& devices) {
         // Save the session to this device
         const auto sessionId = currentOutboundMegolmSession->sessionId();
         const auto _sessionKey = currentOutboundMegolmSession->sessionKey();
@@ -536,11 +552,8 @@ public:
         const auto sessionKey = std::get<QByteArray>(_sessionKey);
         const auto senderKey = q->connection()->olmAccount()->identityKeys().curve25519;
 
-        // Send to key to ourself at this device
-        addInboundGroupSession(senderKey, sessionId, sessionKey);
-
         // Send the session to other people
-        sendRoomKeyToDevices(sessionId, sessionKey);
+        sendRoomKeyToDevices(sessionId, sessionKey, devices, currentOutboundMegolmSession->sessionMessageIndex());
     }
 
 #endif // Quotient_E2EE_ENABLED
@@ -2066,8 +2079,10 @@ QString Room::Private::sendEvent(RoomEventPtr&& event)
     if (q->usesEncryption()) {
         if (!hasValidMegolmSession() || shouldRotateMegolmSession()) {
             createMegolmSession();
-            sendMegolmSession();
         }
+        const auto devicesWithoutKey = getDevicesWithoutKey();
+        sendMegolmSession(devicesWithoutKey);
+
         //TODO check if this is necessary
         //TODO check if we increment the sent message count
         event->setRoomId(id);