sb-woocommerce-email-verification_return-to-original-page.diff

diff --git a/sb-woocommerce-email-verification.php b/sb-woocommerce-email-verification.php
index acf3844..b741676 100755
--- a/sb-woocommerce-email-verification.php
+++ b/sb-woocommerce-email-verification.php
@@ -3,7 +3,7 @@
 	Plugin Name: SB WooCommerce Email Verification
 	Description: Force customers to verify account before login. Reduce spam users.
 	Plugin URI: http://codecanyon.net/user/sbthemes/portfolio?ref=sbthemes
-	Version: 1.3
+	Version: 1.4.1
 	Author: SB Themes
 	Author URI: http://codecanyon.net/user/sbthemes/portfolio?ref=sbthemes
 */
@@ -15,7 +15,7 @@ require_once('admin/sb-admin-panel.php');
 
 class SB_WooCommerce_Email_Verification {
 
-	public $plugin_version 			= '1.3';
+	public $plugin_version 			= '1.4.1';
 	public $db_version 				= '1.0';
 	public $plugin_name		 		= 'WooCommerce Email Verification';
 	public $menu_text 				= 'WooCommerce Email Verification';
@@ -101,11 +101,21 @@ class SB_WooCommerce_Email_Verification {
 	function sb_user_autologout() {
 		if(is_user_logged_in()) {
 			wp_clear_auth_cookie();
-			wp_redirect(get_permalink(get_option('woocommerce_myaccount_page_id')).'?registered=true');
+			wp_redirect(self::verification_link_base().'registered=true');
 			die;
 		}
 	}
 
+	//Uses referrer link if available, or woocommerce myaccount page
+	public static function verification_link_base() {
+		$link_base = isset( $_SERVER['HTTP_REFERER'] ) ? $_SERVER['HTTP_REFERER'] : get_permalink( get_option('woocommerce_myaccount_page_id') );
+		$link_base = str_replace('customer-logout/', '', $link_base);
+		$link_base = preg_replace('/_wpnonce=[a-f,0-9]+/', '&', $link_base); // remove old nonce
+		$link_base = rtrim(preg_replace('/&+/', '&', $link_base),'&?');
+		$link_base .= strpos('?', $link_base) !== false ? '&' : '?';
+		return $link_base;
+	}
+
 	//Auto logout after registration on checkout
 	function sb_user_checkout_autologout() {
 		if(is_user_logged_in()) {
@@ -136,7 +146,14 @@ class SB_WooCommerce_Email_Verification {
 						$wpdb->query($wpdb->prepare("update ".$wpdb->prefix."users set user_activation_key = %s where ID = %d", $key, $user->ID));
 					}
 				}
-				wp_redirect(get_permalink(get_option('woocommerce_myaccount_page_id')).'?verify='.$verify);
+
+				$link_base = '//' . $_SERVER['HTTP_HOST'] . str_replace('action=verify_account','', $_SERVER['REQUEST_URI']);
+				$link_base = str_replace('user_login='.$_GET['user_login'], '', $link_base);
+				$link_base = str_replace('key='.$_GET['key'], '', $link_base);
+				$link_base = rtrim(preg_replace('/&+/', '&', $link_base),'&?');
+				$link_base .= strpos('?', $link_base) !== false ? '&' : '?';
+
+				wp_redirect($link_base.'verify='.$verify);
 				die;
 			}
 		}
@@ -154,7 +171,7 @@ class SB_WooCommerce_Email_Verification {
 		}
 	}
 
-	//Registratin success message
+	//Registration success message
 	function sb_reg_success_message() {
 		if(isset($_GET['registered']) && $_GET['registered'] == 'true') {
 			$settings = $this->sb_admin->get_settings(); ?>
@@ -181,7 +198,7 @@ class SB_WooCommerce_Email_Verification {
 		if(count($common_values) == 0) {
 			$is_approved = get_user_meta($userdata->ID, 'sbwev-approve-user', true);
 			if(!$is_approved || $is_approved == 0) {
-				$regenerate_link = get_permalink(get_option('woocommerce_myaccount_page_id')).'?action=resend_key&user_login='.$userdata->user_login.'&nonce='.wp_create_nonce('resend_key'.$userdata->user_login);
+				$regenerate_link = self::verification_link_base().'action=resend_key&user_login='.$userdata->user_login.'&nonce='.wp_create_nonce('resend_key'.$userdata->user_login);
 				$userdata = new WP_Error(
 					'sb_confirmation_error',
 					__( str_replace('{{resend_verification_link}}', '<a href="'.$regenerate_link.'">'.$settings['resend_link_text'].'</a>', $settings['fail_login']), 'woocommerce' )
@@ -242,7 +259,7 @@ class SB_WooCommerce_Email_Verification {
 							}
 
 						}
-						wp_redirect(get_permalink(get_option('woocommerce_myaccount_page_id')).'?resend='.$resend);
+						wp_redirect(self::verification_link_base().'resend='.$resend);
 						die;
 					}
 				}
@@ -340,17 +357,11 @@ function get_verification_link($user_login) {
 	global $wpdb;
 	$key = wp_generate_password(20, false);
 	$wpdb->query($wpdb->prepare("update ".$wpdb->prefix."users set user_activation_key = %s where user_login = %s", $key, $user_login));
-	return $verification_link = get_permalink(get_option('woocommerce_myaccount_page_id')).'?action=verify_account&user_login='.urlencode($user_login).'&key='.urlencode($key);
-}
-
-
-
-
-
-
-
 
+	$verification_link = SB_WooCommerce_Email_Verification::verification_link_base().'action=verify_account&user_login='.urlencode($user_login).'&key='.urlencode($key);
 
+	return $verification_link;
+}