Adds way to redirect user for connect flow

class.jetpack.php

@@ -378,6 +378,14 @@ class Jetpack {
 	 */
 	public static $plugin_upgrade_lock_key = 'jetpack_upgrade_lock';
 
+	/**
+	 * Constant for login redirect key.
+	 *
+	 * @var string
+	 * @since 8.4.0
+	 */
+	public static $jetpack_redirect_login = 'jetpack_connect_login_redirect';
+
 	/**
 	 * Holds the singleton instance of this class
 	 *
@@ -641,6 +649,9 @@ private function __construct() {
 
 		add_action( 'jetpack_event_log', array( 'Jetpack', 'log' ), 10, 2 );
 
+		add_filter( 'login_url', array( $this, 'login_url' ), 10, 2 );
+		add_action( 'login_init', array( $this, 'login_init' ) );
+
 		add_filter( 'determine_current_user', array( $this, 'wp_rest_authenticate' ) );
 		add_filter( 'rest_authentication_errors', array( $this, 'wp_rest_authentication_errors' ) );
 
@@ -4019,6 +4030,45 @@ function plugin_action_links( $actions ) {
 		return array_merge( $jetpack_home, $actions );
 	}
 
+	/**
+	 * Filters the login URL to include the registration flow in case the user isn't logged in.
+	 *
+	 * @param string $login_url The wp-login URL.
+	 * @param string $redirect  URL to redirect users after logging in.
+	 * @since Jetpack 8.4
+	 * @return string
+	 */
+	public function login_url( $login_url, $redirect ) {
+		parse_str( wp_parse_url( $redirect, PHP_URL_QUERY ), $redirect_parts );
+		if ( ! empty( $redirect_parts[ self::$jetpack_redirect_login ] ) ) {
+			$login_url = add_query_arg( self::$jetpack_redirect_login, 'true', $login_url );
+		}
+		return $login_url;
+	}
+
+	/**
+	 * Redirects non-authenticated users to authenticate with Calypso if redirect flag is set.
+	 *
+	 * @since Jetpack 8.4
+	 */
+	public function login_init() {
+		// phpcs:ignore WordPress.Security.NonceVerification
+		if ( ! empty( $_GET[ self::$jetpack_redirect_login ] ) ) {
+			add_filter( 'allowed_redirect_hosts', array( &$this, 'allow_wpcom_environments' ) );
+			wp_safe_redirect(
+				add_query_arg(
+					array(
+						'forceInstall' => 1,
+						'url'          => rawurlencode( get_site_url() ),
+					),
+					// @todo provide way to go to specific calypso env.
+					self::get_calypso_host() . 'jetpack/connect'
+				)
+			);
+			exit;
+		}
+	}
+
 	/*
 	 * Registration flow:
 	 * 1 - ::admin_page_load() action=register
@@ -7120,6 +7170,28 @@ public static function get_calypso_env() {
 		return '';
 	}
 
+	/**
+	 * Returns the hostname with protocol for Calypso.
+	 * Used for developing Jetpack with Calypso.
+	 *
+	 * @since 8.4.0
+	 *
+	 * @return string Calypso host.
+	 */
+	public static function get_calypso_host() {
+		$calypso_env = self::get_calypso_env();
+		switch ( $calypso_env ) {
+			case 'development':
+				return 'http://calypso.localhost:3000/';
+			case 'wpcalypso':
+				return 'https://wpcalypso.wordpress.com/';
+			case 'horizon':
+				return 'https://horizon.wordpress.com/';
+			default:
+				return 'https://wordpress.com/';
+		}
+	}
+
 	/**
 	 * Checks whether or not TOS has been agreed upon.
 	 * Will return true if a user has clicked to register, or is already connected.

tests/php/general/test_class.jetpack.php

@@ -1208,4 +1208,78 @@ public function partner_code_provider() {
 		);
 	}
 
+	/**
+	 * Tests login URL only adds redirect param when redirect param is in original request.
+	 *
+	 * @since 8.4.0
+	 * @return void
+	 */
+	public function test_login_url_add_redirect() {
+		$login_url = wp_login_url( '/wp-admin' );
+		$this->assertFalse( strpos( $login_url, Jetpack::$jetpack_redirect_login ) );
+
+		$login_url = wp_login_url( '/wp-admin?' . Jetpack::$jetpack_redirect_login . '=true' );
+		parse_str( wp_parse_url( $login_url, PHP_URL_QUERY ), $login_parts );
+		$this->assertArraySubset( array( Jetpack::$jetpack_redirect_login => 'true' ), $login_parts, true );
+	}
+
+	/**
+	 * Tests login redirect sending users to Calypso when redirect param is set.
+	 *
+	 * @since 8.4.0
+	 * @return void
+	 */
+	public function test_login_init_redirect() {
+		tests_add_filter(
+			'wp_redirect',
+			function( $location ) {
+				$expected_location = add_query_arg(
+					array(
+						'forceInstall' => 1,
+						'url'          => rawurlencode( get_site_url() ),
+					),
+					'https://wordpress.com/jetpack/connect'
+				);
+				$this->assertEquals( $location, $expected_location );
+				throw new Exception(); // Cause an exception, as we don't want to run exit.
+			}
+		);
+
+		// Remove core filters that add headers.
+		remove_filter( 'login_init', 'wp_admin_headers' );
+		remove_filter( 'login_init', 'send_frame_options_header' );
+
+		// Run it once and no exception is thrown.
+		do_action( 'login_init' );
+
+		$this->expectException( Exception::class );
+		$_GET[ Jetpack::$jetpack_redirect_login ] = 'true';
+		do_action( 'login_init' ); // Now expect an exception.
+	}
+
+	/**
+	 * Tests getting the correct Calypso host.
+	 *
+	 * @since 8.4.0
+	 * @return void
+	 */
+	public function test_get_calypso_host() {
+		// No env.
+		$this->assertEquals( 'https://wordpress.com/', Jetpack::get_calypso_host() );
+
+		$_GET['calypso_env'] = 'development';
+		$this->assertEquals( 'http://calypso.localhost:3000/', Jetpack::get_calypso_host() );
+
+		$_GET['calypso_env'] = 'wpcalypso';
+		$this->assertEquals( 'https://wpcalypso.wordpress.com/', Jetpack::get_calypso_host() );
+
+		$_GET['calypso_env'] = 'horizon';
+		$this->assertEquals( 'https://horizon.wordpress.com/', Jetpack::get_calypso_host() );
+
+		$_GET['calypso_env'] = 'stage';
+		$this->assertEquals( 'https://wordpress.com/', Jetpack::get_calypso_host() );
+
+		$_GET['calypso_env'] = 'production';
+		$this->assertEquals( 'https://wordpress.com/', Jetpack::get_calypso_host() );
+	}
 } // end class