SSH Agent: Track which database owns a key for remove-on-lock

Fixes keepassxreboot#4532

src/gui/DatabaseWidget.cpp

@@ -221,8 +221,8 @@ DatabaseWidget::DatabaseWidget(QSharedPointer<Database> db, QWidget* parent)
 
 #ifdef WITH_XC_SSHAGENT
     if (sshAgent()->isEnabled()) {
-        connect(this, SIGNAL(databaseLocked()), sshAgent(), SLOT(databaseModeChanged()));
-        connect(this, SIGNAL(databaseUnlocked()), sshAgent(), SLOT(databaseModeChanged()));
+        connect(this, SIGNAL(databaseLockRequested()), sshAgent(), SLOT(databaseLocked()));
+        connect(this, SIGNAL(databaseUnlocked()), sshAgent(), SLOT(databaseUnlocked()));
     }
 #endif
 
@@ -739,7 +739,7 @@ void DatabaseWidget::addToAgent()
 
     OpenSSHKey key;
     if (settings.toOpenSSHKey(currentEntry, key, true)) {
-        SSHAgent::instance()->addIdentity(key, settings);
+        SSHAgent::instance()->addIdentity(key, settings, database()->uuid());
     } else {
         m_messageWidget->showMessage(key.errorString(), MessageWidget::Error);
     }

src/gui/entry/EditEntryWidget.cpp

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

src/sshagent/SSHAgent.cpp

@@ -218,18 +218,22 @@ bool SSHAgent::sendMessagePageant(const QByteArray& in, QByteArray& out)
  * Add the identity to the SSH agent.
  *
  * @param key identity / key to add
- * @param lifetime time after which the key should expire
- * @param confirm ask for confirmation before adding the key
- * @param removeOnLock autoremove from agent when the Database is locked
+ * @param settings constraints (lifetime, confirm), remove-on-lock
+ * @param databaseUuid database that owns the key for remove-on-lock
  * @return true on success
  */
-bool SSHAgent::addIdentity(OpenSSHKey& key, KeeAgentSettings& settings)
+bool SSHAgent::addIdentity(OpenSSHKey& key, const KeeAgentSettings& settings, const QUuid& databaseUuid)
 {
     if (!isAgentRunning()) {
         m_error = tr("No agent running, cannot add identity.");
         return false;
     }
 
+    if (m_addedKeys.contains(key) && m_addedKeys[key].first != databaseUuid) {
+        m_error = tr("Key identity ownership conflict. Refusing to add.");
+        return false;
+    }
+
     QByteArray requestData;
     BinaryStream request(&requestData);
 
@@ -269,7 +273,7 @@ bool SSHAgent::addIdentity(OpenSSHKey& key, KeeAgentSettings& settings)
 
     OpenSSHKey keyCopy = key;
     keyCopy.clearPrivate();
-    m_addedKeys[keyCopy] = settings.removeAtDatabaseClose();
+    m_addedKeys[keyCopy] = qMakePair(databaseUuid, settings.removeAtDatabaseClose());
     return true;
 }
 
@@ -373,7 +377,7 @@ bool SSHAgent::listIdentities(QList<QSharedPointer<OpenSSHKey>>& list)
  * @param loaded is the key laoded
  * @return true on success
  */
