Make Key objects aware of their database IDs

Bug: T242031
Depends-On: I1db9b04a42783b8b64ed69f1f950c794c8659209
Change-Id: I0d8d0a42ce627387949dbbbb32fc318088b3538e

maintenance/disableOATHAuthForUser.php

@@ -14,8 +14,8 @@ require_once "$IP/maintenance/Maintenance.php";
 class DisableOATHAuthForUser extends Maintenance {
 	public function __construct() {
 		parent::__construct();
-		$this->addDescription( 'Remove OATHAuth from a specific user' );
-		$this->addArg( 'user', 'The username to remove OATHAuth from.' );
+		$this->addDescription( 'Remove all two-factor authentication devices from a specific user' );
+		$this->addArg( 'user', 'The username to remove 2FA devices from.' );
 		$this->requireExtension( 'OATHAuth' );
 	}
 
@@ -31,17 +31,17 @@ class DisableOATHAuthForUser extends Maintenance {
 		$repo = OATHAuthServices::getInstance()->getUserRepository();
 		$oathUser = $repo->findByUser( $user );
 		if ( !$oathUser->isTwoFactorAuthEnabled() ) {
-			$this->fatalError( "User $username doesn't have OATHAuth enabled!" );
+			$this->fatalError( "User $username does not have two-factor authentication enabled!" );
 		}
 
-		$repo->remove( $oathUser, 'Maintenance script', false );
+		$repo->removeAll( $oathUser, 'Maintenance script', false );
 		// Kill all existing sessions.
 		// If this request to disable 2FA was social-engineered by an attacker,
 		// the legitimate user will hopefully log in again to the wiki, and notice that the second factor
 		// is missing or different, and alert the operators.
 		SessionManager::singleton()->invalidateSessionsForUser( $user );
 
-		$this->output( "OATHAuth disabled for $username.\n" );
+		$this->output( "Two-factor authentication disabled for $username.\n" );
 	}
 }
 

src/HTMLForm/TOTPDisableForm.php

@@ -71,7 +71,7 @@ class TOTPDisableForm extends OATHAuthOOUIHTMLForm {
 		}
 
 		$this->oathUser->setKeys();
-		$this->oathRepo->remove( $this->oathUser, $this->getRequest()->getIP(), true );
+		$this->oathRepo->removeAll( $this->oathUser, $this->getRequest()->getIP(), true );
 
 		return true;
 	}

src/IAuthKey.php

@@ -7,11 +7,15 @@ use stdClass;
 
 interface IAuthKey extends JsonSerializable {
 
+	/**
+	 * @return int|null the ID of this key in the oathauth_devices table, or null if this key has not been saved yet
+	 */
+	public function getId(): ?int;
+
 	/**
 	 * @param array|stdClass $data
 	 * @param OATHUser $user
 	 * @return bool
 	 */
 	public function verify( $data, OATHUser $user );
-
 }

src/Key/TOTPKey.php

