From 204bad46856796dc45d169e3d83c31574462567f Mon Sep 17 00:00:00 2001
From: Weston Ruter <westonruter@git.wordpress.org>
Date: Sun, 1 Nov 2015 06:39:50 +0000
Subject: [PATCH] Customize: Return user to referring URL when leaving
 Customizer in absence of `return` query param.

When referring URL is not available, default returning user to frontend URL instead of admin URL. Themes page is updated to include the `return` path in Customizer links.

Props McGuive7, westonruter.
Fixes #32637.


git-svn-id: https://develop.svn.wordpress.org/trunk@35483 602fd350-edb4-49c9-b593-d223f7449a82
---
 src/wp-admin/includes/theme.php                | 12 +++++++++++-
 src/wp-admin/themes.php                        |  2 +-
 src/wp-includes/class-wp-customize-manager.php |  7 ++++---
 tests/phpunit/tests/customize/manager.php      | 13 +++++++++++--
 4 files changed, 27 insertions(+), 7 deletions(-)

diff --git a/src/wp-admin/includes/theme.php b/src/wp-admin/includes/theme.php
index c4b8b0c3bf..0a0ea9e9c5 100644
--- a/src/wp-admin/includes/theme.php
+++ b/src/wp-admin/includes/theme.php
@@ -525,6 +525,16 @@ function wp_prepare_themes_for_js( $themes = null ) {
 			$parents[ $slug ] = $theme->parent()->get_stylesheet();
 		}
 
+		$customize_action = null;
+		if ( current_user_can( 'edit_theme_options' ) && current_user_can( 'customize' ) ) {
+			$customize_action = esc_url( add_query_arg(
+				array(
+					'return' => urlencode( esc_url_raw( wp_unslash( $_SERVER['REQUEST_URI'] ) ) ),
+				),
+				wp_customize_url( $slug )
+			) );
+		}
+
 		$prepared_themes[ $slug ] = array(
 			'id'           => $slug,
 			'name'         => $theme->display( 'Name' ),
@@ -540,7 +550,7 @@ function wp_prepare_themes_for_js( $themes = null ) {
 			'update'       => get_theme_update_available( $theme ),
 			'actions'      => array(
 				'activate' => current_user_can( 'switch_themes' ) ? wp_nonce_url( admin_url( 'themes.php?action=activate&amp;stylesheet=' . $encoded_slug ), 'switch-theme_' . $slug ) : null,
-				'customize' => ( current_user_can( 'edit_theme_options' ) && current_user_can( 'customize' ) ) ? wp_customize_url( $slug ) : null,
+				'customize' => $customize_action,
 				'delete'   => current_user_can( 'delete_themes' ) ? wp_nonce_url( admin_url( 'themes.php?action=delete&amp;stylesheet=' . $encoded_slug ), 'delete-theme_' . $slug ) : null,
 			),
 		);
diff --git a/src/wp-admin/themes.php b/src/wp-admin/themes.php
index edbe97ab3f..bf65884fb2 100644
--- a/src/wp-admin/themes.php
+++ b/src/wp-admin/themes.php
@@ -387,7 +387,7 @@ $can_install = current_user_can( 'install_themes' );
 
 	<# if ( data.active ) { #>
 		<# if ( data.actions.customize ) { #>
-			<a class="button button-primary customize load-customize hide-if-no-customize" href="{{ data.actions.customize }}"><?php _e( 'Customize' ); ?></a>
+			<a class="button button-primary customize load-customize hide-if-no-customize" href="{{{ data.actions.customize }}}"><?php _e( 'Customize' ); ?></a>
 		<# } #>
 	<# } else { #>
 		<a class="button button-secondary activate" href="{{{ data.actions.activate }}}"><?php _e( 'Activate' ); ?></a>
diff --git a/src/wp-includes/class-wp-customize-manager.php b/src/wp-includes/class-wp-customize-manager.php
index dd71ede82d..55acd10f56 100644
--- a/src/wp-includes/class-wp-customize-manager.php
+++ b/src/wp-includes/class-wp-customize-manager.php
@@ -1473,14 +1473,15 @@ final class WP_Customize_Manager {
 	 * @return string URL for link to close Customizer.
 	 */
 	public function get_return_url() {
+		$referer = wp_get_referer();
 		if ( $this->return_url ) {
 			$return_url = $this->return_url;
+		} else if ( $referer ) {
+			$return_url = $referer;
 		} else if ( $this->preview_url ) {
 			$return_url = $this->preview_url;
-		} else if ( current_user_can( 'edit_theme_options' ) || current_user_can( 'switch_themes' ) ) {
-			$return_url = admin_url( 'themes.php' );
 		} else {
-			$return_url = admin_url();
+			$return_url = home_url( '/' );
 		}
 		return $return_url;
 	}
diff --git a/tests/phpunit/tests/customize/manager.php b/tests/phpunit/tests/customize/manager.php
index 717b234324..8e6fe2a2a1 100644
--- a/tests/phpunit/tests/customize/manager.php
+++ b/tests/phpunit/tests/customize/manager.php
@@ -252,16 +252,25 @@ class Tests_WP_Customize_Manager extends WP_UnitTestCase {
 	 */
 	function test_return_url() {
 		wp_set_current_user( self::factory()->user->create( array( 'role' => 'author' ) ) );
-		$this->assertEquals( get_admin_url(), $this->manager->get_return_url() );
+		$this->assertEquals( home_url( '/' ), $this->manager->get_return_url() );
 
 		wp_set_current_user( self::factory()->user->create( array( 'role' => 'administrator' ) ) );
 		$this->assertTrue( current_user_can( 'edit_theme_options' ) );
-		$this->assertEquals( admin_url( 'themes.php' ), $this->manager->get_return_url() );
+		$this->assertEquals( home_url( '/' ), $this->manager->get_return_url() );
 
 		$preview_url = home_url( '/foo/' );
 		$this->manager->set_preview_url( $preview_url );
 		$this->assertEquals( $preview_url, $this->manager->get_return_url() );
 
+		$url = home_url( '/referred/' );
+		$_SERVER['HTTP_REFERER'] = wp_slash( $url );
+		$this->assertEquals( $url, $this->manager->get_return_url() );
+
+		$url = 'http://badreferer.example.com/';
+		$_SERVER['HTTP_REFERER'] = wp_slash( $url );
+		$this->assertNotEquals( $url, $this->manager->get_return_url() );
+		$this->assertEquals( $preview_url, $this->manager->get_return_url() );
+
 		$this->manager->set_return_url( admin_url( 'edit.php?trashed=1' ) );
 		$this->assertEquals( admin_url( 'edit.php' ), $this->manager->get_return_url() );
 	}