Make the TOTP disable form only remove that single key

Bug: T242031
Change-Id: Iad07292cc96537e8ebd72da65e8f1e922cba3eca

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,11 +52,25 @@ class TOTPDisableForm extends OATHAuthOOUIHTMLForm {
 			return [ 'oathauth-throttled', Message::durationParam( 60 ) ];
 		}
 
-		$module = $this->oathUser->getModule();
+		foreach ( $this->oathUser->getKeys() as $key ) {
-		if (
+			if ( !( $key instanceof TOTPKey ) ) {
-			( $module instanceof TOTP ) &&
+				continue;
-			!$module->verify( $this->oathUser, [ 'token' => $formData['token'] ] )
+			}
-		) {
+
+			if ( !$key->verify( [ 'token' => $formData['token'] ], $this->oathUser ) ) {
+				continue;
+			}
+
+			$this->oathRepo->removeKey(
+				$this->oathUser,
+				$key,
+				$this->getRequest()->getIP(),
+				true
+			);
+
+			return true;
+		}
+
 		LoggerFactory::getInstance( 'authentication' )->info(
 			'OATHAuth {user} failed to provide a correct token while disabling 2FA from {clientip}', [
 				'user' => $this->getUser()->getName(),
@@ -69,10 +83,4 @@ class TOTPDisableForm extends OATHAuthOOUIHTMLForm {
 
 		return [ 'oathauth-failedtovalidateoath' ];
 	}
-
-		$this->oathUser->setKeys();
-		$this->oathRepo->removeAll( $this->oathUser, $this->getRequest()->getIP(), true );
-
-		return true;
-	}
 }

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;
 	}
 

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

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
 		);