@@ -41,7 +41,10 @@ use Psr\Log\LoggerInterface;
  * @ingroup Extensions
  */
 class TOTPKey implements IAuthKey {
-	/** @var array Two-factor binary secret */
+	/** @var int|null */
+	private ?int $id;
+
+	/** @var array Two factor binary secret */
 	private $secret;
 
 	/** @var string[] List of recovery codes */
@@ -53,6 +56,7 @@ class TOTPKey implements IAuthKey {
 	 */
 	public static function newFromRandom() {
 		$object = new self(
+			null,
 			Base32::encode( random_bytes( 10 ) ),
 			[]
 		);
@@ -70,14 +74,17 @@ class TOTPKey implements IAuthKey {
 		if ( !isset( $data['secret'] ) || !isset( $data['scratch_tokens'] ) ) {
 			return null;
 		}
-		return new static( $data['secret'], $data['scratch_tokens'] );
+		return new static( $data['id'] ?? null, $data['secret'], $data['scratch_tokens'] );
 	}
 
 	/**
+	 * @param int|null $id the database id of this key
 	 * @param string $secret
 	 * @param array $recoveryCodes
 	 */
-	public function __construct( $secret, array $recoveryCodes ) {
+	public function __construct( ?int $id, $secret, array $recoveryCodes ) {
+		$this->id = $id;
+
 		// Currently hardcoded values; might be used in the future
 		$this->secret = [
 			'mode' => 'hotp',
@@ -88,6 +95,13 @@ class TOTPKey implements IAuthKey {
 		$this->recoveryCodes = array_values( $recoveryCodes );
 	}
 
+	/**
+	 * @return int|null
+	 */
+	public function getId(): ?int {
+		return $this->id;
+	}
+
 	/**
 	 * @return string
 	 */

src/OATHUserRepository.php

@@ -76,41 +76,7 @@ class OATHUserRepository implements LoggerAwareInterface {
 			$uid = $this->centralIdLookupFactory->getLookup()
 				->centralIdFromLocalUser( $user );
 			$oathUser = new OATHUser( $user, $uid );
-
-			$res = $this->dbProvider
-				->getReplicaDatabase( 'virtual-oathauth' )
-				->newSelectQueryBuilder()
-				->select( [
-					'oad_data',
-					'oat_name',
-				] )
-				->from( 'oathauth_devices' )
-				->join( 'oathauth_types', null, [ 'oat_id = oad_type' ] )
-				->where( [ 'oad_user' => $uid ] )
-				->caller( __METHOD__ )
-				->fetchResultSet();
-
-			$module = null;
-
-			foreach ( $res as $row ) {
-				if ( $module && $row->oat_name !== $module->getName() ) {
-					// Not supported by current application-layer code.
-					throw new RuntimeException( "user {$uid} has multiple different oathauth modules defined" );
-				}
-
-				if ( !$module ) {
-					$module = $this->moduleRegistry->getModuleByKey( $row->oat_name );
-					$oathUser->setModule( $module );
-
-					if ( !$module ) {
-						throw new MWException( 'oathauth-module-invalid' );
-					}
-				}
-
-				$keyData = FormatJson::decode( $row->oad_data, true );
-				$oathUser->addKey( $module->newKey( $keyData ) );
-			}
-
+			$this->loadKeysFromDatabase( $oathUser );
 			$this->cache->set( $user->getName(), $oathUser );
 		}
 		return $oathUser;
@@ -128,7 +94,6 @@ class OATHUserRepository implements LoggerAwareInterface {
 		}
 		$prevUser = $this->findByUser( $user->getUser() );
 		$userId = $this->centralIdLookupFactory->getLookup()->centralIdFromLocalUser( $user->getUser() );
-
 		$moduleId = $this->moduleRegistry->getModuleId( $user->getModule()->getName() );
 
 		$dbw = $this->dbProvider->getPrimaryDatabase( 'virtual-oathauth' );
@@ -141,20 +106,22 @@ class OATHUserRepository implements LoggerAwareInterface {
 			->caller( __METHOD__ )
 			->execute();
 
-		$insert = $dbw->newInsertQueryBuilder()
-			->insertInto( 'oathauth_devices' )
-			->caller( __METHOD__ );
 		foreach ( $user->getKeys() as $key ) {
-			$insert->row( [
-				'oad_user' => $userId,
-				'oad_type' => $moduleId,
-				'oad_data' => FormatJson::encode( $key->jsonSerialize() )
-			] );
+			$dbw->newInsertQueryBuilder()
+				->insertInto( 'oathauth_devices' )
+				->row( [
+					'oad_user' => $userId,
+					'oad_type' => $moduleId,
+					'oad_data' => FormatJson::encode( $key->jsonSerialize() )
+				] )
+				->caller( __METHOD__ )
+				->execute();
 		}
-		$insert->execute();
 
 		$dbw->endAtomic( __METHOD__ );
 
+		$this->loadKeysFromDatabase( $user );
+
 		$userName = $user->getUser()->getName();
 		$this->cache->set( $userName, $user );
 
@@ -209,7 +176,7 @@ class OATHUserRepository implements LoggerAwareInterface {
 
 		$hasExistingKey = $user->isTwoFactorAuthEnabled();
 
-		$key = $module->newKey( $keyData );
+		$key = $module->newKey( $keyData + [ 'id' => $id ] );
 		$user->addKey( $key );
 
 		$this->logger->info( 'OATHAuth {oathtype} key {key} added for {user} from {clientip}', [
@@ -231,8 +198,19 @@ class OATHUserRepository implements LoggerAwareInterface {
 	 * @param OATHUser $user
 	 * @param string $clientInfo
 	 * @param bool $self Whether the user disabled the 2FA themselves
+	 *
+	 * @deprecated since 1.41, use removeAll() instead
 	 */
 	public function remove( OATHUser $user, $clientInfo, bool $self ) {
+		$this->removeAll( $user, $clientInfo, $self );
+	}
+
+	/**
+	 * @param OATHUser $user
+	 * @param string $clientInfo
+	 * @param bool $self Whether they disabled it themselves
+	 */
+	public function removeAll( OATHUser $user, $clientInfo, bool $self ) {
 		$userId = $this->centralIdLookupFactory->getLookup()
 			->centralIdFromLocalUser( $user->getUser() );
 		$this->dbProvider->getPrimaryDatabase( 'virtual-oathauth' )
@@ -260,4 +238,47 @@ class OATHUserRepository implements LoggerAwareInterface {
 
 		Manager::notifyDisabled( $user, $self );
 	}
+
+	private function loadKeysFromDatabase( OATHUser $user ): void {
+		$uid = $this->centralIdLookupFactory->getLookup()
+			->centralIdFromLocalUser( $user->getUser() );
+
+		$res = $this->dbProvider
+			->getReplicaDatabase( 'virtual-oathauth' )
+			->newSelectQueryBuilder()
+			->select( [
+				'oad_id',
+				'oad_data',
+				'oat_name',
+			] )
+			->from( 'oathauth_devices' )
+			->join( 'oathauth_types', null, [ 'oat_id = oad_type' ] )
+			->where( [ 'oad_user' => $uid ] )
+			->caller( __METHOD__ )
+			->fetchResultSet();
+
+		$module = null;
+
+		// Clear stored key list before loading keys
+		$user->disable();
+
+		foreach ( $res as $row ) {
+			if ( $module && $row->oat_name !== $module->getName() ) {
+				// Not supported by current application-layer code.
+				throw new RuntimeException( "User {$uid} has multiple different two-factor modules defined" );
+			}
+
+			if ( !$module ) {
+				$module = $this->moduleRegistry->getModuleByKey( $row->oat_name );
+				$user->setModule( $module );
+
+				if ( !$module ) {
+					throw new MWException( 'oathauth-module-invalid' );
+				}
+			}
+
+			$keyData = FormatJson::decode( $row->oad_data, true );
+			$user->addKey( $module->newKey( $keyData + [ 'id' => (int)$row->oad_id ] ) );
+		}
+	}
 }

src/Special/DisableOATHForUser.php

@@ -142,7 +142,7 @@ class DisableOATHForUser extends FormSpecialPage {
 			return [ 'oathauth-throttled', Message::durationParam( 60 ) ];
 		}
 
-		$this->userRepo->remove( $oathUser, $this->getRequest()->getIP(), false );
+		$this->userRepo->removeAll( $oathUser, $this->getRequest()->getIP(), false );
 
 		// messages used: logentry-oath-disable-other, log-action-oath-disable-other
 		$logEntry = new ManualLogEntry( 'oath', 'disable-other' );