From fc149bab3cc6388895588985b352c1bd2413152b Mon Sep 17 00:00:00 2001
From: Christoph Wurst <christoph@winzerhof-wurst.at>
Date: Tue, 31 Jul 2018 06:43:44 +0200
Subject: [PATCH] Fix duplicate inserts in the 2fa provider registry DAO

Signed-off-by: Christoph Wurst <christoph@winzerhof-wurst.at>
---
 .../Db/ProviderUserAssignmentDao.php          | 28 ++++++++++++-------
 .../Db/ProviderUserAssignmentDaoTest.php      | 19 +++++++++++++
 2 files changed, 37 insertions(+), 10 deletions(-)

diff --git a/lib/private/Authentication/TwoFactorAuth/Db/ProviderUserAssignmentDao.php b/lib/private/Authentication/TwoFactorAuth/Db/ProviderUserAssignmentDao.php
index 5b3f284943..0e9cd8045c 100644
--- a/lib/private/Authentication/TwoFactorAuth/Db/ProviderUserAssignmentDao.php
+++ b/lib/private/Authentication/TwoFactorAuth/Db/ProviderUserAssignmentDao.php
@@ -1,6 +1,6 @@
 <?php
 
-declare(strict_types = 1);
+declare(strict_types=1);
 
 /**
  * @copyright 2018 Christoph Wurst <christoph@winzerhof-wurst.at>
@@ -59,7 +59,7 @@ class ProviderUserAssignmentDao {
 		$result = $query->execute();
 		$providers = [];
 		foreach ($result->fetchAll() as $row) {
-			$providers[$row['provider_id']] = 1 === (int) $row['enabled'];
+			$providers[$row['provider_id']] = 1 === (int)$row['enabled'];
 		}
 		$result->closeCursor();
 
@@ -72,15 +72,23 @@ class ProviderUserAssignmentDao {
 	public function persist(string $providerId, string $uid, int $enabled) {
 		$qb = $this->conn->getQueryBuilder();
 
-		// TODO: concurrency? What if (providerId, uid) private key is inserted
-		//       twice at the same time?
-		$query = $qb->insert(self::TABLE_NAME)->values([
-			'provider_id' => $qb->createNamedParameter($providerId),
-			'uid' => $qb->createNamedParameter($uid),
-			'enabled' => $qb->createNamedParameter($enabled, IQueryBuilder::PARAM_INT),
-		]);
+		// First, try to update an existing entry
+		$updateQuery = $qb->update(self::TABLE_NAME)
+			->set('enabled', $qb->createNamedParameter($enabled))
+			->where($qb->expr()->eq('provider_id', $qb->createNamedParameter($providerId)))
+			->andWhere($qb->expr()->eq('uid', $qb->createNamedParameter($uid)));
+		$updatedRows = $updateQuery->execute();
 
-		$query->execute();
+		// If this (providerId, UID) key tuple is new, we have to insert it
+		if (0 === (int)$updatedRows) {
+			$insertQuery = $qb->insert(self::TABLE_NAME)->values([
+				'provider_id' => $qb->createNamedParameter($providerId),
+				'uid' => $qb->createNamedParameter($uid),
+				'enabled' => $qb->createNamedParameter($enabled, IQueryBuilder::PARAM_INT),
+			]);
+
+			$insertQuery->execute();
+		}
 	}
 
 }
diff --git a/tests/lib/Authentication/TwoFactorAuth/Db/ProviderUserAssignmentDaoTest.php b/tests/lib/Authentication/TwoFactorAuth/Db/ProviderUserAssignmentDaoTest.php
index d1fed0c3a6..6323c2f0aa 100644
--- a/tests/lib/Authentication/TwoFactorAuth/Db/ProviderUserAssignmentDaoTest.php
+++ b/tests/lib/Authentication/TwoFactorAuth/Db/ProviderUserAssignmentDaoTest.php
@@ -93,4 +93,23 @@ class ProviderUserAssignmentDaoTest extends TestCase {
 		$this->assertCount(1, $data);
 	}
 
+	public function testPersistTwice() {
+		$qb = $this->dbConn->getQueryBuilder();
+
+		$this->dao->persist('twofactor_totp', 'user123', 0);
+		$this->dao->persist('twofactor_totp', 'user123', 1);
+
+		$q = $qb
+			->select('*')
+			->from(ProviderUserAssignmentDao::TABLE_NAME)
+			->where($qb->expr()->eq('provider_id', $qb->createNamedParameter('twofactor_totp')))
+			->andWhere($qb->expr()->eq('uid', $qb->createNamedParameter('user123')))
+			->andWhere($qb->expr()->eq('enabled', $qb->createNamedParameter(1)));
+		$res = $q->execute();
+		$data = $res->fetchAll();
+		$res->closeCursor();
+
+		$this->assertCount(1, $data);
+	}
+
 }