Jetpack Setup: Improve web connection flow for outdated plugin

Modules/Sources/Networking/Remote/JetpackConnectionRemote.swift

@@ -4,11 +4,12 @@ import Foundation
 ///
 public final class JetpackConnectionRemote: Remote {
     private let siteURL: String
-
+    private let network: Network
     private var accountConnectionURL: URL?
 
     public init(siteURL: String, network: Network) {
         self.siteURL = siteURL
+        self.network = network
         super.init(network: network)
     }
 
@@ -49,6 +50,41 @@ public final class JetpackConnectionRemote: Remote {
         enqueue(request, mapper: mapper, completion: completion)
     }
 
+    /// Registers Jetpack site connection by requesting the input URL while disabling automatic redirection,
+    /// and returns the URL in the requested redirection.
+    /// To simplify redirection manipulation, we'll use a `URLSession` here instead of `Network`.
+    ///
+    public func registerJetpackSiteConnection(with url: URL, completion: @escaping (Result<URL, Error>) -> Void) {
+
+        let configuration = URLSessionConfiguration.default
+        for cookie in network.session.configuration.httpCookieStorage?.cookies ?? [] {
+            configuration.httpCookieStorage?.setCookie(cookie)
+        }
+
+        let session = URLSession(configuration: configuration, delegate: self, delegateQueue: nil)
+        do {
+            let request = try URLRequest(url: url, method: .get)
+            let task = session.dataTask(with: request) { [weak self] data, response, error in
+                if let result = self?.accountConnectionURL {
+                    DispatchQueue.main.async {
+                        completion(.success(result))
+                    }
+                    return
+                }
+                // We don't expect any response here since we'll cancel the task as soon as a redirect request is received.
+                // So always complete with a failure here.
+                let returnedError = error ?? JetpackConnectionError.accountConnectionURLNotFound
+                DispatchQueue.main.async {
+                    completion(.failure(returnedError))
+                }
+                return
+            }
+            task.resume()
+        } catch {
+            completion(.failure(error))
+        }
+    }
+
     /// Fetches the connection state with the site's Jetpack for the authenticated user.
     ///
     public func fetchJetpackConnectionData(completion: @escaping (Result<JetpackConnectionData, Error>) -> Void) {
@@ -84,6 +120,24 @@ public final class JetpackConnectionRemote: Remote {
     }
 }
 
+// MARK: - URLSessionDataDelegate conformance
+//
+extension JetpackConnectionRemote: URLSessionDataDelegate {
+    public func urlSession(_ session: URLSession,
+                           task: URLSessionTask,
+                           willPerformHTTPRedirection response: HTTPURLResponse,
+                           newRequest request: URLRequest) async -> URLRequest? {
+        // Disables redirection if the request is to load the Jetpack account connection URL
+        if let url = request.url,
+            url.absoluteString.hasPrefix(Constants.jetpackAccountConnectionURL) {
+            accountConnectionURL = url
+            task.cancel()
+            return nil
+        }
+        return request
+    }
+}
+
 /// periphery: ignore - used in test module and on the UI layer
 public enum JetpackConnectionError: Error, Equatable {
     case malformedURL

Modules/Sources/Yosemite/Actions/JetpackConnectionAction.swift

@@ -13,7 +13,8 @@ public enum JetpackConnectionAction: Action {
     /// Updates Jetpack the plugin for the current site.
     case activateJetpackPlugin(completion: (Result<Void, Error>) -> Void)
     /// Fetches the URL used for setting up Jetpack connection.
-    case fetchJetpackConnectionURL(completion: (Result<URL, Error>) -> Void)
+    case fetchJetpackConnectionURL(usingApplicationPassword: Bool,
+                                   completion: (Result<URL, Error>) -> Void)
     /// Fetches connection state with the given site's Jetpack.
     case fetchJetpackConnectionData(completion: (Result<JetpackConnectionData, Error>) -> Void)
     /// Establishes site-level connection and returns WordPress.com blog ID.

Modules/Sources/Yosemite/Stores/JetpackConnectionStore.swift

@@ -41,8 +41,9 @@ public final class JetpackConnectionStore: DeauthenticatedStore {
             installJetpackPlugin(completion: completion)
         case .activateJetpackPlugin(let completion):
             activateJetpackPlugin(completion: completion)
-        case .fetchJetpackConnectionURL(let completion):
-            fetchJetpackConnectionURL(completion: completion)
+        case let .fetchJetpackConnectionURL(usingApplicationPassword, completion):
+            fetchJetpackConnectionURL(usingApplicationPassword: usingApplicationPassword,
+                                      completion: completion)
         case .fetchJetpackConnectionData(let completion):
             fetchJetpackConnectionData(completion: completion)
         case .registerSite(let completion):
@@ -88,8 +89,26 @@ private extension JetpackConnectionStore {
         })
     }
 
-    func fetchJetpackConnectionURL(completion: @escaping (Result<URL, Error>) -> Void) {
-        jetpackConnectionRemote?.fetchJetpackConnectionURL(completion: completion)
+    func fetchJetpackConnectionURL(usingApplicationPassword: Bool,
+                                   completion: @escaping (Result<URL, Error>) -> Void) {
+        guard !usingApplicationPassword else {
+            jetpackConnectionRemote?.fetchJetpackConnectionURL(completion: completion)
+            return
+        }
+        jetpackConnectionRemote?.fetchJetpackConnectionURL { [weak self] result in
+            guard let self else { return }
+            switch result {
+            case .success(let url):
+                // If we get the account connection URL, return it immediately.
+                if url.absoluteString.hasPrefix(Constants.jetpackAccountConnectionURL) {
+                    return completion(.success(url))
+                }
+                // Otherwise, request the url with redirection disabled and retrieve the URL in LOCATION header
+                self.jetpackConnectionRemote?.registerJetpackSiteConnection(with: url, completion: completion)
+            case .failure(let error):
+                completion(.failure(error))
+            }
+        }
     }
 
     func fetchJetpackConnectionData(completion: @escaping (Result<JetpackConnectionData, Error>) -> Void) {
@@ -152,3 +171,9 @@ private extension JetpackConnectionStore {
         self.accountRemote = remote
     }
 }
+
+private extension JetpackConnectionStore {
+    enum Constants {
+        static let jetpackAccountConnectionURL = "https://jetpack.wordpress.com/jetpack.authorize"
+    }
+}

Modules/Tests/YosemiteTests/Stores/JetpackConnectionStoreTests.swift

@@ -167,7 +167,7 @@ final class JetpackConnectionStoreTests: XCTestCase {
 
         // When
         let result: Result<URL, Error> = waitFor { promise in
-            let action = JetpackConnectionAction.fetchJetpackConnectionURL { result in
+            let action = JetpackConnectionAction.fetchJetpackConnectionURL(usingApplicationPassword: true) { result in
                 promise(result)
             }
             store.onAction(action)
@@ -192,7 +192,7 @@ final class JetpackConnectionStoreTests: XCTestCase {
 
         // When
         let result: Result<URL, Error> = waitFor { promise in
-            let action = JetpackConnectionAction.fetchJetpackConnectionURL { result in
+            let action = JetpackConnectionAction.fetchJetpackConnectionURL(usingApplicationPassword: true) { result in
                 promise(result)
             }
             store.onAction(action)

WooCommerce/Classes/Authentication/Jetpack Setup/LoginJetpackSetupCoordinator.swift

@@ -44,10 +44,14 @@ final class LoginJetpackSetupCoordinator: Coordinator {
 //
 private extension LoginJetpackSetupCoordinator {
     func showSetupSteps() {
+        guard let credentials = stores.sessionManager.defaultCredentials,
+              case .wpcom = credentials else {
+            fatalError("Unexpected error: No WPCom credentials found for setting up Jetpack")
+        }
         let setupUI = JetpackSetupHostingController(
             siteURL: siteURL,
             connectionOnly: connectionOnly,
-            wpcomCredentials: stores.sessionManager.defaultCredentials,
+            wpcomCredentials: credentials,
             onStoreNavigation: { [weak self] connectedEmail in
             guard let self, let email = connectedEmail else { return }
             if email != self.stores.sessionManager.defaultAccount?.email {

WooCommerce/Classes/Authentication/Jetpack Setup/Native Jetpack Setup/JetpackSetupView.swift

@@ -7,11 +7,13 @@ import protocol WooFoundation.Analytics
 final class JetpackSetupHostingController: UIHostingController<JetpackSetupView> {
     private let viewModel: JetpackSetupViewModel
     private let authentication: Authentication
-    private let wpcomCredentials: Credentials?
+    private let wpcomCredentials: Credentials
+
+    private var connectionWebView: UINavigationController?
 
     init(siteURL: String,
          connectionOnly: Bool,
-         wpcomCredentials: Credentials?,
+         wpcomCredentials: Credentials,
          stores: StoresManager = ServiceLocator.stores,
          authentication: Authentication = ServiceLocator.authenticationManager,
          analytics: Analytics = ServiceLocator.analytics,
@@ -27,7 +29,10 @@ final class JetpackSetupHostingController: UIHostingController<JetpackSetupView>
         super.init(rootView: JetpackSetupView(viewModel: viewModel))
 
         rootView.webViewPresentationHandler = { [weak self] in
-            self?.presentJetpackConnectionWebView()
+            guard let url = self?.viewModel.jetpackConnectionURL else {
+                return
+            }
+            self?.presentJetpackConnectionWebView(with: url)
         }
 
         rootView.supportHandler = { [weak self] in
@@ -64,21 +69,20 @@ final class JetpackSetupHostingController: UIHostingController<JetpackSetupView>
 
     @objc
     private func dismissWebView() {
-        dismiss(animated: true)
+        connectionWebView?.dismiss(animated: true)
+        connectionWebView = nil
     }
 
-    private func presentJetpackConnectionWebView() {
-        guard let connectionURL = viewModel.jetpackConnectionURL else {
-            return
-        }
-
-        let webViewModel = JetpackConnectionWebViewModel(initialURL: connectionURL,
+    private func presentJetpackConnectionWebView(with url: URL) {
+        let webViewModel = JetpackConnectionWebViewModel(initialURL: url,
                                                          siteURL: viewModel.siteURL,
                                                          completion: { [weak self] in
             guard let self else { return }
             self.viewModel.shouldPresentWebView = false
             self.viewModel.didAuthorizeJetpackConnection()
             self.dismissView()
+        }, onAuthorization: { [weak self] url in
+            self?.presentJetpackConnectionWebView(with: url)
         }, onFailure: { [weak self] errorCode in
             guard let self else { return }
             self.viewModel.shouldPresentWebView = false
@@ -88,13 +92,20 @@ final class JetpackSetupHostingController: UIHostingController<JetpackSetupView>
             guard let self else { return }
             self.viewModel.jetpackConnectionInterrupted = true
         })
+
         let webView = AuthenticatedWebViewController(viewModel: webViewModel, extraCredentials: wpcomCredentials)
         webView.navigationItem.leftBarButtonItem = UIBarButtonItem(title: Localization.cancel,
                                                                    style: .plain,
                                                                    target: self,
                                                                    action: #selector(self.dismissWebView))
-        let navigationController = UINavigationController(rootViewController: webView)
-        self.present(navigationController, animated: true)
+        if let connectionWebView {
+            /// Replace the web view to avoid unnecessary navigations
+            connectionWebView.viewControllers = [webView]
+        } else {
+            let navigationController = UINavigationController(rootViewController: webView)
+            self.present(navigationController, animated: true)
+            self.connectionWebView = navigationController
+        }
     }
 
     private func presentSupport() {

WooCommerce/Classes/Authentication/Jetpack Setup/Native Jetpack Setup/JetpackSetupViewModel.swift

@@ -15,7 +15,7 @@ final class JetpackSetupViewModel: ObservableObject {
 
     private let stores: StoresManager
     private let storeNavigationHandler: (_ connectedEmail: String?) -> Void
-    private let wpcomCredentials: Credentials?
+    private let wpcomCredentials: Credentials
     private var isPluginActivated = false
     private var connectionType = WooAnalyticsEvent.JetpackSetup.ConnectionType.native
 
@@ -95,7 +95,7 @@ final class JetpackSetupViewModel: ObservableObject {
 
     init(siteURL: String,
          connectionOnly: Bool,
-         wpcomCredentials: Credentials?,
+         wpcomCredentials: Credentials,
          stores: StoresManager = ServiceLocator.stores,
          analytics: Analytics = ServiceLocator.analytics,
          delayBeforeRetry: Double = Constants.delayBeforeRetry,
@@ -107,7 +107,7 @@ final class JetpackSetupViewModel: ObservableObject {
         self.analytics = analytics
         self.setupSteps = connectionOnly ? [.connection, .done] : JetpackInstallStep.allCases
         self.storeNavigationHandler = onStoreNavigation
-        self.siteConnectionURL = URL(string: String(format: Constants.jetpackInstallString, siteURL, Constants.mobileRedirectURL))
+        self.siteConnectionURL = URL(string: String(format: Constants.jetpackInstallString, siteURL))
         self.delayBeforeRetry = delayBeforeRetry
     }
 
@@ -264,7 +264,13 @@ private extension JetpackSetupViewModel {
         currentSetupStep = .connection
         connectionType = .web
         trackSetup()
-        let action = JetpackConnectionAction.fetchJetpackConnectionURL { [weak self] result in
+        let usingApplicationPassword: Bool = {
+            if case .some(.applicationPassword) = stores.sessionManager.defaultCredentials {
+                return true
+            }
+            return false
+        }()
+        let action = JetpackConnectionAction.fetchJetpackConnectionURL(usingApplicationPassword: usingApplicationPassword) { [weak self] result in
             guard let self else { return }
             switch result {
             case .success(let url):
@@ -428,11 +434,6 @@ private extension JetpackSetupViewModel {
     }
 
     func finalizeSiteConnection(blogID: Int64, provisionResponse: JetpackConnectionProvisionResponse) {
-        guard let wpcomCredentials, case .wpcom = wpcomCredentials else {
-            /// WPCom credentials are necessary to finalize connection through API
-            /// If this is unavailable, fall back to the web flow.
-            return startConnectionWithWebView()
-        }
         let network = AlamofireNetwork(credentials: wpcomCredentials)
         stores.dispatch(JetpackConnectionAction.finalizeConnection(
             siteID: blogID,
@@ -558,8 +559,7 @@ extension JetpackSetupViewModel {
         static let errorCodeNoWPComUser = 99
         static let errorUserInfoReason = "reason"
         static let errorUserInfoNoWPComUser = "No connected WP.com user found"
-        static let jetpackInstallString = "https://wordpress.com/jetpack/connect?url=%@&mobile_redirect=%@&from=mobile"
-        static let mobileRedirectURL = "woocommerce://jetpack-connected"
+        static let jetpackInstallString = "%@/wp-admin/admin.php?page=jetpack"
         static let accountConnectionURL = "https://jetpack.wordpress.com/jetpack.authorize"
     }
 }

WooCommerce/Classes/Authentication/Navigation Exceptions/JetpackConnectionWebViewModel.swift

@@ -14,6 +14,7 @@ final class JetpackConnectionWebViewModel: AuthenticatedWebViewModel {
     /// Failure handler with an optional error code if available
     let failureHandler: (Int?) -> Void
     let dismissalHandler: () -> Void
+    let authorizationHandler: (URL) -> Void
 
     private let stores: StoresManager
     private let analytics: Analytics
@@ -25,6 +26,7 @@ final class JetpackConnectionWebViewModel: AuthenticatedWebViewModel {
          stores: StoresManager = ServiceLocator.stores,
          analytics: Analytics = ServiceLocator.analytics,
          completion: @escaping () -> Void,
+         onAuthorization: @escaping (URL) -> Void = { _ in },
          onFailure: @escaping (Int?) -> Void = { _ in },
          onDismissal: @escaping () -> Void = {}) {
         self.title = title
@@ -35,6 +37,7 @@ final class JetpackConnectionWebViewModel: AuthenticatedWebViewModel {
         self.completionHandler = completion
         self.failureHandler = onFailure
         self.dismissalHandler = onDismissal
+        self.authorizationHandler = onAuthorization
     }
 
     func handleDismissal() {
@@ -56,7 +59,15 @@ final class JetpackConnectionWebViewModel: AuthenticatedWebViewModel {
 
     func decidePolicy(for navigationURL: URL) async -> WKNavigationActionPolicy {
         let url = navigationURL.absoluteString
-        if handleCompletionIfPossible(url) {
+        if url.contains(Constants.authorizationURL),
+            initialURL?.absoluteString.contains(Constants.authorizationURL) == false {
+            await MainActor.run { [weak self] in
+                guard let self else { return }
+                shouldIgnoreDismissalHandling = true
+                authorizationHandler(navigationURL)
+            }
+            return .cancel
+        } else if handleCompletionIfPossible(url) {
             return .cancel
         }
         return .allow
@@ -104,6 +115,7 @@ private extension JetpackConnectionWebViewModel {
     enum Constants {
         static let mobileRedirectURL = "woocommerce://jetpack-connected"
         static let plansPage = "https://wordpress.com/jetpack/connect/plans"
+        static let authorizationURL = "jetpack.wordpress.com/jetpack.authorize"
     }
 
     enum Localization {