diff --git a/class-two-factor-core.php b/class-two-factor-core.php index 61215108..420cb4e1 100644 --- a/class-two-factor-core.php +++ b/class-two-factor-core.php @@ -447,6 +447,42 @@ public static function get_available_providers_for_user( $user = null ) { return $configured_providers; } + /** + * Fetch the provider for the request based on the user preferences. + * + * @param int|WP_User $user Optonal. User ID, or WP_User object of the the user. Defaults to current user. + * @param null|string|object $preferred_provider Optional. The name of the provider, the provider, or empty. + * @return null|object The provider + */ + public static function get_provider_for_user( $user = null, $preferred_provider = null ) { + $user = self::fetch_user( $user ); + if ( ! $user ) { + return null; + } + + // If a specific provider instance is passed, process it just as the key. + if ( $preferred_provider && $preferred_provider instanceof Two_Factor_Provider ) { + $preferred_provider = $preferred_provider->get_key(); + } + + // Default to the currently logged in provider. + if ( ! $preferred_provider && get_current_user_id() === $user->ID ) { + $session = self::get_current_user_session(); + if ( ! empty( $session['two-factor-provider'] ) ) { + $preferred_provider = $session['two-factor-provider']; + } + } + + if ( is_string( $preferred_provider ) ) { + $providers = self::get_available_providers_for_user( $user ); + if ( isset( $providers[ $preferred_provider ] ) ) { + return $providers[ $preferred_provider ]; + } + } + + return self::get_primary_provider_for_user( $user ); + } + /** * Gets the Two-Factor Auth provider for the specified|current user. * @@ -733,10 +769,9 @@ public static function clear_password_reset_notice( $user ) { * @param string|object $provider An override to the provider. */ public static function login_html( $user, $login_nonce, $redirect_to, $error_msg = '', $provider = null, $action = 'validate_2fa' ) { - if ( empty( $provider ) ) { - $provider = self::get_primary_provider_for_user( $user->ID ); - } elseif ( is_string( $provider ) && method_exists( $provider, 'get_instance' ) ) { - $provider = call_user_func( array( $provider, 'get_instance' ) ); + $provider = self::get_provider_for_user( $user, $provider ); + if ( ! $provider ) { + wp_die( __( 'Cheatin’ uh?', 'two-factor' ) ); } $provider_key = $provider->get_key(); @@ -1095,10 +1130,7 @@ public static function is_user_rate_limited( $user ) { * @return int|false The last time the two-factor was validated on success, false if not currently using a 2FA session. */ public static function is_current_user_session_two_factor() { - $user_id = get_current_user_id(); - $token = wp_get_session_token(); - $manager = WP_Session_Tokens::get_instance( $user_id ); - $session = $manager->get( $token ); + $session = self::get_current_user_session(); if ( empty( $session['two-factor-login'] ) ) { return false; @@ -1123,8 +1155,8 @@ public static function current_user_can_update_two_factor_options( $context = 'd return false; } - // If the current user is not a two-factor user, not having a two-factor session is okay. - if ( ! self::is_user_using_two_factor( $user_id ) && ! $is_two_factor_session ) { + // If the current user is not using two-factor, they can adjust the settings. + if ( ! self::is_user_using_two_factor( $user_id ) ) { return true; } @@ -1224,11 +1256,9 @@ public static function _login_form_validate_2fa( $user, $nonce = '', $provider = return; } - $providers = self::get_available_providers_for_user( $user ); - if ( $provider && isset( $providers[ $provider ] ) ) { - $provider = $providers[ $provider ]; - } else { - $provider = self::get_primary_provider_for_user( $user->ID ); + $provider = self::get_provider_for_user( $user, $provider ); + if ( ! $provider ) { + wp_die( __( 'Cheatin’ uh?', 'two-factor' ) ); } // Run the provider processing. @@ -1347,21 +1377,10 @@ public static function _login_form_revalidate_2fa( $provider = '', $redirect_to return; } - $user = wp_get_current_user(); - $session_token = wp_get_session_token(); - $session_manager = WP_Session_Tokens::get_instance( $user->ID ); - $session = $session_manager->get( $session_token ); - $providers = self::get_available_providers_for_user( $user ); - - // Default to the currently logged in provider. - if ( ! $provider && ! empty( $session['two-factor-provider'] ) ) { - $provider = $session['two-factor-provider']; - } - - if ( $provider && isset( $providers[ $provider ] ) ) { - $provider = $providers[ $provider ]; - } else { - $provider = self::get_primary_provider_for_user( $user->ID ); + $user = wp_get_current_user(); + $provider = self::get_provider_for_user( $user, $provider ); + if ( ! $provider ) { + wp_die( __( 'Cheatin’ uh?', 'two-factor' ) ); } // Run the provider processing. @@ -1378,10 +1397,11 @@ public static function _login_form_revalidate_2fa( $provider = '', $redirect_to return; } - $session['two-factor-provider'] = get_class( $provider ); - $session['two-factor-login'] = time(); - - $session_manager->update( $session_token, $session ); + // Update the session metadata with the revalidation details. + self::update_current_user_session( array( + 'two-factor-provider' => $provider->get_key(), + 'two-factor-login' => time(), + ) ); // Must be global because that's how login_header() uses it. global $interim_login; @@ -1825,27 +1845,72 @@ public static function user_two_factor_options_update( $user_id ) { update_user_meta( $user_id, self::PROVIDER_USER_META_KEY, $new_provider ); } - // Have we enabled new providers? Set this as a 2FA session, so they can continue to edit. - if ( - ! $existing_providers && - $enabled_providers && - ! self::is_current_user_session_two_factor() && - $user_id === get_current_user_id() - ) { - $token = wp_get_session_token(); - if ( $token ) { - $manager = WP_Session_Tokens::get_instance( $user_id ); - $session = $manager->get( $token ); - - $session['two-factor-provider'] = ''; // Set the key, but not the provider, as no provider has been used yet. - $session['two-factor-login'] = time(); - - $manager->update( $token, $session ); + // Have we changed the two-factor settings for the current user? Alter their session metadata. + if ( $user_id === get_current_user_id() ) { + + if ( $enabled_providers && ! $existing_providers && ! self::is_current_user_session_two_factor() ) { + // We've enabled two-factor from a non-two-factor session, set the key but not the provider, as no provider has been used yet. + self::update_current_user_session( array( + 'two-factor-provider' => '', + 'two-factor-login' => time(), + ) ); + } elseif ( $existing_providers && ! $enabled_providers ) { + // We've disabled two-factor, remove session metadata. + self::update_current_user_session( array( + 'two-factor-provider' => null, + 'two-factor-login' => null, + ) ); } } } } + /** + * Update the current user session metadata. + * + * Any values set in $data that are null will be removed from the user session metadata. + * + * @param array $data The data to append/remove from the current session. + * @return bool + */ + public static function update_current_user_session( $data = array() ) { + $user_id = get_current_user_id(); + $token = wp_get_session_token(); + if ( ! $user_id || ! $token ) { + return false; + } + + $manager = WP_Session_Tokens::get_instance( $user_id ); + $session = $manager->get( $token ); + + // Add any session data. + $session = array_merge( $session, $data ); + + // Remove any set null fields. + foreach ( array_filter( $data, 'is_null' ) as $key => $null ) { + unset( $session[ $key ] ); + } + + return $manager->update( $token, $session ); + } + + /** + * Fetch the current user session metadata. + * + * @return false|array The session array, false on error. + */ + public static function get_current_user_session() { + $user_id = get_current_user_id(); + $token = wp_get_session_token(); + if ( ! $user_id || ! $token ) { + return false; + } + + $manager = WP_Session_Tokens::get_instance( $user_id ); + + return $manager->get( $token ); + } + /** * Should the login session persist between sessions. * diff --git a/tests/class-two-factor-core.php b/tests/class-two-factor-core.php index 1bc181d2..c844aec8 100644 --- a/tests/class-two-factor-core.php +++ b/tests/class-two-factor-core.php @@ -862,6 +862,42 @@ public function test_enabling_two_factor_is_factored_session() { } + /** + * Validate that disabling all providers results in a non-two-factor session. + * + * @covers Two_Factor_Core::current_user_can_update_two_factor_options() + * @covers Two_Factor_Core::user_two_factor_options_update() + */ + public function test_disabling_two_factor_is_not_factored_session() { + $user = self::factory()->user->create_and_get(); + + // Setup a 2FA session. + wp_set_current_user( $user->ID ); + wp_set_auth_cookie( $user->ID ); + + $key = '_nonce_user_two_factor_options'; + $nonce = wp_create_nonce( 'user_two_factor_options' ); + $_POST[ $key ] = $nonce; + $_REQUEST[ $key ] = $nonce; + + $_POST[ Two_Factor_Core::ENABLED_PROVIDERS_USER_META_KEY ] = [ 'Two_Factor_Dummy' => 'Two_Factor_Dummy' ]; + + Two_Factor_Core::user_two_factor_options_update( $user->ID ); + + $this->assertNotFalse( Two_Factor_Core::is_current_user_session_two_factor() ); + + // Disable all providers, and test that the session is invalidated. + $_POST[ Two_Factor_Core::ENABLED_PROVIDERS_USER_META_KEY ] = []; + Two_Factor_Core::user_two_factor_options_update( $user->ID ); + + $this->assertFalse( Two_Factor_Core::is_current_user_session_two_factor() ); + + $session = Two_Factor_Core::get_current_user_session(); + + $this->assertArrayNotHasKey( 'two-factor-login', $session ); + $this->assertArrayNotHasKey( 'two-factor-provider', $session ); + } + /** * Validate that a non-2fa login doesn't set the session two-factor data. * @@ -1083,4 +1119,144 @@ public function test_current_user_can_update_two_factor_options() { $this->assertFalse( Two_Factor_Core::current_user_can_update_two_factor_options() ); } + /** + * Test the fetch & update session methods. + * + * @covers Two_Factor_Core::get_current_user_session + * @covers Two_Factor_Core::update_current_user_session + */ + public function test_session_getter_setter() { + $user = $this->get_dummy_user( array( 'Two_Factor_Dummy' => 'Two_Factor_Dummy' ) ); + + // Fetch the session, it's not yet a thing. + $session = Two_Factor_Core::get_current_user_session(); + $this->assertFalse( $session ); + + // Set the cookie without going through two-factor, and fill in $_COOKIE, for session handler. + wp_set_auth_cookie( $user->ID ); + + // Fetch the session, check it exists. + $session = Two_Factor_Core::get_current_user_session(); + $this->assertNotEmpty( $session ); + + // Check setting keys works. + $this->assertArrayNotHasKey( 'test-key', $session ); + + // Set the key + Two_Factor_Core::update_current_user_session( array( + 'test-key' => true, + 'test-key-two' => true, + ) ); + + // Retrieve the session again, and verify it's updated. + $session = Two_Factor_Core::get_current_user_session(); + $this->assertNotEmpty( $session ); + $this->assertArrayHasKey( 'test-key', $session ); + $this->assertArrayHasKey( 'test-key-two', $session ); + + // Remove the key by setting it to null + Two_Factor_Core::update_current_user_session( array( + 'test-key' => null + ) ); + + // Check the key is no longer there. + $session = Two_Factor_Core::get_current_user_session(); + $this->assertNotEmpty( $session ); + $this->assertArrayNotHasKey( 'test-key', $session ); + $this->assertArrayHasKey( 'test-key-two', $session ); + } + + /** + * Test get_provider_for_user() + * + * @covers Two_Factor_Core::get_provider_for_user + */ + public function test_get_provider_for_user() { + $other_user = $this->get_dummy_user( array( 'Two_Factor_Dummy' => 'Two_Factor_Dummy' ) ); + $user = $this->get_dummy_user( array( 'Two_Factor_Dummy' => 'Two_Factor_Dummy' ) ); + + // Set the cookie without going through two-factor, and fill in $_COOKIE. + wp_set_auth_cookie( $user->ID ); + + // Setup the current session as 2fa'd + Two_Factor_Core::update_current_user_session( array( + 'two-factor-provider' => 'Two_Factor_Dummy', + 'two-factor-login' => time() + ) ); + + $dummy = Two_Factor_Dummy::get_instance(); + $email = Two_Factor_Email::get_instance(); + + // Ensure the provider returned is the primary for the user. + $this->assertEquals( + 'Two_Factor_Dummy', + // Using get_class() to verify the actual class, rather than the provider key. + get_class( Two_Factor_Core::get_provider_for_user( $user ) ) + ); + + // Validate that passing a specific provider that's not enabled, returns their primary provider. + $this->assertSame( + $dummy, + Two_Factor_Core::get_provider_for_user( $user, $email ) + ); + + // Validate that upon requesting an invalid provider, valid data comes back. + $this->assertEquals( + 'Two_Factor_Dummy', + Two_Factor_Core::get_provider_for_user( $user, new stdClass )->get_key() + ); + $this->assertEquals( + 'Two_Factor_Dummy', + Two_Factor_Core::get_provider_for_user( $user, false )->get_key() + ); + $this->assertEquals( + 'Two_Factor_Dummy', + Two_Factor_Core::get_provider_for_user( $user, 'Unknown_Provider' )->get_key() + ); + + // Validate that a provider not setup for the user is not returned. + $this->assertEquals( + 'Two_Factor_Dummy', + Two_Factor_Core::get_provider_for_user( $user, 'Two_Factor_Email' )->get_key() + ); + $this->assertEquals( + 'Two_Factor_Dummy', + Two_Factor_Core::get_provider_for_user( $user, $email )->get_key() + ); + + // Validate that even when another provider is set as primary, if the user is the current user it returns their last-used. + Two_Factor_Core::enable_provider_for_user( $user->ID, 'Two_Factor_Email' ); + $this->assertEquals( + 'Two_Factor_Dummy', + Two_Factor_Core::get_provider_for_user( $user )->get_key() + ); + + // Update the session to say that Email was last-used. + Two_Factor_Core::update_current_user_session( array( + 'two-factor-provider' => $email->get_key(), + ) ); + + // Validate it's now the default for the current session. + $this->assertEquals( + 'Two_Factor_Email', + Two_Factor_Core::get_provider_for_user( $user )->get_key() + ); + + // Ensure that passing a specific provider class in comes back out. + $this->assertSame( + $dummy, + Two_Factor_Core::get_provider_for_user( $user, $dummy ) + ); + $this->assertSame( + $email, + Two_Factor_Core::get_provider_for_user( $user, $email ) + ); + + // Validate that the current user session doesn't affect fetching it for a different user. + $this->assertEquals( + 'Two_Factor_Dummy', + Two_Factor_Core::get_provider_for_user( $other_user )->get_key() + ); + } + }