Improve authentication logic for authenticated web view (#15164)

itsmeichigo · web-flow · commit 6e32098b2163 · 2025-02-20T16:29:19.000+07:00
diff --git a/RELEASE-NOTES.txt b/RELEASE-NOTES.txt
@@ -11,6 +11,7 @@
 - [*] Background image upload: Fix missing error notice in iPhones [https://github.com/woocommerce/woocommerce-ios/pull/15117]
 - [*] Background image upload: Show a notice when the user leaves product details while uploads are pending [https://github.com/woocommerce/woocommerce-ios/pull/15134]
 - [*] Filters applied in product selector no longer affect the main product list screen. [https://github.com/woocommerce/woocommerce-ios/pull/14764]
+- [Internal] Improve authentication logic for authenticated web view [https://github.com/woocommerce/woocommerce-ios/pull/15164]
 - [*] Better error messages if Application Password login is disabled on user's website. [https://github.com/woocommerce/woocommerce-ios/pull/15031]
 - [**] Product Images: Update error handling [https://github.com/woocommerce/woocommerce-ios/pull/15105]
 - [*] Merchants can mark and filter favorite products for quicker access. [https://github.com/woocommerce/woocommerce-ios/pull/14597]
diff --git a/WooCommerce/Classes/Authentication/AuthenticatedWebViewController.swift b/WooCommerce/Classes/Authentication/AuthenticatedWebViewController.swift
@@ -1,15 +1,17 @@
 import Combine
 import UIKit
 import WebKit
+import Yosemite
 import class Networking.UserAgent
 import struct WordPressAuthenticator.WordPressOrgCredentials
-import enum Yosemite.Credentials
 
 /// A web view which is authenticated for WordPress.com, when possible.
 ///
 final class AuthenticatedWebViewController: UIViewController {
 
+    private let currentSite: Site?
     private let viewModel: AuthenticatedWebViewModel
+    private let authenticationFlow: WebViewAuthenticationFlow
 
     private lazy var activityIndicator: UIActivityIndicatorView = {
         let activityIndicator = UIActivityIndicatorView(style: .medium)
@@ -45,12 +47,15 @@ final class AuthenticatedWebViewController: UIViewController {
 
     private let wpcomCredentials: Credentials?
 
+    private var isFirstNavigation = true
 
-    init(viewModel: AuthenticatedWebViewModel, extraCredentials: Credentials? = nil) {
+    init(stores: StoresManager = ServiceLocator.stores,
+         viewModel: AuthenticatedWebViewModel,
+         extraCredentials: Credentials? = nil) {
         self.viewModel = viewModel
-        let currentCredentials = ServiceLocator.stores.sessionManager.defaultCredentials
+        let currentCredentials = stores.sessionManager.defaultCredentials
 
-        self.siteCredentials = {
+        let siteCredentials: WordPressOrgCredentials? = {
             if case let .wporg(username, password, siteAddress) = extraCredentials {
                 return WordPressOrgCredentials(username: username,
                                                password: password,
@@ -65,14 +70,29 @@ final class AuthenticatedWebViewController: UIViewController {
             return nil
         }()
 
-        self.wpcomCredentials = {
+        let wpcomCredentials: Credentials? = {
             if case .wpcom = extraCredentials {
                 return extraCredentials
             } else if case .wpcom = currentCredentials {
                 return currentCredentials
             }
             return nil
         }()
+
+        let currentSite = stores.sessionManager.defaultSite
+
+        self.authenticationFlow = {
+            guard let currentSite else {
+                return WebViewAuthenticationFlow.none
+            }
+            return viewModel.authenticationFlow(currentSite: currentSite,
+                                                wpcomCredentialsAvailable: wpcomCredentials != nil,
+                                                wporgCredentialsAvailable: siteCredentials != nil)
+        }()
+        self.currentSite = currentSite
+        self.wpcomCredentials = wpcomCredentials
+        self.siteCredentials = siteCredentials
+
         super.init(nibName: nil, bundle: nil)
 
         if let initialURL = viewModel.initialURL,
@@ -93,6 +113,7 @@ final class AuthenticatedWebViewController: UIViewController {
         configureWebView()
         configureActivityIndicator()
         configureProgressBar()
+        observeWebView()
         startLoading()
     }
 
@@ -144,7 +165,7 @@ private extension AuthenticatedWebViewController {
         ])
     }
 
-    func startLoading() {
+    func observeWebView() {
         webView.publisher(for: \.estimatedProgress)
             .sink { [weak self] progress in
                 if progress == 1 {
@@ -157,42 +178,98 @@ private extension AuthenticatedWebViewController {
 
         webView.publisher(for: \.url)
             .sink { [weak self] url in
-                guard let self else { return }
-                let initialURL = self.viewModel.initialURL
-                // avoids infinite loop if the initial url happens to be the nonce retrieval path.
-                if url?.absoluteString.contains(WKWebView.wporgNoncePath) == true,
-                   initialURL?.absoluteString.contains(WKWebView.wporgNoncePath) != true {
-                    self.loadContent()
-                } else {
-                    self.viewModel.handleRedirect(for: url)
-                }
+                guard let url else { return }
+                self?.handleRedirect(for: url)
             }
             .store(in: &subscriptions)
+    }
 
-        if let siteCredentials, let request = try? webView.authenticateForWPOrg(with: siteCredentials) {
-            webView.load(request)
-        } else {
-            loadContent()
+    /// Authentication logic differs depending on the destination URL and the current site.
+    /// More information: pe5sF9-3Si-p2
+    ///
+    func startLoading() {
+        guard let url = viewModel.initialURL else {
+            return
+        }
+
+        switch authenticationFlow {
+        case .wpcom:
+            authenticateWPComAndLoadContent(url: url)
+        case .jetpackSSO:
+            authenticateSSOAndLoadContent(url: url)
+        case .siteCredentials:
+            authenticateUsingSiteCredentialsAndLoadContent(url: url)
+        case .none:
+            loadContent(url: url)
         }
     }
 }
 
 // MARK: - Helper methods
 private extension AuthenticatedWebViewController {
-    func loadContent() {
-        guard let url = viewModel.initialURL else {
+    func authenticateWPComAndLoadContent(url: URL) {
+        guard let wpcomCredentials, case .wpcom = wpcomCredentials else {
+            return loadContent(url: url)
+        }
+        do {
+            try webView.authenticateForWPComAndRedirect(to: url, credentials: wpcomCredentials)
+        } catch {
+            loadContent(url: url)
+        }
+    }
+
+    func authenticateSSOAndLoadContent(url: URL) {
+        let tempURL = WooConstants.URLs.wpcomTempRedirectURL.asURL()
+        authenticateWPComAndLoadContent(url: tempURL)
+    }
+
+    func authenticateUsingSiteCredentialsAndLoadContent(url: URL) {
+        guard let siteCredentials, let request = try? webView.authenticateForWPOrg(with: siteCredentials) else {
+            return loadContent(url: url)
+        }
+        webView.load(request)
+    }
+
+    func loadContent(url: URL) {
+        let request = URLRequest(url: url)
+        webView.load(request)
+    }
+
+    func handleRedirect(for url: URL) {
+        guard let initialURL = viewModel.initialURL else {
             return
         }
 
-        /// Authenticate for WP.com automatically if credentials are available.
-        ///
-        if let wpcomCredentials, case .wpcom = wpcomCredentials {
-            webView.authenticateForWPComAndRedirect(to: url, credentials: wpcomCredentials)
-        } else {
-            let request = URLRequest(url: url)
-            webView.load(request)
+        switch url.absoluteString {
+        case WooConstants.URLs.wpcomTempRedirectURL.rawValue:
+            guard let currentSite, let host = URL(string: currentSite.url)?.host else {
+                return loadContent(url: initialURL)
+            }
+            let cookie = HTTPCookie(properties: [
+                .domain: host,
+                .path: "/",
+                .name: Constants.ssoRedirectCookieName,
+                .value: initialURL.absoluteString,
+            ])
+
+            let queryItem = URLQueryItem(name: Constants.actionParam, value: Constants.jetpackSSOAction)
+            guard let cookie, let loginURL = URL(string: currentSite.loginURL)?.appending(queryItems: [queryItem]) else {
+                return loadContent(url: initialURL)
+            }
+            webView.configuration.websiteDataStore.httpCookieStore.setCookie(cookie)
+            loadContent(url: loginURL)
+
+        default:
+            if url.absoluteString.contains(WKWebView.wporgNoncePath) == true,
+               initialURL.absoluteString.contains(WKWebView.wporgNoncePath) != true {
+                // Site credentials login completes, now proceed to load the initial URL.
+                loadContent(url: initialURL)
+            } else {
+                viewModel.handleRedirect(for: url)
+            }
         }
     }
+
 }
 
 extension AuthenticatedWebViewController: WKNavigationDelegate {
@@ -204,7 +281,19 @@ extension AuthenticatedWebViewController: WKNavigationDelegate {
     }
 
     func webView(_ webView: WKWebView, decidePolicyFor navigationResponse: WKNavigationResponse) async -> WKNavigationResponsePolicy {
+        defer {
+            isFirstNavigation = false
+        }
         let response = navigationResponse.response
+        if let initialURL = viewModel.initialURL,
+           viewModel.isAuthenticationFailure(response: response,
+                                             currentSite: currentSite,
+                                             authenticationFlow: authenticationFlow,
+                                             isFirstNavigation: isFirstNavigation) {
+            /// When automatic authentication fails, cancel the navigation and redirect to the original URL instead.
+            loadContent(url: initialURL)
+            return .cancel
+        }
         return await viewModel.decidePolicy(for: response)
     }
 
@@ -244,3 +333,11 @@ extension AuthenticatedWebViewController: WKUIDelegate {
         return nil
     }
 }
+
+private extension AuthenticatedWebViewController {
+    enum Constants {
+        static let actionParam = "action"
+        static let jetpackSSOAction = "jetpack-sso"
+        static let ssoRedirectCookieName = "jetpack_sso_redirect_to"
+    }
+}
diff --git a/WooCommerce/Classes/Authentication/AuthenticatedWebViewModel.swift b/WooCommerce/Classes/Authentication/AuthenticatedWebViewModel.swift
@@ -1,5 +1,6 @@
 import Foundation
 import WebKit
+import struct Yosemite.Site
 
 /// Optional conformance for a `AuthenticatedWebViewModel` implementation to reload a webview asynchronously.
 protocol WebviewReloadable {
@@ -70,4 +71,78 @@ extension AuthenticatedWebViewModel {
         }
         return false
     }
+
+    /// Checks for the appropriate authentication flow based on the current site and the URL to be loaded.
+    ///
+    func authenticationFlow(currentSite: Site,
+                            wpcomCredentialsAvailable: Bool,
+                            wporgCredentialsAvailable: Bool) -> WebViewAuthenticationFlow {
+        guard let initialURL else {
+            return .none
+        }
+
+        if wpcomCredentialsAvailable {
+            if let domain = initialURL.host, Constants.wpcomAcceptedDomains.contains(domain) {
+                return .wpcom
+            } else if currentSite.hasSSOEnabled,
+                      initialURL.absoluteString.hasPrefix(currentSite.url) {
+                return .jetpackSSO
+            }
+        }
+
+        if wporgCredentialsAvailable {
+            return .siteCredentials
+        } else {
+            return .none
+        }
+    }
+
+    /// Returns `true` if the response indicates an authentication failure and `false` otherwise.
+    ///
+    func isAuthenticationFailure(response: URLResponse,
+                                 currentSite: Site?,
+                                 authenticationFlow: WebViewAuthenticationFlow,
+                                 isFirstNavigation: Bool) -> Bool {
+        guard authenticationFlow != .none,
+              isFirstNavigation,
+              let currentSite,
+              let urlResponse = response as? HTTPURLResponse,
+              let url = urlResponse.url else {
+            return false
+        }
+
+        // if the authentication request fails
+        if Constants.errorResponseCodes ~= urlResponse.statusCode {
+            return true
+        }
+
+        let isAuthenticationFailure: Bool = {
+            switch authenticationFlow {
+            case .none:
+                return false
+            case .siteCredentials:
+                return url.absoluteString == currentSite.loginURL
+            case .jetpackSSO, .wpcom:
+                return url.absoluteString == WooConstants.URLs.loginWPCom.rawValue
+            }
+        }()
+
+        guard isAuthenticationFailure else {
+            return false
+        }
+
+        return true
+    }
+}
+
+enum WebViewAuthenticationFlow {
+    case wpcom
+    case jetpackSSO
+    case siteCredentials
+    case none
+}
+
+private enum Constants {
+    static let errorResponseCodes = 400...599
+    static let wpcomAcceptedDomains = ["wordpress.com", "wp.com", "jetpack.com", "woocommerce.com", "jetpack.wordpress.com"]
 }
diff --git a/WooCommerce/Classes/Extensions/WKWebView+Authenticated.swift b/WooCommerce/Classes/Extensions/WKWebView+Authenticated.swift
@@ -27,13 +27,9 @@ extension WKWebView {
         return try URLEncoding.default.encode(request, with: parameters)
     }
 
-    func authenticateForWPComAndRedirect(to url: URL, credentials: Credentials?) {
+    func authenticateForWPComAndRedirect(to url: URL, credentials: Credentials?) throws {
         customUserAgent = UserAgent.defaultUserAgent
-        do {
-            try load(authenticatedPostData(with: credentials, redirectTo: url))
-        } catch {
-            DDLogError("⛔️ Cannot load the authenticated web view on WPCom")
-        }
+        try load(authenticatedPostData(with: credentials, redirectTo: url))
     }
 
     private func authenticatedPostData(with credentials: Credentials?, redirectTo url: URL) throws -> URLRequest {
diff --git a/WooCommerce/Classes/System/WooConstants.swift b/WooCommerce/Classes/System/WooConstants.swift
@@ -250,6 +250,10 @@ extension WooConstants {
         ///
         case pointOfSaleDocumentation = "https://woocommerce.com/document/woo-mobile-app-point-of-sale-mode/"
 
+        /// Temporary redirect URL for authenticated web view when authenticating WPCom automatically
+        ///
+        case wpcomTempRedirectURL = "https://wordpress.com/mobile-redirect"
+
 #if DEBUG
         case orderCreationFeedback = "https://automattic.survey.fm/woo-app-order-creation-testing"
 #else
diff --git a/WooCommerce/Classes/ViewRelated/Products/Edit Product/ProductFormViewController.swift b/WooCommerce/Classes/ViewRelated/Products/Edit Product/ProductFormViewController.swift
@@ -1065,13 +1065,6 @@ private extension ProductFormViewController {
             return
         }
 
-        let stores = ServiceLocator.stores
-        guard let site = stores.sessionManager.defaultSite,
-            stores.shouldAuthenticateAdminPage(for: site) else {
-            WebviewHelper.launch(url.absoluteString, with: self)
-            return
-        }
-
         let viewModel = DefaultAuthenticatedWebViewModel(title: product.name, initialURL: url)
         let controller = AuthenticatedWebViewController(viewModel: viewModel)
         let navigationController = UINavigationController(rootViewController: controller)
diff --git a/WooCommerce/WooCommerce.xcodeproj/project.pbxproj b/WooCommerce/WooCommerce.xcodeproj/project.pbxproj
diff --git a/WooCommerce/WooCommerceTests/Authentication/AuthenticatedWebViewModelTests.swift b/WooCommerce/WooCommerceTests/Authentication/AuthenticatedWebViewModelTests.swift

Original file line number	Diff line number	Diff line change
`@@ -27,13 +27,9 @@ extension WKWebView {`
`27`	`27`	`return try URLEncoding.default.encode(request, with: parameters)`
`28`	`28`	`}`
`29`	`29`
`30`		`- func authenticateForWPComAndRedirect(to url: URL, credentials: Credentials?) {`
	`30`	`+ func authenticateForWPComAndRedirect(to url: URL, credentials: Credentials?) throws {`
`31`	`31`	`customUserAgent = UserAgent.defaultUserAgent`
`32`		`- do {`
`33`		`- try load(authenticatedPostData(with: credentials, redirectTo: url))`
`34`		`- } catch {`
`35`		`- DDLogError("⛔️ Cannot load the authenticated web view on WPCom")`
`36`		`- }`
	`32`	`+ try load(authenticatedPostData(with: credentials, redirectTo: url))`
`37`	`33`	`}`
`38`	`34`
`39`	`35`	`private func authenticatedPostData(with credentials: Credentials?, redirectTo url: URL) throws -> URLRequest {`