Block creating authentication cookies when inappropriate to do so.

class-two-factor-core.php

@@ -57,6 +57,13 @@ class Two_Factor_Core {
 	 */
 	private static $password_auth_tokens = array();
 
+	/**
+	 * Keep track of who wp_set_auth_cookie() is running for.
+	 *
+	 * @var int
+	 */
+	private static $last_auth_cookie_user = 0;
+
 	/**
 	 * Set up filters and actions.
 	 *
@@ -68,6 +75,10 @@ public static function add_hooks( $compat ) {
 		add_action( 'plugins_loaded', array( __CLASS__, 'load_textdomain' ) );
 		add_action( 'init', array( __CLASS__, 'get_providers' ) );
 		add_action( 'wp_login', array( __CLASS__, 'wp_login' ), 10, 2 );
+		add_action( 'clear_auth_cookie', array( __CLASS__, 'clear_auth_cookie' ) );
+		add_action( 'set_auth_cookie', array( __CLASS__, 'set_auth_cookie' ), 10, 4 );
+		add_action( 'send_auth_cookies', array( __CLASS__, 'send_auth_cookies' ), 10, 2 );
+		add_action( 'login_form_show_2fa', array( __CLASS__, 'login_form_show_2fa' ) );
 		add_action( 'login_form_validate_2fa', array( __CLASS__, 'login_form_validate_2fa' ) );
 		add_action( 'login_form_backup_2fa', array( __CLASS__, 'backup_2fa' ) );
 		add_action( 'show_user_profile', array( __CLASS__, 'user_two_factor_options' ) );
@@ -433,13 +444,67 @@ public static function wp_login( $user_login, $user ) {
 		// Invalidate the current login session to prevent from being re-used.
 		self::destroy_current_session_for_user( $user );
 
-		// Also clear the cookies which are no longer valid.
+		// Clear any cookies which are no longer valid.
 		wp_clear_auth_cookie();
 
 		self::show_two_factor_login( $user );
 		exit;
 	}
 
+	/**
+	 * Maybe abort sending authentication cookies, and prompt for two-factor instead.
+	 *
+	 * @param bool $send_cookies Whether to send the authentication cookies.
+	 * @param int  $user_id      The User ID having cookies sent for. WordPress 6.2+.
+	 * @return bool Filtered `$send_cookies`.
+	 */
+	public static function send_auth_cookies( $send_cookies, $user_id = null ) {
+		if ( ! $send_cookies ) {
+			return $send_cookies;
+		}
+
+		// WordPress < 6.2 compat. See https://core.trac.wordpress.org/ticket/56971.
+		if ( is_null( $user_id ) ) {
+			$user_id = self::$last_auth_cookie_user;
+		}
+
+		// Don't send auth cookies for an authenticated user, unless they've passed the two-factor check.
+		if (
+			$send_cookies &&
+			$user_id &&
+			self::is_user_using_two_factor( $user_id )
+		) {
+			$send_cookies = false;
+
+			// If we're not on wp-login.php, redirect there. This is for when `wp_set_auth_cookie()` is called outside of wp-login.php.
+			if ( ! did_action( 'login_init' ) ) {
+				$user = get_userdata( $user_id );
+				self::redirect_to_two_factor( $user );
+				exit;
+			}
+		}
+
+		return $send_cookies;
+	}
+
+	/**
+	 * Keep track of the last user a cookie was generated for.
+	 *
+	 * This is used for when context is not provided via send_auth_cookies.
+	 */
+	public static function set_auth_cookie( $auth_cookie, $expire, $expiration, $user_id ) {
+		self::$last_auth_cookie_user = $user_id;
+	}
+
+	/**
+	 * Keep track of the last user a cookie was generated for.
+	 *
+	 * When clearing cookies, there is not user context.
+	 */
+	public static function clear_auth_cookie() {
+		self::$last_auth_cookie_user = 0;
+	}
+
 	/**
 	 * Destroy the known password-based authentication sessions for the current user.
 	 *
@@ -530,6 +595,60 @@ public static function show_two_factor_login( $user ) {
 		self::login_html( $user, $login_nonce['key'], $redirect_to );
 	}
 
+	/**
+	 * Display the Two Factor login.
+	 */
+	public static function login_form_show_2fa() {
+		$wp_auth_id  = filter_input( INPUT_GET, 'wp-auth-id', FILTER_SANITIZE_NUMBER_INT );
+		$nonce       = filter_input( INPUT_GET, 'wp-auth-nonce', FILTER_CALLBACK, array( 'options' => 'sanitize_key' ) );
+		$redirect_to = ! empty( $_REQUEST['redirect_to'] ) ? $_REQUEST['redirect_to'] : wp_get_referer();
+		$user        = get_user_by( 'id', $wp_auth_id );
+
+		if ( ! $wp_auth_id || ! $nonce ) {
+			return;
+		}
+
+		$user = get_userdata( $wp_auth_id );
+		if ( ! $user ) {
+			return;
+		}
+
+		if ( true !== self::verify_login_nonce( $user->ID, $nonce ) ) {
+			wp_safe_redirect( home_url() );
+			exit;
+		}
+
+		self::login_html( $user, $nonce, $redirect_to );
+		exit;
+	}
+
+	/**
+	 * Redirect the user to the two-factor authentication.
+	 */
+	public static function redirect_to_two_factor( $user ) {
+		if ( ! $user ) {
+			$user = wp_get_current_user();
+		}
+
+		$login_nonce = self::create_login_nonce( $user->ID );
+		if ( ! $login_nonce ) {
+			wp_die( esc_html__( 'Failed to create a login nonce.', 'two-factor' ) );
+		}
+
+		$redirect_to = isset( $_REQUEST['redirect_to'] ) ? $_REQUEST['redirect_to'] : '';
+
+		wp_safe_redirect( self::login_url(
+			array(
+				'action'        => 'show_2fa',
+				'wp-auth-id'    => $user->ID,
+				'wp-auth-nonce' => $login_nonce['key'],
+				'redirect_to'   => $redirect_to,
+			)
+		) );
+		exit;
+	}
+
+
 	/**
 	 * Display the Backup code 2fa screen.
 	 *
@@ -897,6 +1016,9 @@ public static function login_form_validate_2fa() {
 			$rememberme = true;
 		}
 
+		// Allow authentication cookies to be set for this user.
+		remove_action( 'send_auth_cookies', array( __CLASS__, 'send_auth_cookies' ) );
+
 		wp_set_auth_cookie( $user->ID, $rememberme );
 
 		do_action( 'two_factor_user_authenticated', $user );