Fix disabling TOTP keys with scratch tokens

The current implementation of OATHUserRepository::persist() causes every
key to get a new ID when it's saved. This, combined with ::removeKey()
which compares keys by ID, means that using recovery codes to disable
TOTP is broken since TOTPKey calls persist() to mark the code as saved
just before the key is deleted.

In this patch I've chosen to add a new ::updateKey() method instead of
fixing ::persist(). This is more in line with the other new APIs in
OATHUserRepository (namely ::createKey() and ::removeKey()), and is
something I've been planning to do eventually - this bug just made that
a bit more urgent. ::persist() should be dropped once WebAuthn has been
updated too.

Tests are also updated - OATHUserRepositoryTest now updates the key
before deleting it and there's a new TOTPDisableFormTest to test the
entire disabling process.

Bug: T363548
Change-Id: I86ddc8e5bfc9cf74c587ffdff523f559c5a3c08c
(cherry picked from commit 0dad2c7031)

src/Key/TOTPKey.php

@@ -189,11 +189,9 @@ class TOTPKey implements IAuthKey {
 					'clientip' => $clientIP,
 				] );
 
-				$userRepo = OATHAuthServices::getInstance()->getUserRepository();
-				// TODO: support for multiple keys
-				$user->setKeys( [ $this ] );
-				$userRepo->persist( $user, $clientIP );
-
+				OATHAuthServices::getInstance()
+					->getUserRepository()
+					->updateKey( $user, $this );
 				return true;
 			}
 		}

src/OATHUserRepository.php

@@ -194,6 +194,35 @@ class OATHUserRepository implements LoggerAwareInterface {
 		return $key;
 	}
 
+	/**
+	 * Saves an existing key in the database.
+	 *
+	 * @param OATHUser $user
+	 * @param IAuthKey $key
+	 * @return void
+	 */
+	public function updateKey( OATHUser $user, IAuthKey $key ): void {
+		$keyId = $key->getId();
+		if ( !$keyId ) {
+			throw new InvalidArgumentException( 'updateKey() can only be used with already existing keys' );
+		}
+
+		$userId = $this->centralIdLookupFactory->getLookup()
+			->centralIdFromLocalUser( $user->getUser() );
+
+		$dbw = $this->dbProvider->getPrimaryDatabase( 'virtual-oathauth' );
+		$dbw->newUpdateQueryBuilder()
+			->table( 'oathauth_devices' )
+			->set( [ 'oad_data' => FormatJson::encode( $key->jsonSerialize() ) ] )
+			->where( [ 'oad_user' => $userId, 'oad_id' => $keyId ] )
+			->execute();
+
+		$this->logger->info( 'OATHAuth key {keyId} updated for {user}', [
+			'keyId' => $keyId,
+			'user' => $user->getUser()->getName(),
+		] );
+	}
+
 	/**
 	 * @param OATHUser $user
 	 * @param IAuthKey $key

tests/phpunit/integration/HTMLForm/TOTPDisableFormTest.php

@@ -0,0 +1,104 @@
+<?php
+/**
+ * This program is free software; you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation; either version 2 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License along
+ * with this program; if not, write to the Free Software Foundation, Inc.,
+ * 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
+ * http://www.gnu.org/copyleft/gpl.html
+ *
+ * @file
+ */
+
+namespace MediaWiki\Extension\OATHAuth\Tests\Integration\HTMLForm;
+
+use Base32\Base32;
+use jakobo\HOTP\HOTP;
+use MediaWiki\Context\RequestContext;
+use MediaWiki\Extension\OATHAuth\HTMLForm\TOTPDisableForm;
+use MediaWiki\Extension\OATHAuth\Key\TOTPKey;
+use MediaWiki\Extension\OATHAuth\OATHAuthServices;
+use MediaWikiIntegrationTestCase;
+use Wikimedia\TestingAccessWrapper;
+
+/**
+ * @author Taavi Väänänen <hi@taavi.wtf>
+ * @group Database
+ * @coversDefaultClass \MediaWiki\Extension\OATHAuth\HTMLForm\TOTPDisableForm
+ */
+class TOTPDisableFormTest extends MediaWikiIntegrationTestCase {
+	/**
+	 * @return array
+	 * @phan-return array{0:TOTPDisableForm,1:TOTPKey,2:MediaWiki\Extension\OATHAuth\OATHUser}
+	 */
+	private function setupFormAndKey(): array {
+		$user = $this->getTestUser()->getUser();
+		$repository = OATHAuthServices::getInstance( $this->getServiceContainer() )
+			->getUserRepository();
+		$oathUser = $repository->findByUser( $user );
+		$module = OATHAuthServices::getInstance( $this->getServiceContainer() )
+			->getModuleRegistry()
+			->getModuleByKey( 'totp' );
+
+		$key = $repository->createKey(
+			$oathUser,
+			$module,
+			TOTPKey::newFromRandom()->jsonSerialize(),
+			'127.0.0.1'
+		);
+
+		$form = new TOTPDisableForm(
+			$oathUser,
+			$repository,
+			$module,
+			RequestContext::getMain(),
+		);
+
+		return [ $form, $key, $oathUser ];
+	}
+
+	/**
+	 * @covers ::onSubmit
+	 */
+	public function testSubmitInvalidCode(): void {
+		[ $form ] = $this->setupFormAndKey();
+		$this->assertEquals(
+			[ 'oathauth-failedtovalidateoath' ],
+			$form->onSubmit( [ 'token' => 'wrong' ] ),
+		);
+	}
+
+	/**
+	 * @covers ::onSubmit
+	 */
+	public function testSubmitCorrectToken(): void {
+		[ $form, $key, $oathUser ] = $this->setupFormAndKey();
+
+		$secret = TestingAccessWrapper::newFromObject( $key )->secret;
+		$correctToken = HOTP::generateByTime(
+			Base32::decode( $secret['secret'] ),
+			$secret['period'],
+		)->toHOTP( 6 );
+
+		$this->assertTrue( $form->onSubmit( [ 'token' => $correctToken ] ) );
+		$this->assertEquals( [], $oathUser->getKeys() );
+	}
+
+	/**
+	 * @covers ::onSubmit
+	 */
+	public function testSubmitScratchToken(): void {
+		[ $form, $key, $oathUser ] = $this->setupFormAndKey();
+
+		$this->assertTrue( $form->onSubmit( [ 'token' => $key->getScratchTokens()[0] ] ) );
+		$this->assertEquals( [], $oathUser->getKeys() );
+	}
+}

