diff --git a/class-two-factor-core.php b/class-two-factor-core.php index 99d3595c..8cab506a 100644 --- a/class-two-factor-core.php +++ b/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; + } + + /** + * Reset 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 );