diff --git a/src/Key/TOTPKey.php b/src/Key/TOTPKey.php index 3663c422..908b7ad8 100644 --- a/src/Key/TOTPKey.php +++ b/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; } } diff --git a/src/OATHUserRepository.php b/src/OATHUserRepository.php index dd1d550f..83d1bae3 100644 --- a/src/OATHUserRepository.php +++ b/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 diff --git a/tests/phpunit/integration/HTMLForm/TOTPDisableFormTest.php b/tests/phpunit/integration/HTMLForm/TOTPDisableFormTest.php new file mode 100644 index 00000000..2e22b63f --- /dev/null +++ b/tests/phpunit/integration/HTMLForm/TOTPDisableFormTest.php @@ -0,0 +1,104 @@ + + * @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() ); + } +} diff --git a/tests/phpunit/integration/OATHUserRepositoryTest.php b/tests/phpunit/integration/OATHUserRepositoryTest.php index 8a245776..7f6dff0d 100644 --- a/tests/phpunit/integration/OATHUserRepositoryTest.php +++ b/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 @@ -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,