From c6a621d31c3d34b5d15c67df8e8c396e6ca761e3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Taavi=20V=C3=A4=C3=A4n=C3=A4nen?= <hi@taavi.wtf>
Date: Tue, 31 Jan 2023 11:45:26 +0200
Subject: [PATCH] Make the TOTP disable form only remove that single key

Bug: T242031
Change-Id: Iad07292cc96537e8ebd72da65e8f1e922cba3eca
---
 src/HTMLForm/TOTPDisableForm.php              | 44 +++++++++------
 src/OATHUser.php                              |  4 +-
 src/OATHUserRepository.php                    | 55 +++++++++++++++++++
 .../integration/OATHUserRepositoryTest.php    |  3 +-
 4 files changed, 85 insertions(+), 21 deletions(-)

diff --git a/src/HTMLForm/TOTPDisableForm.php b/src/HTMLForm/TOTPDisableForm.php
index 9ee0b7b4..a235e003 100644
--- a/src/HTMLForm/TOTPDisableForm.php
+++ b/src/HTMLForm/TOTPDisableForm.php
@@ -2,7 +2,7 @@
 
 namespace MediaWiki\Extension\OATHAuth\HTMLForm;
 
-use MediaWiki\Extension\OATHAuth\Module\TOTP;
+use MediaWiki\Extension\OATHAuth\Key\TOTPKey;
 use MediaWiki\Logger\LoggerFactory;
 use Message;
 use MWException;
@@ -52,27 +52,35 @@ class TOTPDisableForm extends OATHAuthOOUIHTMLForm {
 			return [ 'oathauth-throttled', Message::durationParam( 60 ) ];
 		}
 
-		$module = $this->oathUser->getModule();
-		if (
-			( $module instanceof TOTP ) &&
-			!$module->verify( $this->oathUser, [ 'token' => $formData['token'] ] )
-		) {
-			LoggerFactory::getInstance( 'authentication' )->info(
-				'OATHAuth {user} failed to provide a correct token while disabling 2FA from {clientip}', [
-					'user' => $this->getUser()->getName(),
-					'clientip' => $this->getRequest()->getIP(),
-				]
+		foreach ( $this->oathUser->getKeys() as $key ) {
+			if ( !( $key instanceof TOTPKey ) ) {
+				continue;
+			}
+
+			if ( !$key->verify( [ 'token' => $formData['token'] ], $this->oathUser ) ) {
+				continue;
+			}
+
+			$this->oathRepo->removeKey(
+				$this->oathUser,
+				$key,
+				$this->getRequest()->getIP(),
+				true
 			);
 
-			// Increase rate limit counter for failed request
-			$this->getUser()->pingLimiter( 'badoath' );
-
-			return [ 'oathauth-failedtovalidateoath' ];
+			return true;
 		}
 
-		$this->oathUser->setKeys();
-		$this->oathRepo->removeAll( $this->oathUser, $this->getRequest()->getIP(), true );
+		LoggerFactory::getInstance( 'authentication' )->info(
+			'OATHAuth {user} failed to provide a correct token while disabling 2FA from {clientip}', [
+				'user' => $this->getUser()->getName(),
+				'clientip' => $this->getRequest()->getIP(),
+			]
+		);
 
-		return true;
+		// Increase rate limit counter for failed request
+		$this->getUser()->pingLimiter( 'badoath' );
+
+		return [ 'oathauth-failedtovalidateoath' ];
 	}
 }
diff --git a/src/OATHUser.php b/src/OATHUser.php
index 12ec67fe..9ac7bde8 100644
--- a/src/OATHUser.php
+++ b/src/OATHUser.php
@@ -81,9 +81,9 @@ class OATHUser {
 	/**
 	 * Get the key associated with this user.
 	 *
-	 * @return IAuthKey[]|array
+	 * @return IAuthKey[]
 	 */
-	public function getKeys() {
+	public function getKeys(): array {
 		return $this->keys;
 	}
 
diff --git a/src/OATHUserRepository.php b/src/OATHUserRepository.php
index 4100af08..dd1d550f 100644
--- a/src/OATHUserRepository.php
+++ b/src/OATHUserRepository.php
@@ -194,6 +194,61 @@ class OATHUserRepository implements LoggerAwareInterface {
 		return $key;
 	}
 
+	/**
+	 * @param OATHUser $user
+	 * @param IAuthKey $key
+	 * @param string $clientInfo
+	 * @param bool $self Whether they disabled it themselves
+	 */
+	public function removeKey( OATHUser $user, IAuthKey $key, string $clientInfo, bool $self ) {
+		$keyId = $key->getId();
+		if ( !$keyId ) {
+			throw new InvalidArgumentException( 'A non-persisted key cannot be removed' );
+		}
+
+		$userId = $this->centralIdLookupFactory->getLookup()
+			->centralIdFromLocalUser( $user->getUser() );
+		$this->dbProvider->getPrimaryDatabase( 'virtual-oathauth' )
+			->newDeleteQueryBuilder()
+			->deleteFrom( 'oathauth_devices' )
+			->where( [ 'oad_user' => $userId, 'oad_id' => $keyId ] )
+			->caller( __METHOD__ )
+			->execute();
+
+		// TODO: figure this out from the key itself
+		// After calling ->disable(), getModule() will return null so this
+		// has to be done before.
+		$keyType = $user->getModule()->getName();
+
+		// Remove the key from the user object
+		$user->setKeys(
+			array_values(
+				array_filter(
+					$user->getKeys(),
+					static function ( IAuthKey $key ) use ( $keyId ) {
+						return $key->getId() !== $keyId;
+					}
+				)
+			)
+		);
+
+		if ( !$user->getKeys() ) {
+			$user->setModule( null );
+		}
+
+		$userName = $user->getUser()->getName();
+		$this->cache->delete( $userName );
+
+		$this->logger->info( 'OATHAuth removed {oathtype} key {key} for {user} from {clientip}', [
+			'key' => $keyId,
+			'user' => $userName,
+			'clientip' => $clientInfo,
+			'oathtype' => $keyType,
+		] );
+
+		Manager::notifyDisabled( $user, $self );
+	}
+
 	/**
 	 * @param OATHUser $user
 	 * @param string $clientInfo
diff --git a/tests/phpunit/integration/OATHUserRepositoryTest.php b/tests/phpunit/integration/OATHUserRepositoryTest.php
index ea6ed7ec..8a245776 100644
--- a/tests/phpunit/integration/OATHUserRepositoryTest.php
+++ b/tests/phpunit/integration/OATHUserRepositoryTest.php
@@ -93,8 +93,9 @@ class OATHUserRepositoryTest extends MediaWikiIntegrationTestCase {
 		// Test looking it up again from the database
 		$this->assertArrayEquals( [ $key ], $repository->findByUser( $user )->getKeys() );
 
-		$repository->remove(
+		$repository->removeKey(
 			$oathUser,
+			$key,
 			'127.0.0.1',
 			true
 		);