Jetpack Setup: Update connection step to use APIs where possible

Modules/Sources/Networking/Mapper/JetpackConnectionProvisionMapper.swift

@@ -18,4 +18,11 @@ public struct JetpackConnectionProvisionResponse: Decodable {
     public let userId: Int64
     public let scope: String
     public let secret: String
+
+    /// periphery: ignore - used in UI module
+    public init(userId: Int64, scope: String, secret: String) {
+        self.userId = userId
+        self.scope = scope
+        self.secret = secret
+    }
 }

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

@@ -44,7 +44,11 @@ final class LoginJetpackSetupCoordinator: Coordinator {
 //
 private extension LoginJetpackSetupCoordinator {
     func showSetupSteps() {
-        let setupUI = JetpackSetupHostingController(siteURL: siteURL, connectionOnly: connectionOnly, onStoreNavigation: { [weak self] connectedEmail in
+        let setupUI = JetpackSetupHostingController(
+            siteURL: siteURL,
+            connectionOnly: connectionOnly,
+            wpcomCredentials: stores.sessionManager.defaultCredentials,
+            onStoreNavigation: { [weak self] connectedEmail in
             guard let self, let email = connectedEmail else { return }
             if email != self.stores.sessionManager.defaultAccount?.email {
                 // if the user authorized Jetpack with a different account, support them to log in with that account.

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

@@ -7,22 +7,23 @@ import protocol WooFoundation.Analytics
 final class JetpackSetupHostingController: UIHostingController<JetpackSetupView> {
     private let viewModel: JetpackSetupViewModel
     private let authentication: Authentication
-    private let connectionWebViewCredentials: Credentials?
+    private let wpcomCredentials: Credentials?
 
     init(siteURL: String,
          connectionOnly: Bool,
-         connectionWebViewCredentials: Credentials? = nil,
+         wpcomCredentials: Credentials?,
          stores: StoresManager = ServiceLocator.stores,
          authentication: Authentication = ServiceLocator.authenticationManager,
          analytics: Analytics = ServiceLocator.analytics,
          onStoreNavigation: @escaping (String?) -> Void) {
         self.viewModel = JetpackSetupViewModel(siteURL: siteURL,
                                                connectionOnly: connectionOnly,
+                                               wpcomCredentials: wpcomCredentials,
                                                stores: stores,
                                                analytics: analytics,
                                                onStoreNavigation: onStoreNavigation)
         self.authentication = authentication
-        self.connectionWebViewCredentials = connectionWebViewCredentials
+        self.wpcomCredentials = wpcomCredentials
         super.init(rootView: JetpackSetupView(viewModel: viewModel))
 
         rootView.webViewPresentationHandler = { [weak self] in
@@ -92,7 +93,7 @@ final class JetpackSetupHostingController: UIHostingController<JetpackSetupView>
             guard let self else { return }
             self.viewModel.jetpackConnectionInterrupted = true
         })
-        let webView = AuthenticatedWebViewController(viewModel: webViewModel, extraCredentials: connectionWebViewCredentials)
+        let webView = AuthenticatedWebViewController(viewModel: webViewModel, extraCredentials: wpcomCredentials)
         webView.navigationItem.leftBarButtonItem = UIBarButtonItem(title: Localization.cancel,
                                                                    style: .plain,
                                                                    target: self,
@@ -327,10 +328,3 @@ private extension JetpackSetupView {
         static let interruptedConnectionActionHandlerDelayTime: Double = 0.3
     }
 }
-
-struct JetpackSetupView_Previews: PreviewProvider {
-    static var previews: some View {
-        JetpackSetupView(viewModel: JetpackSetupViewModel(siteURL: "https://test.com", connectionOnly: true))
-        JetpackSetupView(viewModel: JetpackSetupViewModel(siteURL: "https://test.com", connectionOnly: false))
-    }
-}

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

@@ -3,6 +3,7 @@ import UIKit
 import Yosemite
 import enum Alamofire.AFError
 import enum Networking.NetworkError
+import class Networking.AlamofireNetwork
 import protocol WooFoundation.Analytics
 
 /// View model for `JetpackSetupView`.
@@ -14,6 +15,7 @@ final class JetpackSetupViewModel: ObservableObject {
 
     private let stores: StoresManager
     private let storeNavigationHandler: (_ connectedEmail: String?) -> Void
+    private let wpcomCredentials: Credentials?
 
     @Published private(set) var setupSteps: [JetpackInstallStep]
 
@@ -101,12 +103,14 @@ final class JetpackSetupViewModel: ObservableObject {
 
     init(siteURL: String,
          connectionOnly: Bool,
+         wpcomCredentials: Credentials?,
          stores: StoresManager = ServiceLocator.stores,
          analytics: Analytics = ServiceLocator.analytics,
          delayBeforeRetry: Double = Constants.delayBeforeRetry,
          onStoreNavigation: @escaping (String?) -> Void = { _ in}) {
         self.siteURL = siteURL
         self.connectionOnly = connectionOnly
+        self.wpcomCredentials = wpcomCredentials
         self.stores = stores
         self.analytics = analytics
         self.setupSteps = connectionOnly ? [.connection, .done] : JetpackInstallStep.allCases
@@ -138,14 +142,14 @@ final class JetpackSetupViewModel: ObservableObject {
 
     func startSetup() {
         if connectionOnly {
-            fetchJetpackConnectionURL()
+            checkJetpackConnection(afterConnection: false)
         } else {
             retrieveJetpackPluginDetails()
         }
     }
 
     func didAuthorizeJetpackConnection() {
-        checkJetpackConnection()
+        checkJetpackConnection(afterConnection: true)
     }
 
     func didEncounterErrorDuringConnection(code: Int?) {
@@ -217,7 +221,7 @@ private extension JetpackSetupViewModel {
                 if plugin.status == .inactive {
                     self.activateJetpack()
                 } else {
-                    self.fetchJetpackConnectionURL()
+                    self.checkJetpackConnection(afterConnection: false)
                 }
             case .failure(let error):
                 DDLogError("⛔️ Error retrieving Jetpack: \(error)")
@@ -268,7 +272,7 @@ private extension JetpackSetupViewModel {
             switch result {
             case .success:
                 self.trackSetupDuringLogin(.loginJetpackSetupActivationSuccessful)
-                self.fetchJetpackConnectionURL()
+                self.checkJetpackConnection(afterConnection: false)
             case .failure(let error):
                 self.trackSetupDuringLogin(.loginJetpackSetupActivationFailed, failure: error)
                 self.trackSetupAfterLogin(failure: error)
@@ -280,6 +284,9 @@ private extension JetpackSetupViewModel {
         stores.dispatch(action)
     }
 
+    /// Jetpack connection flow using web view.
+    /// Used only for sites with Jetpack plugin versions lower than 14.4.
+    ///
     func fetchJetpackConnectionURL() {
         currentSetupStep = .connection
         trackSetupAfterLogin()
@@ -308,50 +315,6 @@ private extension JetpackSetupViewModel {
         stores.dispatch(action)
     }
 
-    func checkJetpackConnection(retryCount: Int = 0) {
-        guard retryCount <= Constants.maxRetryCount else {
-            setupFailed = true
-            if let setupError {
-                analytics.track(.loginJetpackSetupErrorCheckingJetpackConnection, withError: setupError)
-            }
-            return
-        }
-        currentConnectionStep = .inProgress
-        let action = JetpackConnectionAction.fetchJetpackConnectionData { [weak self] result in
-            guard let self else { return }
-            switch result {
-            case .success(let connectionData):
-                guard let connectedEmail = connectionData.currentUser.wpcomUser?.email else {
-                    DDLogWarn("⚠️ Cannot find connected WPcom user")
-                    let missingWpcomUserError = NSError(domain: Constants.errorDomain,
-                                                        code: Constants.errorCodeNoWPComUser,
-                                                        userInfo: [Constants.errorUserInfoReason: Constants.errorUserInfoNoWPComUser])
-                    self.setupError = missingWpcomUserError
-                    self.trackSetupDuringLogin(.loginJetpackSetupCannotFindWPCOMUser, failure: missingWpcomUserError)
-                    // Retry fetching user in case Jetpack sync takes some time.
-                    DispatchQueue.main.asyncAfter(deadline: .now() + delayBeforeRetry) { [weak self] in
-                        self?.checkJetpackConnection(retryCount: retryCount + 1)
-                    }
-                    return
-                }
-
-                self.jetpackConnectedEmail = connectedEmail
-                self.currentConnectionStep = .authorized
-                self.currentSetupStep = .done
-
-                self.trackSetupDuringLogin(.loginJetpackSetupAllStepsMarkedDone)
-                self.trackSetupAfterLogin()
-            case .failure(let error):
-                DDLogError("⛔️ Error checking Jetpack connection: \(error)")
-                self.setupError = error
-                DispatchQueue.main.asyncAfter(deadline: .now() + delayBeforeRetry) { [weak self] in
-                    self?.checkJetpackConnection(retryCount: retryCount + 1)
-                }
-            }
-        }
-        stores.dispatch(action)
-    }
-
     func updateErrorMessage() {
         guard let setupErrorCode else {
             setupErrorDetail = .init(setupErrorMessage: Localization.genericErrorMessage,
@@ -377,6 +340,158 @@ private extension JetpackSetupViewModel {
     }
 }
 
+// MARK: Handle connection steps
+// Ref: pe5sF9-401-p2
+private extension JetpackSetupViewModel {
+    func checkJetpackConnection(afterConnection: Bool, retryCount: Int = 0) {
+        guard retryCount <= Constants.maxRetryCount else {
+            return didFailJetpackConnection()
+        }
+        currentConnectionStep = .inProgress
+        let action = JetpackConnectionAction.fetchJetpackConnectionData { [weak self] result in
+            guard let self else { return }
+            switch result {
+            case .success(let connectionData):
+                if afterConnection {
+                    checkConnectedUser(data: connectionData, retryCount: retryCount)
+                } else {
+                    handleJetpackConnectionData(connectionData)
+                }
+            case .failure(let error):
+                DDLogError("⛔️ Error checking Jetpack connection: \(error)")
+                self.setupError = error
+                DispatchQueue.main.asyncAfter(deadline: .now() + delayBeforeRetry) { [weak self] in
+                    self?.checkJetpackConnection(afterConnection: afterConnection, retryCount: retryCount + 1)
+                }
+            }
+        }
+        stores.dispatch(action)
+    }
+
+    func checkConnectedUser(data: JetpackConnectionData, retryCount: Int = 0) {
+        let connectedEmail = data.currentUser.wpcomUser?.email
+        if let connectedEmail {
+            return didCompleteJetpackConnection(connectedEmail: connectedEmail)
+        }
+
+        DDLogWarn("⚠️ Cannot find connected WPcom user")
+        let missingWpcomUserError = NSError(domain: Constants.errorDomain,
+                                            code: Constants.errorCodeNoWPComUser,
+                                            userInfo: [Constants.errorUserInfoReason: Constants.errorUserInfoNoWPComUser])
+        setupError = missingWpcomUserError
+        trackSetupDuringLogin(.loginJetpackSetupCannotFindWPCOMUser, failure: missingWpcomUserError)
+        // Retry fetching user in case Jetpack sync takes some time.
+        DispatchQueue.main.asyncAfter(deadline: .now() + delayBeforeRetry) { [weak self] in
+            self?.checkJetpackConnection(afterConnection: true, retryCount: retryCount + 1)
+        }
+    }
+
+    func handleJetpackConnectionData(_ data: JetpackConnectionData) {
+        if let connectedEmail = data.currentUser.wpcomUser?.email {
+            return didCompleteJetpackConnection(connectedEmail: connectedEmail)
+        }
+
+        if let isRegistered = data.isRegistered {
+            return handleSiteRegisterResult(isRegistered: isRegistered, blogID: data.blogID)
+        }
+
+        /// Check site info if `isRegistered` is unavailable.
+        stores.dispatch(WordPressSiteAction.fetchSiteInfo(siteURL: siteURL, completion: { [weak self] result in
+            guard let self else { return }
+            switch result {
+            case .success(let site):
+                if site.isJetpackThePluginInstalled {
+                    /// `isRegistered` is unavailable due to outdated Jetpack. Proceed with web flow.
+                    fetchJetpackConnectionURL()
+                } else {
+                    /// For Jetpack-connected sites, `isRegistered` is not returned. Check for `connectionOwner` instead.
+                    handleSiteRegisterResult(isRegistered: data.connectionOwner != nil, blogID: data.blogID)
+                }
+            case .failure(let error):
+                DDLogWarn("⛔️ Cannot fetch site info")
+                setupError = error
+                didFailJetpackConnection()
+            }
+        }))
+    }
+
+    func handleSiteRegisterResult(isRegistered: Bool, blogID: Int64?) {
+        if let blogID, isRegistered {
+            provisionSiteConnection(blogID: blogID)
+        } else {
+            registerSiteConnection()
+        }
+    }
+
+    func registerSiteConnection() {
+        stores.dispatch(JetpackConnectionAction.registerSite(completion: { [weak self] result in
+            guard let self else { return }
+            switch result {
+            case .success(let blogID):
+                provisionSiteConnection(blogID: blogID)
+            case .failure(let error):
+                setupError = error
+                didFailJetpackConnection()
+            }
+        }))
+    }
+
+    func provisionSiteConnection(blogID: Int64) {
+        currentSetupStep = .connection
+        trackSetupAfterLogin()
+        stores.dispatch(JetpackConnectionAction.provisionConnection(completion: { [weak self] result in
+            guard let self else { return }
+            switch result {
+            case .success(let response):
+                finalizeSiteConnection(blogID: blogID, provisionResponse: response)
+            case .failure(let error):
+                setupError = error
+                didFailJetpackConnection()
+            }
+        }))
+    }
+
+    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 fetchJetpackConnectionURL()
+        }
+        let network = AlamofireNetwork(credentials: wpcomCredentials)
+        stores.dispatch(JetpackConnectionAction.finalizeConnection(
+            siteID: blogID,
+            siteURL: siteURL,
+            provisionResponse: provisionResponse,
+            network: network
+        ) { [weak self] result in
+            guard let self else { return }
+            switch result {
+            case .success:
+                checkJetpackConnection(afterConnection: true)
+            case .failure(let error):
+                setupError = error
+                didFailJetpackConnection()
+            }
+        })
+    }
+
+    func didCompleteJetpackConnection(connectedEmail: String) {
+        jetpackConnectedEmail = connectedEmail
+        currentConnectionStep = .authorized
+        currentSetupStep = .done
+
+        trackSetupDuringLogin(.loginJetpackSetupAllStepsMarkedDone)
+        trackSetupAfterLogin()
+    }
+
+    func didFailJetpackConnection() {
+        setupFailed = true
+        if let setupError {
+            analytics.track(.loginJetpackSetupErrorCheckingJetpackConnection, withError: setupError)
+        }
+    }
+}
+
 // MARK: Subtypes
 //
 extension JetpackSetupViewModel {

WooCommerce/Classes/ViewRelated/JetpackSetup/JetpackSetupCoordinator.swift

@@ -222,7 +222,7 @@ private extension JetpackSetupCoordinator {
         }
         let setupUI = JetpackSetupHostingController(siteURL: site.url,
                                                     connectionOnly: requiresConnectionOnly,
-                                                    connectionWebViewCredentials: credentials,
+                                                    wpcomCredentials: credentials,
                                                     onStoreNavigation: { [weak self] _ in
             DDLogInfo("🎉 Jetpack setup completes!")
             self?.rootViewController.topmostPresentedViewController.dismiss(animated: true, completion: {
@@ -320,12 +320,7 @@ private extension JetpackSetupCoordinator {
 
     @MainActor
     func fetchJetpackConnectionData() async throws -> JetpackConnectionData {
-        /// Jetpack setup will fail anyway without admin role, so check that first.
-        let roles = stores.sessionManager.defaultRoles
-        guard roles.contains(.administrator) else {
-            throw JetpackCheckError.missingPermission
-        }
-        return try await withCheckedThrowingContinuation { continuation in
+        try await withCheckedThrowingContinuation { continuation in
             let action = JetpackConnectionAction.fetchJetpackConnectionData { result in
                 continuation.resume(with: result)
             }
@@ -489,11 +484,13 @@ private extension JetpackSetupCoordinator {
 }
 
 // MARK: - Subtypes
-private extension JetpackSetupCoordinator {
+extension JetpackSetupCoordinator {
     enum JetpackCheckError: Int, Error {
         case missingPermission = 403
     }
+}
 
+private extension JetpackSetupCoordinator {
     enum Constants {
         static let magicLinkUrlHostname = "magic-login"
     }