Inject load balancers into some authentication providers

Change-Id: Ie2407cdebf1bf565b0db2f0a6bd0f5dec043a1b9
2021-05-05 17:02:51 +00:00 · 2021-05-05 17:02:51 +00:00 · e11939f1d0
commit e11939f1d0
parent 334dfff732
8 changed files with 116 additions and 62 deletions
--- a/includes/DefaultSettings.php
+++ b/includes/DefaultSettings.php
@ -5301,6 +5301,9 @@ $wgAuthManagerAutoConfig = [
 		// probably auto-insert themselves in the wrong place.
 		MediaWiki\Auth\TemporaryPasswordPrimaryAuthenticationProvider::class => [
 			'class' => MediaWiki\Auth\TemporaryPasswordPrimaryAuthenticationProvider::class,
+			'services' => [
+				'DBLoadBalancer',
+			],
 			'args' => [ [
 				// Fall through to LocalPasswordPrimaryAuthenticationProvider
 				'authoritative' => false,
@ -5309,6 +5312,9 @@ $wgAuthManagerAutoConfig = [
 		],
 		MediaWiki\Auth\LocalPasswordPrimaryAuthenticationProvider::class => [
 			'class' => MediaWiki\Auth\LocalPasswordPrimaryAuthenticationProvider::class,
+			'services' => [
+				'DBLoadBalancer',
+			],
 			'args' => [ [
 				// Last one should be authoritative, or else the user will get
 				// a less-than-helpful error message (something like "supplied
@ -5335,6 +5341,9 @@ $wgAuthManagerAutoConfig = [
 		// ],
 		MediaWiki\Auth\EmailNotificationSecondaryAuthenticationProvider::class => [
 			'class' => MediaWiki\Auth\EmailNotificationSecondaryAuthenticationProvider::class,
+			'services' => [
+				'DBLoadBalancer',
+			],
 			'sort' => 200,
 		],
 	],
--- a/includes/auth/EmailNotificationSecondaryAuthenticationProvider.php
+++ b/includes/auth/EmailNotificationSecondaryAuthenticationProvider.php
@ -2,6 +2,8 @@

 namespace MediaWiki\Auth;

+use Wikimedia\Rdbms\ILoadBalancer;
+
 /**
 * Handles email notification / email address confirmation for account creation.
 *
@ -14,15 +16,20 @@ class EmailNotificationSecondaryAuthenticationProvider
 	/** @var bool */
 	protected $sendConfirmationEmail;

+	/** @var ILoadBalancer */
+	private $loadBalancer;
+
 	/**
+	 * @param ILoadBalancer $loadBalancer
 	 * @param array $params
 	 *  - sendConfirmationEmail: (bool) send an email asking the user to confirm their email
 	 *    address after a successful registration
 	 */
-	public function __construct( $params = [] ) {
+	public function __construct( ILoadBalancer $loadBalancer, $params = [] ) {
 		if ( isset( $params['sendConfirmationEmail'] ) ) {
 			$this->sendConfirmationEmail = (bool)$params['sendConfirmationEmail'];
 		}
+		$this->loadBalancer = $loadBalancer;
 	}

 	protected function postInitSetup() {
@ -47,7 +54,7 @@ class EmailNotificationSecondaryAuthenticationProvider
 			&& !$this->manager->getAuthenticationSessionData( 'no-email' )
 		) {
 			// TODO show 'confirmemail_oncreate'/'confirmemail_sendfailed' message
-			wfGetDB( DB_PRIMARY )->onTransactionCommitOrIdle(
+			$this->loadBalancer->getConnectionRef( DB_PRIMARY )->onTransactionCommitOrIdle(
 				function () use ( $user ) {
 					$user = $user->getInstanceForUpdate();
 					$status = $user->sendConfirmationMail();
--- a/includes/auth/LocalPasswordPrimaryAuthenticationProvider.php
+++ b/includes/auth/LocalPasswordPrimaryAuthenticationProvider.php
@ -22,6 +22,7 @@
 namespace MediaWiki\Auth;

 use User;
+use Wikimedia\Rdbms\ILoadBalancer;

 /**
 * A primary authentication provider that uses the password field in the 'user' table.
@ -35,15 +36,20 @@ class LocalPasswordPrimaryAuthenticationProvider
 	/** @var bool If true, this instance is for legacy logins only. */
 	protected $loginOnly = false;

+	/** @var ILoadBalancer */
+	private $loadBalancer;
+
 	/**
+	 * @param ILoadBalancer $loadBalancer
 	 * @param array $params Settings
 	 *  - loginOnly: If true, the local passwords are for legacy logins only:
 	 *    the local password will be invalidated when authentication is changed
 	 *    and new users will not have a valid local password set.
 	 */
-	public function __construct( $params = [] ) {
+	public function __construct( ILoadBalancer $loadBalancer, $params = [] ) {
 		parent::__construct( $params );
 		$this->loginOnly = !empty( $params['loginOnly'] );
+		$this->loadBalancer = $loadBalancer;
 	}

 	/**
@ -95,7 +101,7 @@ class LocalPasswordPrimaryAuthenticationProvider
 			'user_id', 'user_password', 'user_password_expires',
 		];

-		$dbr = wfGetDB( DB_REPLICA );
+		$dbr = $this->loadBalancer->getConnectionRef( DB_REPLICA );
 		$row = $dbr->selectRow(
 			'user',
 			$fields,
@ -140,8 +146,8 @@ class LocalPasswordPrimaryAuthenticationProvider
 		if ( $this->getPasswordFactory()->needsUpdate( $pwhash ) ) {
 			$newHash = $this->getPasswordFactory()->newFromPlaintext( $req->password );
 			$fname = __METHOD__;
-			\DeferredUpdates::addCallableUpdate( static function () use ( $newHash, $oldRow, $fname ) {
-				$dbw = wfGetDB( DB_PRIMARY );
+			\DeferredUpdates::addCallableUpdate( function () use ( $newHash, $oldRow, $fname ) {
+				$dbw = $this->loadBalancer->getConnectionRef( DB_PRIMARY );
 				$dbw->update(
 					'user',
 					[ 'user_password' => $newHash->toString() ],
@ -166,7 +172,7 @@ class LocalPasswordPrimaryAuthenticationProvider
 			return false;
 		}

-		$dbr = wfGetDB( DB_REPLICA );
+		$dbr = $this->loadBalancer->getConnectionRef( DB_REPLICA );
 		$row = $dbr->selectRow(
 			'user',
 			[ 'user_password' ],
@ -193,7 +199,7 @@ class LocalPasswordPrimaryAuthenticationProvider
 		}

 		list( $db, $options ) = \DBAccessObjectUtils::getDBOptions( $flags );
-		return (bool)wfGetDB( $db )->selectField(
+		return (bool)$this->loadBalancer->getConnectionRef( $db )->selectField(
 			[ 'user' ],
 			'user_id',
 			[ 'user_name' => $username ],
@ -218,7 +224,7 @@ class LocalPasswordPrimaryAuthenticationProvider

 			$username = User::getCanonicalName( $req->username, 'usable' );
 			if ( $username !== false ) {
-				$row = wfGetDB( DB_PRIMARY )->selectRow(
+				$row = $this->loadBalancer->getConnectionRef( DB_PRIMARY )->selectRow(
 					'user',
 					[ 'user_id' ],
 					[ 'user_name' => $username ],
@ -260,7 +266,7 @@ class LocalPasswordPrimaryAuthenticationProvider
 		}

 		if ( $pwhash ) {
-			$dbw = wfGetDB( DB_PRIMARY );
+			$dbw = $this->loadBalancer->getConnectionRef( DB_PRIMARY );
 			$dbw->update(
 				'user',
 				[
--- a/includes/auth/TemporaryPasswordPrimaryAuthenticationProvider.php
+++ b/includes/auth/TemporaryPasswordPrimaryAuthenticationProvider.php
@ -25,6 +25,7 @@ use MediaWiki\MediaWikiServices;
 use SpecialPage;
 use User;
 use Wikimedia\IPUtils;
+use Wikimedia\Rdbms\ILoadBalancer;

 /**
 * A primary authentication provider that uses the temporary password field in
@ -52,14 +53,18 @@ class TemporaryPasswordPrimaryAuthenticationProvider
 	/** @var bool */
 	protected $allowRequiringEmail = null;

+	/** @var ILoadBalancer */
+	private $loadBalancer;
+
 	/**
+	 * @param ILoadBalancer $loadBalancer
 	 * @param array $params
 	 *  - emailEnabled: (bool) must be true for the option to email passwords to be present
 	 *  - newPasswordExpiry: (int) expiraton time of temporary passwords, in seconds
 	 *  - passwordReminderResendTime: (int) cooldown period in hours until a password reminder can
 	 *    be sent to the same user again,
 	 */
-	public function __construct( $params = [] ) {
+	public function __construct( ILoadBalancer $loadBalancer, $params = [] ) {
 		parent::__construct( $params );

 		if ( isset( $params['emailEnabled'] ) ) {
@ -74,6 +79,7 @@ class TemporaryPasswordPrimaryAuthenticationProvider
 		if ( isset( $params['allowRequiringEmailForResets'] ) ) {
 			$this->allowRequiringEmail = $params['allowRequiringEmailForResets'];
 		}
+		$this->loadBalancer = $loadBalancer;
 	}

 	protected function postInitSetup() {
@ -136,7 +142,7 @@ class TemporaryPasswordPrimaryAuthenticationProvider
 			return AuthenticationResponse::newAbstain();
 		}

-		$dbr = wfGetDB( DB_REPLICA );
+		$dbr = $this->loadBalancer->getConnectionRef( DB_REPLICA );
 		$row = $dbr->selectRow(
 			'user',
 			[
@ -185,7 +191,7 @@ class TemporaryPasswordPrimaryAuthenticationProvider
 			return false;
 		}

-		$dbr = wfGetDB( DB_REPLICA );
+		$dbr = $this->loadBalancer->getConnectionRef( DB_REPLICA );
 		$row = $dbr->selectRow(
 			'user',
 			[ 'user_newpassword', 'user_newpass_time' ],
@ -214,7 +220,7 @@ class TemporaryPasswordPrimaryAuthenticationProvider
 		}

 		list( $db, $options ) = \DBAccessObjectUtils::getDBOptions( $flags );
-		return (bool)wfGetDB( $db )->selectField(
+		return (bool)$this->loadBalancer->getConnectionRef( $db )->selectField(
 			[ 'user' ],
 			'user_id',
 			[ 'user_name' => $username ],
@ -240,7 +246,7 @@ class TemporaryPasswordPrimaryAuthenticationProvider
 			return \StatusValue::newGood( 'ignored' );
 		}

-		$row = wfGetDB( DB_PRIMARY )->selectRow(
+		$row = $this->loadBalancer->getConnectionRef( DB_PRIMARY )->selectRow(
 			'user',
 			[ 'user_id', 'user_newpass_time' ],
 			[ 'user_name' => $username ],
@ -297,7 +303,7 @@ class TemporaryPasswordPrimaryAuthenticationProvider
 			return;
 		}

-		$dbw = wfGetDB( DB_PRIMARY );
+		$dbw = $this->loadBalancer->getConnectionRef( DB_PRIMARY );

 		$sendMail = false;
 		if ( $req->action !== AuthManager::ACTION_REMOVE &&
@ -396,7 +402,7 @@ class TemporaryPasswordPrimaryAuthenticationProvider

 		if ( $mailpassword ) {
 			// Send email after DB commit
-			wfGetDB( DB_PRIMARY )->onTransactionCommitOrIdle(
+			$this->loadBalancer->getConnectionRef( DB_PRIMARY )->onTransactionCommitOrIdle(
 				function () use ( $user, $creator, $req ) {
 					$this->sendNewAccountEmail( $user, $creator, $req->password );
 				},
--- a/tests/common/TestSetup.php
+++ b/tests/common/TestSetup.php
@ -108,12 +108,18 @@ class TestSetup {
 			'primaryauth' => [
 				[
 					'class' => MediaWiki\Auth\TemporaryPasswordPrimaryAuthenticationProvider::class,
+					'services' => [
+						'DBLoadBalancer',
+					],
 					'args' => [ [
 						'authoritative' => false,
 					] ],
 				],
 				[
 					'class' => MediaWiki\Auth\LocalPasswordPrimaryAuthenticationProvider::class,
+					'services' => [
+						'DBLoadBalancer',
+					],
 					'args' => [ [
 						'authoritative' => true,
 					] ],
--- a/tests/phpunit/includes/auth/EmailNotificationSecondaryAuthenticationProviderTest.php
+++ b/tests/phpunit/includes/auth/EmailNotificationSecondaryAuthenticationProviderTest.php
@ -3,44 +3,57 @@
 namespace MediaWiki\Auth;

 use Config;
+use HashConfig;
 use MediaWiki\MediaWikiServices;
 use MediaWiki\User\UserNameUtils;
 use Psr\Container\ContainerInterface;
 use Psr\Log\LoggerInterface;
 use Psr\Log\NullLogger;
+use TestLogger;
 use Wikimedia\TestingAccessWrapper;

 /**
 * @covers \MediaWiki\Auth\EmailNotificationSecondaryAuthenticationProvider
+ * @group Database
 */
 class EmailNotificationSecondaryAuthenticationProviderTest extends \MediaWikiIntegrationTestCase {
+
+	/**
+	 * @param array $options
+	 * @return EmailNotificationSecondaryAuthenticationProvider
+	 */
+	private function getProvider( array $options = [] ) : EmailNotificationSecondaryAuthenticationProvider {
+		$services = $this->getServiceContainer();
+		$provider = new EmailNotificationSecondaryAuthenticationProvider(
+			$options['loadBalancer'] ?? $services->getDBLoadBalancer(),
+			$options // make things easier for tests by using the same options
+		);
+		$provider->init(
+			$options['logger'] ?? new TestLogger(),
+			$options['authManager'] ?? $this->createNoOpMock( AuthManager::class ),
+			$options['hookContainer'] ?? $this->createHookContainer(),
+			$options['config'] ?? new HashConfig( [] ),
+			$options['userNameUtils'] ?? $this->createNoOpMock( UserNameUtils::class )
+		);
+		return $provider;
+	}
+
 	public function testConstructor() {
 		$config = new \HashConfig( [
 			'EnableEmail' => true,
 			'EmailAuthentication' => true,
 		] );

-		$provider = new EmailNotificationSecondaryAuthenticationProvider();
-		$provider->init(
-			$this->createNoOpMock( NullLogger::class ),
-			$this->createNoOpMock( AuthManager::class ),
-			$this->createHookContainer(),
-			$config,
-			$this->createNoOpMock( UserNameUtils::class )
-		);
+		$provider = $this->getProvider( [
+			'config' => $config,
+		] );
 		$providerPriv = TestingAccessWrapper::newFromObject( $provider );
 		$this->assertTrue( $providerPriv->sendConfirmationEmail );

-		$provider = new EmailNotificationSecondaryAuthenticationProvider( [
+		$provider = $this->getProvider( [
+			'config' => $config,
 			'sendConfirmationEmail' => false,
 		] );
-		$provider->init(
-			$this->createNoOpMock( NullLogger::class ),
-			$this->createNoOpMock( AuthManager::class ),
-			$this->createHookContainer(),
-			$config,
-			$this->createNoOpMock( UserNameUtils::class )
-		);
 		$providerPriv = TestingAccessWrapper::newFromObject( $provider );
 		$this->assertFalse( $providerPriv->sendConfirmationEmail );
 	}
@ -51,7 +64,7 @@ class EmailNotificationSecondaryAuthenticationProviderTest extends \MediaWikiInt
 	 * @param AuthenticationRequest[] $expected
 	 */
 	public function testGetAuthenticationRequests( $action, $expected ) {
-		$provider = new EmailNotificationSecondaryAuthenticationProvider( [
+		$provider = $this->getProvider( [
 			'sendConfirmationEmail' => true,
 		] );
 		$this->assertSame( $expected, $provider->getAuthenticationRequests( $action, [] ) );
@ -68,7 +81,7 @@ class EmailNotificationSecondaryAuthenticationProviderTest extends \MediaWikiInt
 	}

 	public function testBeginSecondaryAuthentication() {
-		$provider = new EmailNotificationSecondaryAuthenticationProvider( [
+		$provider = $this->getProvider( [
 			'sendConfirmationEmail' => true,
 		] );
 		$this->assertEquals( AuthenticationResponse::newAbstain(),
@ -117,28 +130,19 @@ class EmailNotificationSecondaryAuthenticationProviderTest extends \MediaWikiInt
 			->willReturnSelf();
 		$userNotExpectsConfirmation->expects( $this->never() )->method( 'sendConfirmationMail' );

-		$provider = new EmailNotificationSecondaryAuthenticationProvider( [
+		$provider = $this->getProvider( [
 			'sendConfirmationEmail' => false,
+			'authManager' => $authManager,
+			'hookContainer' => $hookContainer,
+			'userNameUtils' => $userNameUtils
 		] );
-		$provider->init(
-			$this->createNoOpMock( NullLogger::class ),
-			$authManager,
-			$hookContainer,
-			$this->createNoOpAbstractMock( Config::class ),
-			$userNameUtils
-		);
 		$provider->beginSecondaryAccountCreation( $userNotExpectsConfirmation, $creator, [] );

-		$provider = new EmailNotificationSecondaryAuthenticationProvider( [
+		$provider = $this->getProvider( [
 			'sendConfirmationEmail' => true,
+			'authManager' => $authManager,
+			'userNameUtils' => $userNameUtils
 		] );
-		$provider->init(
-			$this->createNoOpMock( NullLogger::class ),
-			$authManager,
-			$this->createHookContainer(),
-			$this->createNoOpAbstractMock( Config::class ),
-			$userNameUtils
-		);
 		$provider->beginSecondaryAccountCreation( $userWithoutEmail, $creator, [] );
 		$provider->beginSecondaryAccountCreation( $userExpectsConfirmation, $creator, [] );

--- a/tests/phpunit/includes/auth/LocalPasswordPrimaryAuthenticationProviderTest.php
+++ b/tests/phpunit/includes/auth/LocalPasswordPrimaryAuthenticationProviderTest.php
@ -35,7 +35,7 @@ class LocalPasswordPrimaryAuthenticationProviderTest extends \MediaWikiIntegrati
 		}
 		$config = new \MultiConfig( [
 			$this->config,
-			MediaWikiServices::getInstance()->getMainConfig()
+			$mwServices->getMainConfig()
 		] );

 		// We need a real HookContainer since testProviderChangeAuthenticationData()
@ -62,7 +62,10 @@ class LocalPasswordPrimaryAuthenticationProviderTest extends \MediaWikiIntegrati
 		$this->validity = \Status::newGood();
 		$provider = $this->getMockBuilder( LocalPasswordPrimaryAuthenticationProvider::class )
 			->onlyMethods( [ 'checkPasswordValidity' ] )
-			->setConstructorArgs( [ [ 'loginOnly' => $loginOnly ] ] )
+			->setConstructorArgs( [
+				$mwServices->getDBLoadBalancer(),
+				[ 'loginOnly' => $loginOnly ]
+			] )
 			->getMock();

 		$provider->method( 'checkPasswordValidity' )
@ -154,7 +157,9 @@ class LocalPasswordPrimaryAuthenticationProviderTest extends \MediaWikiIntegrati
 		$this->config->set( 'PasswordExpireGrace', 100 );
 		$this->config->set( 'InvalidPasswordReset', true );

-		$provider = new LocalPasswordPrimaryAuthenticationProvider();
+		$provider = new LocalPasswordPrimaryAuthenticationProvider(
+			$this->getServiceContainer()->getDBLoadBalancer()
+		);
 		$provider->init(
 			new NullLogger(),
 			$this->manager,
--- a/tests/phpunit/includes/auth/TemporaryPasswordPrimaryAuthenticationProviderTest.php
+++ b/tests/phpunit/includes/auth/TemporaryPasswordPrimaryAuthenticationProviderTest.php
@ -10,6 +10,8 @@ use Wikimedia\ScopedCallback;
 use Wikimedia\TestingAccessWrapper;

 /**
+ * TODO clean up and reduce duplication
+ *
 * @group AuthManager
 * @group Database
 * @covers \MediaWiki\Auth\TemporaryPasswordPrimaryAuthenticationProvider
@ -64,7 +66,7 @@ class TemporaryPasswordPrimaryAuthenticationProviderTest extends \MediaWikiInteg
 		$mockedMethods[] = 'checkPasswordValidity';
 		$provider = $this->getMockBuilder( TemporaryPasswordPrimaryAuthenticationProvider::class )
 			->onlyMethods( $mockedMethods )
-			->setConstructorArgs( [ $params ] )
+			->setConstructorArgs( [ $mwServices->getDBLoadBalancer(), $params ] )
 			->getMock();
 		$provider->method( 'checkPasswordValidity' )
 			->will( $this->returnCallback( function () {
@ -122,7 +124,11 @@ class TemporaryPasswordPrimaryAuthenticationProviderTest extends \MediaWikiInteg
 			'AllowRequiringEmailForResets' => false,
 		] );

-		$p = TestingAccessWrapper::newFromObject( new TemporaryPasswordPrimaryAuthenticationProvider() );
+		$p = TestingAccessWrapper::newFromObject(
+			new TemporaryPasswordPrimaryAuthenticationProvider(
+				$this->getServiceContainer()->getDBLoadBalancer()
+			)
+		);
 		$p->init(
 			$this->createNoOpMock( NullLogger::class ),
 			$this->createNoOpMock( AuthManager::class ),
@ -134,12 +140,17 @@ class TemporaryPasswordPrimaryAuthenticationProviderTest extends \MediaWikiInteg
 		$this->assertSame( 100, $p->newPasswordExpiry );
 		$this->assertSame( 101, $p->passwordReminderResendTime );

-		$p = TestingAccessWrapper::newFromObject( new TemporaryPasswordPrimaryAuthenticationProvider( [
-			'emailEnabled' => true,
-			'newPasswordExpiry' => 42,
-			'passwordReminderResendTime' => 43,
-			'allowRequiringEmailForResets' => true,
-		] ) );
+		$p = TestingAccessWrapper::newFromObject(
+			new TemporaryPasswordPrimaryAuthenticationProvider(
+				$this->getServiceContainer()->getDBLoadBalancer(),
+				[
+					'emailEnabled' => true,
+					'newPasswordExpiry' => 42,
+					'passwordReminderResendTime' => 43,
+					'allowRequiringEmailForResets' => true,
+				]
+			)
+		);
 		$p->init(
 			$this->createNoOpMock( NullLogger::class ),
 			$this->createNoOpMock( AuthManager::class ),