tests/phpunit/integration/OATHUserRepositoryTest.php

@@ -29,6 +29,7 @@ use MediaWiki\User\CentralId\CentralIdLookupFactory;
 use MediaWikiIntegrationTestCase;
 use Psr\Log\LoggerInterface;
 use Wikimedia\Rdbms\IConnectionProvider;
+use Wikimedia\TestingAccessWrapper;
 
 /**
  * @author Taavi Väänänen <hi@taavi.wtf>
@@ -39,6 +40,7 @@ class OATHUserRepositoryTest extends MediaWikiIntegrationTestCase {
 	/**
 	 * @covers ::findByUser
 	 * @covers ::createKey
+	 * @covers ::updateKey
 	 * @covers ::remove
 	 */
 	public function testLookupCreateRemoveKey(): void {
@@ -73,6 +75,7 @@ class OATHUserRepositoryTest extends MediaWikiIntegrationTestCase {
 		$this->assertEquals( [], $oathUser->getKeys() );
 		$this->assertNull( $oathUser->getModule() );
 
+		/** @var TOTPKey $key */
 		$key = $repository->createKey(
 			$oathUser,
 			$module,
@@ -93,6 +96,15 @@ class OATHUserRepositoryTest extends MediaWikiIntegrationTestCase {
 		// Test looking it up again from the database
 		$this->assertArrayEquals( [ $key ], $repository->findByUser( $user )->getKeys() );
 
+		// Use a scratch code, which causes the key to be updated.
+		TestingAccessWrapper::newFromObject( $key )->recoveryCodes = [ 'new scratch tokens' ];
+		$repository->updateKey( $oathUser, $key );
+
+		$this->assertEquals(
+			[ 'new scratch tokens' ],
+			$repository->findByUser( $user )->getKeys()[0]->getScratchTokens()
+		);
+
 		$repository->removeKey(
 			$oathUser,
 			$key,