Track SSH agent key ownership by the root group uuid

src/gui/DatabaseWidget.cpp

@@ -842,7 +842,7 @@ void DatabaseWidget::addToAgent()
     SSHAgent* agent = SSHAgent::instance();
     OpenSSHKey key;
     if (settings.toOpenSSHKey(currentEntry, key, true)) {
-        if (!agent->addIdentity(key, settings, database()->uuid())) {
+        if (!agent->addIdentity(key, settings, database()->rootGroup()->uuid())) {
             m_messageWidget->showMessage(agent->errorString(), MessageWidget::Error);
         }
     } else {

src/gui/entry/EditEntryWidget.cpp

@@ -886,7 +886,7 @@ void EditEntryWidget::addKeyToAgent()
     KeeAgentSettings settings;
     toKeeAgentSettings(settings);
 
-    if (!sshAgent()->addIdentity(key, settings, m_db->uuid())) {
+    if (!sshAgent()->addIdentity(key, settings, m_db->rootGroup()->uuid())) {
         showMessage(sshAgent()->errorString(), MessageWidget::Error);
         return;
     }

src/sshagent/SSHAgent.cpp

@@ -264,7 +264,8 @@ bool SSHAgent::sendMessagePageant(const QByteArray& in, QByteArray& out)
  *
  * @param key identity / key to add
  * @param settings constraints (lifetime, confirm), remove-on-lock
- * @param databaseUuid database that owns the key for remove-on-lock
+ * @param databaseUuid database that owns the key for remove-on-lock: its root group uuid,
+ *                     which is stored in the file and survives a reload (Database::uuid() does not)
  * @return true on success
  */
 bool SSHAgent::addIdentity(OpenSSHKey& key, const KeeAgentSettings& settings, const QUuid& databaseUuid)
@@ -535,7 +536,7 @@ void SSHAgent::databaseLocked(const QSharedPointer<Database>& db)
 
     auto it = m_addedKeys.begin();
     while (it != m_addedKeys.end()) {
-        if (it.value().first != db->uuid()) {
+        if (it.value().first != db->rootGroup()->uuid()) {
             ++it;
             continue;
         }
@@ -578,7 +579,7 @@ void SSHAgent::databaseUnlocked(const QSharedPointer<Database>& db)
 
         // Add key to agent; ignore errors if we have previously added the key
         bool known_key = m_addedKeys.contains(key);
-        if (!addIdentity(key, settings, db->uuid()) && !known_key) {
+        if (!addIdentity(key, settings, db->rootGroup()->uuid()) && !known_key) {
             emit error(m_error);
         }
     }

tests/TestSSHAgent.cpp

@@ -18,12 +18,15 @@
 #include "TestSSHAgent.h"
 #include "config-keepassx-tests.h"
 #include "core/Config.h"
+#include "core/Group.h"
 #include "crypto/Crypto.h"
+#include "keys/PasswordKey.h"
 #include "sshagent/KeeAgentSettings.h"
 #include "sshagent/OpenSSHKeyGen.h"
 #include "sshagent/SSHAgent.h"
 
 #include <QElapsedTimer>
+#include <QScopeGuard>
 #include <QTest>
 
 QTEST_GUILESS_MAIN(TestSSHAgent)
@@ -84,9 +87,9 @@ void TestSSHAgent::initTestCase()
                                       "MEBQY=\n"
                                       "-----END OPENSSH PRIVATE KEY-----\n");
 
-    const QByteArray keyData = keyString.toLatin1();
+    m_keyData = keyString.toLatin1();
 
-    QVERIFY(m_key.parsePKCS1PEM(keyData));
+    QVERIFY(m_key.parsePKCS1PEM(m_keyData));
 }
 
 void TestSSHAgent::init()
@@ -164,6 +167,135 @@ void TestSSHAgent::testRemoveOnClose()
     QVERIFY(agent.checkIdentity(m_key, keyInAgent) && !keyInAgent);
 }
 
+// A database file holding m_key as an attachment, set to be added to the agent
+// on unlock and removed on lock.
+void TestSSHAgent::writeKeyDatabase(const QString& filePath)
+{
+    auto db = QSharedPointer<Database>::create();
+    auto key = QSharedPointer<CompositeKey>::create();
+    key->addKey(QSharedPointer<PasswordKey>::create("a"));
+    db->setKey(key);
+
+    auto* entry = new Entry();
+    entry->setUuid(QUuid::createUuid());
+    entry->setGroup(db->rootGroup());
+    entry->attachments()->set("id_ed25519", m_keyData);
+
+    KeeAgentSettings settings;
+    settings.setAllowUseOfSshKey(true);
+    settings.setAddAtDatabaseOpen(true);
+    settings.setRemoveAtDatabaseClose(true);
+    settings.setSelectedType("attachment");
+    settings.setAttachmentName("id_ed25519");
+    settings.toEntry(entry);
+
+    QString error;
+    QVERIFY2(db->saveAs(filePath, Database::Atomic, QString(), &error), qPrintable(error));
+}
+
+// Every load of the file builds a new Database object, as a reload from disk
+// (DatabaseWidget::reloadDatabaseFile) or opening a second copy of the file does.
+QSharedPointer<Database> TestSSHAgent::loadKeyDatabase(const QString& filePath)
+{
+    auto db = QSharedPointer<Database>::create();
+    auto key = QSharedPointer<CompositeKey>::create();
+    key->addKey(QSharedPointer<PasswordKey>::create("a"));
+    QString error;
+    if (!db->open(filePath, key, &error)) {
+        qWarning() << error;
+        return {};
+    }
+    return db;
+}
+
+void TestSSHAgent::testRemoveOnLockAfterReload()
+{
+    SSHAgent agent;
+    agent.setEnabled(true);
+    agent.setAuthSockOverride(m_agentSocketFileName);
+    QVERIFY(agent.isAgentRunning());
+    auto cleanup = qScopeGuard([&] { agent.removeIdentity(m_key); });
+
+    TemporaryFile file;
+    QVERIFY(file.open());
+    writeKeyDatabase(file.fileName());
+    bool keyInAgent;
+
+    auto db = loadKeyDatabase(file.fileName());
+    QVERIFY(db);
+    agent.databaseUnlocked(db);
+    QVERIFY(agent.checkIdentity(m_key, keyInAgent) && keyInAgent);
+
+    auto reloaded = loadKeyDatabase(file.fileName());
+    QVERIFY(reloaded);
+    db.reset();
+
+    agent.databaseLocked(reloaded);
+    QVERIFY(agent.checkIdentity(m_key, keyInAgent));
+    QVERIFY2(!keyInAgent, "key stayed in the agent after locking a reloaded database");
+}
+
+void TestSSHAgent::testReaddOnUnlockAfterReload()
+{
+    SSHAgent agent;
+    agent.setEnabled(true);
+    agent.setAuthSockOverride(m_agentSocketFileName);
+    QVERIFY(agent.isAgentRunning());
+    auto cleanup = qScopeGuard([&] { agent.removeIdentity(m_key); });
+
+    TemporaryFile file;
+    QVERIFY(file.open());
+    writeKeyDatabase(file.fileName());
+    bool keyInAgent;
+
+    auto db = loadKeyDatabase(file.fileName());
+    QVERIFY(db);
+    agent.databaseUnlocked(db);
+
+    auto reloaded = loadKeyDatabase(file.fileName());
+    QVERIFY(reloaded);
+    db.reset();
+    agent.databaseLocked(reloaded);
+
+    // The user clears the agent (ssh-add -D) and unlocks again.
+    agent.removeIdentity(m_key);
+    agent.databaseUnlocked(reloaded);
+    QCOMPARE(agent.errorString(), QString());
+    QVERIFY(agent.checkIdentity(m_key, keyInAgent));
+    QVERIFY2(keyInAgent, "key was not added back when unlocking a reloaded database");
+}
+
+// Two copies of one file open at once hold the same keys, so the agent treats
+// them as one owner: neither is refused, and locking either removes the keys.
+void TestSSHAgent::testTwoOpenCopiesShareKeys()
+{
+    SSHAgent agent;
+    agent.setEnabled(true);
+    agent.setAuthSockOverride(m_agentSocketFileName);
+    QVERIFY(agent.isAgentRunning());
+    auto cleanup = qScopeGuard([&] { agent.removeIdentity(m_key); });
+
+    TemporaryFile original;
+    QVERIFY(original.open());
+    writeKeyDatabase(original.fileName());
+    TemporaryFile copy;
+    QVERIFY(copy.copyFromFile(original.fileName()));
+    bool keyInAgent;
+
+    auto first = loadKeyDatabase(original.fileName());
+    auto second = loadKeyDatabase(copy.fileName());
+    QVERIFY(first && second);
+
+    agent.databaseUnlocked(first);
+    agent.databaseUnlocked(second);
+    QCOMPARE(agent.errorString(), QString());
+    QVERIFY(agent.checkIdentity(m_key, keyInAgent) && keyInAgent);
+
+    agent.databaseLocked(second);
+    QVERIFY(agent.checkIdentity(m_key, keyInAgent));
+    QVERIFY2(!keyInAgent, "key stayed in the agent after locking the second copy");
+}
+
 void TestSSHAgent::testLifetimeConstraint()
 {
     SSHAgent agent;

tests/TestSSHAgent.h

@@ -18,6 +18,7 @@
 #ifndef TESTSSHAGENT_H
 #define TESTSSHAGENT_H
 
+#include "core/Database.h"
 #include "sshagent/OpenSSHKey.h"
 #include "util/TemporaryFile.h"
 #include <QProcess>
@@ -33,6 +34,9 @@ private slots:
     void testConfiguration();
     void testIdentity();
     void testRemoveOnClose();
+    void testRemoveOnLockAfterReload();
+    void testReaddOnUnlockAfterReload();
+    void testTwoOpenCopiesShareKeys();
     void testLifetimeConstraint();
     void testConfirmConstraint();
     void testToOpenSSHKey();
@@ -45,6 +49,10 @@ private slots:
     QScopedPointer<TemporaryFile> m_agentSocketFile;
     QString m_agentSocketFileName;
     QProcess m_agentProcess;
+    void writeKeyDatabase(const QString& filePath);
+    QSharedPointer<Database> loadKeyDatabase(const QString& filePath);
+
+    QByteArray m_keyData;
     OpenSSHKey m_key;
     QUuid m_uuid;
 };