Merge "Fix disabling TOTP keys with scratch tokens" into REL1_42

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,