-bool SSHAgent::checkIdentity(OpenSSHKey& key, bool& loaded)
+bool SSHAgent::checkIdentity(const OpenSSHKey& key, bool& loaded)
 {
     QList<QSharedPointer<OpenSSHKey>> list;
 
@@ -401,7 +405,7 @@ void SSHAgent::removeAllIdentities()
     auto it = m_addedKeys.begin();
     while (it != m_addedKeys.end()) {
         // Remove key if requested to remove on lock
-        if (it.value()) {
+        if (it.value().second) {
             OpenSSHKey key = it.key();
             removeIdentity(key);
         }
@@ -419,33 +423,43 @@ void SSHAgent::removeAllIdentities()
 void SSHAgent::setAutoRemoveOnLock(const OpenSSHKey& key, bool autoRemove)
 {
     if (m_addedKeys.contains(key)) {
-        m_addedKeys[key] = autoRemove;
+        m_addedKeys[key].second = autoRemove;
     }
 }
 
-void SSHAgent::databaseModeChanged()
+void SSHAgent::databaseLocked()
 {
     auto* widget = qobject_cast<DatabaseWidget*>(sender());
     if (!widget) {
         return;
     }
 
-    if (widget->isLocked()) {
-        auto it = m_addedKeys.begin();
-        while (it != m_addedKeys.end()) {
-            OpenSSHKey key = it.key();
-            if (it.value()) {
-                if (!removeIdentity(key)) {
-                    emit error(m_error);
-                }
-                it = m_addedKeys.erase(it);
-            } else {
-                // don't remove it yet
-                m_addedKeys[key] = false;
-                ++it;
+    QUuid databaseUuid = widget->database()->uuid();
+
+    auto it = m_addedKeys.begin();
+    while (it != m_addedKeys.end()) {
+        if (it.value().first != databaseUuid) {
+            ++it;
+            continue;
+        }
+        OpenSSHKey key = it.key();
+        if (it.value().second) {
+            if (!removeIdentity(key)) {
+                emit error(m_error);
             }
+            it = m_addedKeys.erase(it);
+        } else {
+            // don't remove it yet
+            m_addedKeys[key].second = false;
+            ++it;
         }
+    }
+}
 
+void SSHAgent::databaseUnlocked()
+{
+    auto* widget = qobject_cast<DatabaseWidget*>(sender());
+    if (!widget) {
         return;
     }
 
@@ -473,7 +487,7 @@ void SSHAgent::databaseModeChanged()
 
         // Add key to agent; ignore errors if we have previously added the key
         bool known_key = m_addedKeys.contains(key);
-        if (!addIdentity(key, settings) && !known_key) {
+        if (!addIdentity(key, settings, widget->database()->uuid()) && !known_key) {
             emit error(m_error);
         }
     }

src/sshagent/SSHAgent.h

@@ -47,9 +47,9 @@ class SSHAgent : public QObject
 
     const QString errorString() const;
     bool isAgentRunning() const;
-    bool addIdentity(OpenSSHKey& key, KeeAgentSettings& settings);
+    bool addIdentity(OpenSSHKey& key, const KeeAgentSettings& settings, const QUuid& databaseUuid);
     bool listIdentities(QList<QSharedPointer<OpenSSHKey>>& list);
-    bool checkIdentity(OpenSSHKey& key, bool& loaded);
+    bool checkIdentity(const OpenSSHKey& key, bool& loaded);
     bool removeIdentity(OpenSSHKey& key);
     void removeAllIdentities();
     void setAutoRemoveOnLock(const OpenSSHKey& key, bool autoRemove);
@@ -59,7 +59,8 @@ class SSHAgent : public QObject
     void enabledChanged(bool enabled);
 
 public slots:
-    void databaseModeChanged();
+    void databaseLocked();
+    void databaseUnlocked();
 
 private:
     const quint8 SSH_AGENT_FAILURE = 5;
@@ -81,7 +82,7 @@ public slots:
     const quint32 AGENT_COPYDATA_ID = 0x804e50ba;
 #endif
 
-    QHash<OpenSSHKey, bool> m_addedKeys;
+    QHash<OpenSSHKey, QPair<QUuid, bool>> m_addedKeys;
     QString m_error;
 };
 

tests/TestSSHAgent.cpp

@@ -119,9 +119,16 @@ void TestSSHAgent::testIdentity()
     bool keyInAgent;
 
     // test adding a key works
-    QVERIFY(agent.addIdentity(m_key, settings));
+    QVERIFY(agent.addIdentity(m_key, settings, m_uuid));
     QVERIFY(agent.checkIdentity(m_key, keyInAgent) && keyInAgent);
 
+    // test non-conflicting key ownership doesn't throw an error
+    QVERIFY(agent.addIdentity(m_key, settings, m_uuid));
+
+    // test conflicting key ownership throws an error
+    QUuid secondUuid("{11111111-1111-1111-1111-111111111111}");
+    QVERIFY(!agent.addIdentity(m_key, settings, secondUuid));
+
     // test removing a key works
     QVERIFY(agent.removeIdentity(m_key));
     QVERIFY(agent.checkIdentity(m_key, keyInAgent) && !keyInAgent);
@@ -139,7 +146,7 @@ void TestSSHAgent::testRemoveOnClose()
     bool keyInAgent;
 
     settings.setRemoveAtDatabaseClose(true);
-    QVERIFY(agent.addIdentity(m_key, settings));
+    QVERIFY(agent.addIdentity(m_key, settings, m_uuid));
     QVERIFY(agent.checkIdentity(m_key, keyInAgent) && keyInAgent);
     agent.setEnabled(false);
     QVERIFY(agent.checkIdentity(m_key, keyInAgent) && !keyInAgent);
@@ -160,7 +167,7 @@ void TestSSHAgent::testLifetimeConstraint()
     settings.setLifetimeConstraintDuration(2); // two seconds
 
     // identity should be in agent immediately after adding
-    QVERIFY(agent.addIdentity(m_key, settings));
+    QVERIFY(agent.addIdentity(m_key, settings, m_uuid));
     QVERIFY(agent.checkIdentity(m_key, keyInAgent) && keyInAgent);
 
     QElapsedTimer timer;
@@ -193,7 +200,7 @@ void TestSSHAgent::testConfirmConstraint()
 
     settings.setUseConfirmConstraintWhenAdding(true);
 
-    QVERIFY(agent.addIdentity(m_key, settings));
+    QVERIFY(agent.addIdentity(m_key, settings, m_uuid));
 
     // we can't test confirmation itself is working but we can test the agent accepts the key
     QVERIFY(agent.checkIdentity(m_key, keyInAgent) && keyInAgent);

tests/TestSSHAgent.h

@@ -41,6 +41,7 @@ private slots:
     QString m_agentSocketFileName;
     QProcess m_agentProcess;
     OpenSSHKey m_key;
+    QUuid m_uuid;
 };
 
 #endif // TESTSSHAGENT_H