From dbee8261096062aa08cf49299eabb642f8ffdb7c Mon Sep 17 00:00:00 2001
From: Dion Hulse <contact@dd32.id.au>
Date: Mon, 20 Feb 2023 15:00:04 +1000
Subject: [PATCH] Use an anonymous function attached to a callback to set the
 user session information.

---
 class-two-factor-core.php | 24 ++++++++++++------------
 1 file changed, 12 insertions(+), 12 deletions(-)

diff --git a/class-two-factor-core.php b/class-two-factor-core.php
index 174276fb..f1acb590 100644
--- a/class-two-factor-core.php
+++ b/class-two-factor-core.php
@@ -1095,18 +1095,16 @@ public static function login_form_validate_2fa() {
 			$rememberme = true;
 		}
 
-		// Create a new User session
-		$expiration = time() + apply_filters( 'auth_cookie_expiration', ( $rememberme ? 14 : 2 ) * DAY_IN_SECONDS, $user->ID, $rememberme );
-		$manager    = WP_Session_Tokens::get_instance( $user->ID );
-		$token      = $manager->create( $expiration );
-		$session    = $manager->get( $token );
+		$session_information_callback = function( $session, $user_id ) use( $provider, $user ) {
+			if ( $user->ID === $user_id ) {
+				$session['two-factor-login']    = time();
+				$session['two-factor-provider'] = get_class( $provider );
+			}
 
-		// Append the Two Factor session data
-		$session['two-factor-login']    = time();
-		$session['two-factor-provider'] = get_class( $provider );
+			return $session;
+		};
 
-		// Save it in the session and create the cookie with it.
-		$manager->update( $token, $session );
+		add_filter( 'attach_session_information', $session_information_callback, 10, 2 );
 
 		/*
 		 * NOTE: This filter removal is not normally required, this is included for protection against
@@ -1115,9 +1113,11 @@ public static function login_form_validate_2fa() {
 		 */
 		remove_filter( 'send_auth_cookies', '__return_false', PHP_INT_MAX );
 
-		wp_set_auth_cookie( $user->ID, $rememberme, '', $token );
+		wp_set_auth_cookie( $user->ID, $rememberme );
+
+		do_action( 'two_factor_user_authenticated', $user, $provider );
 
-		do_action( 'two_factor_user_authenticated', $user, $provider, $token );
+		remove_filter( 'attach_session_information', $session_information_callback );
 
 		// Must be global because that's how login_header() uses it.
 		global $interim_login;