Support WordPress.com magic link login

WordPress/Classes/Login/WordPressDotComAuthenticator.swift

@@ -4,6 +4,8 @@ import Foundation
 import UIKit
 import WordPressData
 import WordPressShared
+import BuildSettingsKit
+@preconcurrency import Combine
 
 /// Log in or sign up a WordPress.com account via web.
 ///
@@ -47,15 +49,33 @@ struct WordPressDotComAuthenticator {
         case loadingSites(Error)
     }
 
+    private static let callbackNotification = Foundation.Notification.Name(rawValue: "WordPressDotComAuthenticatorCallbackURL")
+
+    static func redirectURI(for scheme: String) -> String {
+        "\(scheme)://oauth2-callback"
+    }
+
+    static func handleAppOpeningURL(_ url: URL, appURLScheme: String = BuildSettings.current.appURLScheme) -> Bool {
+        guard url.scheme == appURLScheme, url.absoluteString.hasPrefix(redirectURI(for: appURLScheme)) else {
+            return false
+        }
+
+        NotificationCenter.default.post(name: callbackNotification, object: url)
+        return true
+    }
+
+    let redirectURIScheme: String
     private let coreDataStack: CoreDataStackSwift
     private let authenticator: ((URL) throws(AuthenticationError) -> URL)?
 
     init(
         coreDataStack: CoreDataStackSwift = ContextManager.shared,
-        authenticator: ((URL) throws(AuthenticationError) -> URL)? = nil
+        authenticator: ((URL) throws(AuthenticationError) -> URL)? = nil,
+        redirectURIScheme: String = BuildSettings.current.appURLScheme
     ) {
         self.coreDataStack = coreDataStack
         self.authenticator = authenticator
+        self.redirectURIScheme = redirectURIScheme
     }
 
     /// Sign in WP.com account.
@@ -172,7 +192,7 @@ struct WordPressDotComAuthenticator {
     ) async throws(AuthenticationError) -> String {
         let clientId = ApiCredentials.client
         let clientSecret = ApiCredentials.secret
-        let redirectURI = "x-wordpress-app://oauth2-callback"
+        let redirectURI = Self.redirectURI(for: redirectURIScheme)
 
         var queries: [String: Any] = [
             "client_id": clientId,
@@ -190,33 +210,85 @@ struct WordPressDotComAuthenticator {
         let authorizeURL = try? URLEncoding.queryString.encode(URLRequest(url: URL(string: "https://public-api.wordpress.com/oauth2/authorize")!), with: queries).url
         guard let authorizeURL else { throw .urlError(URLError(.badURL)) }
 
-        let callbackURL = try await authorize(from: viewController, url: authorizeURL, prefersEphemeralWebBrowserSession: prefersEphemeralWebBrowserSession)
+        let callbackURL = try await authorize(from: viewController, url: authorizeURL, prefersEphemeralWebBrowserSession: prefersEphemeralWebBrowserSession, redirectURI: redirectURI)
 
         return try await handleAuthorizeCallbackURL(callbackURL, clientId: clientId, clientSecret: clientSecret, redirectURI: redirectURI)
     }
 
-    private func authorize(from viewController: UIViewController, url authorizeURL: URL, prefersEphemeralWebBrowserSession: Bool) async throws(AuthenticationError) -> URL {
+    @MainActor
+    private func authorize(from viewController: UIViewController, url authorizeURL: URL, prefersEphemeralWebBrowserSession: Bool, redirectURI: String) async throws(AuthenticationError) -> URL {
         if let authenticator {
             return try authenticator(authorizeURL)
         }
 
-        return try await withCheckedTypedThrowingContinuation { continuation in
-            DispatchQueue.main.async {
-                let provider = WebAuthenticationPresentationAnchorProvider(anchor: viewController.view.window ?? UIWindow())
-                let session = ASWebAuthenticationSession(url: authorizeURL, callbackURLScheme: "x-wordpress-app") { url, error in
-                    let result: Result<URL, AuthenticationError>
-                    if let url {
-                        result = .success(url)
-                    } else {
-                        DDLogWarn("Error from authentication session: \(String(describing: error))")
-                        result = .failure(.cancelled)
-                    }
-                    continuation(result)
+        class CancellableHolder: Cancellable, @unchecked Sendable {
+            var cancellable: AnyCancellable?
+
+            func cancel() {
+                cancellable?.cancel()
+                cancellable = nil
+            }
+        }
+        let cancellable = CancellableHolder()
+
+        // When the login account is unverified, the user receives an email with a login link. Opening the login link takes
+        // the user to Safari, where they can authorize/deny the OAuth request. After authorization/denial, the website
+        // opens the app via the OAuth callback URL. In this scenario, the callback URL is received by the open URL
+        // delegate method instead of the ASWebAuthenticationSession instance.
+        let callbackURLViaOpenAppURL = NotificationCenter.default
+            .publisher(for: Self.callbackNotification)
+            .compactMap {
+                if let url = $0.object as? URL {
+                    return url
+                }
+                return nil
+            }
+            .filter { (url: URL) in
+                url.absoluteString.hasPrefix(redirectURI)
+            }
+            .setFailureType(to: AuthenticationError.self)
+            .first()
+
+        let callbackURLViaWebAuthenticationSession = PassthroughSubject<URL, AuthenticationError>()
+        let provider = WebAuthenticationPresentationAnchorProvider(anchor: viewController.view.window ?? UIWindow())
+        let session = ASWebAuthenticationSession(url: authorizeURL, callbackURLScheme: redirectURIScheme) { url, error in
+            if let url {
+                callbackURLViaWebAuthenticationSession.send(url)
+                callbackURLViaWebAuthenticationSession.send(completion: .finished)
+            } else {
+                DDLogWarn("Error from authentication session: \(String(describing: error))")
+                callbackURLViaWebAuthenticationSession.send(completion: .failure(.cancelled))
+            }
+        }
+        session.prefersEphemeralWebBrowserSession = prefersEphemeralWebBrowserSession
+        session.presentationContextProvider = provider
+        session.start()
+
+        do {
+            return try await withTaskCancellationHandler {
+                try await withCheckedThrowingContinuation { continuation in
+                    cancellable.cancellable = Publishers.Merge(callbackURLViaOpenAppURL, callbackURLViaWebAuthenticationSession)
+                        .first()
+                        .sink(receiveCompletion: { [session] in
+                            session.cancel()
+
+                            if case let .failure(error) = $0 {
+                                continuation.resume(throwing: error)
+                            }
+                        }, receiveValue: { url in
+                            continuation.resume(returning: url)
+                        })
                 }
-                session.prefersEphemeralWebBrowserSession = prefersEphemeralWebBrowserSession
-                session.presentationContextProvider = provider
-                session.start()
+            } onCancel: { [session] in
+                session.cancel()
+                cancellable.cancel()
             }
+        } catch {
+            // We can't get the `AuthenticationError` type from the syntax level, because
+            // the `withTaskCancellationHandler` and `withCheckedThrowingContinuation` functions are not updated to
+            // support typed throw. But the error can only be `AuthenticationError` at runtime.
+            wpAssert(error is AuthenticationError)
+            throw (error as? AuthenticationError) ?? .cancelled
         }
     }
 
@@ -281,19 +353,6 @@ struct WordPressDotComAuthenticator {
     }
 }
 
-/// typed-throw version of `withCheckedThrowingContinuation`
-private func withCheckedTypedThrowingContinuation<T, E: Error>(body: (@escaping ((Result<T, E>) -> Void)) -> Void) async throws(E) -> T {
-    do {
-        return try await withCheckedThrowingContinuation { continuation in
-            body {
-                continuation.resume(with: $0)
-            }
-        }
-    } catch {
-        throw (error as! E)
-    }
-}
-
 private extension WordPressDotComAuthenticator.SignInError {
     var alertMessage: String? {
         switch self {

WordPress/Classes/System/WordPressAppDelegate+openURL.swift

@@ -30,6 +30,10 @@ import BuildSettingsKit
             return JetpackNotificationMigrationService.shared.handleNotificationMigrationOnWordPress()
         }
 
+        if WordPressDotComAuthenticator.handleAppOpeningURL(url) {
+            return true
+        }
+
         guard url.scheme == BuildSettings.current.appURLScheme else {
             return false
         }

WordPress/WordPressTest/Login/WordPressDotComAuthenticatorTests.swift

@@ -16,7 +16,7 @@ class WordPressDotComAuthenticatorTests: CoreDataTestCase {
     func testAuthenticateSuccess() async {
         stubTokenExchange()
 
-        let authenticator = WordPressDotComAuthenticator(authenticator: fakeAuthenticator(callback: ["code": "random"]))
+        let authenticator = WordPressDotComAuthenticator(authenticator: fakeAuthenticator(callback: ["code": "random"]), redirectURIScheme: "testapp")
         do {
             let _ = try await authenticator.authenticate(from: UIViewController(), prefersEphemeralWebBrowserSession: false)
         } catch {
@@ -25,7 +25,7 @@ class WordPressDotComAuthenticatorTests: CoreDataTestCase {
     }
 
     func testAuthenticateWithInvalidCallbackURL() async {
-        let authenticator = WordPressDotComAuthenticator(authenticator: fakeAuthenticator(callback: ["empty": "yes"]))
+        let authenticator = WordPressDotComAuthenticator(authenticator: fakeAuthenticator(callback: ["empty": "yes"]), redirectURIScheme: "testapp")
         do {
             let _ = try await authenticator.authenticate(from: .init(), prefersEphemeralWebBrowserSession: false)
             XCTFail("Unexpected successful result")
@@ -37,7 +37,7 @@ class WordPressDotComAuthenticatorTests: CoreDataTestCase {
     }
 
     func testAuthenticateWithAccessDenied() async {
-        let authenticator = WordPressDotComAuthenticator(authenticator: fakeAuthenticator(callback: ["error": "access_denied"]))
+        let authenticator = WordPressDotComAuthenticator(authenticator: fakeAuthenticator(callback: ["error": "access_denied"]), redirectURIScheme: "testapp")
         do {
             let _ = try await authenticator.authenticate(from: .init(), prefersEphemeralWebBrowserSession: false)
             XCTFail("Unexpected successful result")
@@ -58,7 +58,36 @@ class WordPressDotComAuthenticatorTests: CoreDataTestCase {
         try XCTAssertNil(WPAccount.lookupDefaultWordPressComAccount(in: mainContext))
 
         // When signing in with a WP.com account
-        let authenticator = WordPressDotComAuthenticator(coreDataStack: contextManager, authenticator: fakeAuthenticator(callback: ["code": "random"]))
+        let authenticator = WordPressDotComAuthenticator(coreDataStack: contextManager, authenticator: fakeAuthenticator(callback: ["code": "random"]), redirectURIScheme: "testapp")
+        let accountID = await authenticator.signIn(from: UIViewController(), context: .default)
+        XCTAssertNotNil(accountID)
+
+        // The new WP.com acount should be set as the default account.
+        let isDefaultAccount = try mainContext.existingObject(with: XCTUnwrap(accountID)).isDefaultWordPressComAccount
+        XCTAssertTrue(isDefaultAccount)
+    }
+
+    @MainActor
+    func testSignInViaMagicLink() async throws {
+        stubTokenExchange()
+        stubGetAccountDetails()
+        stubGetSites()
+
+        // Given the app is not signed in with a WP.com account
+        try XCTAssertNil(WPAccount.lookupDefaultWordPressComAccount(in: mainContext))
+
+        // When signing in with a WP.com account via web authentication session using a magic link
+        let authenticator = WordPressDotComAuthenticator(coreDataStack: contextManager, redirectURIScheme: "testapp")
+        // Here we post a notification to simulate this scenario in the production app:
+        // - user taps the magic link in their email
+        // - the link opens Safari and then redirects to the app
+        // - the app posts a notification.
+        Task.detached { @MainActor in
+            try await Task.sleep(for: .milliseconds(100))
+            let handled = WordPressDotComAuthenticator.handleAppOpeningURL(URL(string: "testapp://oauth2-callback?code=random")!, appURLScheme: "testapp")
+            XCTAssertTrue(handled)
+        }
+
         let accountID = await authenticator.signIn(from: UIViewController(), context: .default)
         XCTAssertNotNil(accountID)
 
@@ -79,7 +108,7 @@ class WordPressDotComAuthenticatorTests: CoreDataTestCase {
         AccountService(coreDataStack: contextManager).setDefaultWordPressComAccount(account)
 
         // When signing in with another WP.com account.
-        let authenticator = WordPressDotComAuthenticator(coreDataStack: contextManager, authenticator: fakeAuthenticator(callback: ["code": "random"]))
+        let authenticator = WordPressDotComAuthenticator(coreDataStack: contextManager, authenticator: fakeAuthenticator(callback: ["code": "random"]), redirectURIScheme: "testapp")
         do {
             let _ = try await authenticator.attemptSignIn(from: UIViewController(), context: .jetpackSite(accountEmail: nil))
             XCTFail("Unexpected successful result")
@@ -95,7 +124,7 @@ class WordPressDotComAuthenticatorTests: CoreDataTestCase {
         stubGetAccountDetails()
         stubGetSitesError()
 
-        let authenticator = WordPressDotComAuthenticator(coreDataStack: contextManager, authenticator: fakeAuthenticator(callback: ["code": "random"]))
+        let authenticator = WordPressDotComAuthenticator(coreDataStack: contextManager, authenticator: fakeAuthenticator(callback: ["code": "random"]), redirectURIScheme: "testapp")
         do {
             let _ = try await authenticator.attemptSignIn(from: UIViewController(), context: .default)
             XCTFail("Unexpected successful result")
@@ -117,7 +146,7 @@ class WordPressDotComAuthenticatorTests: CoreDataTestCase {
         stubGetAccountDetails()
 
         // Then `mismatchedEmail` error should be returned
-        let authenticator = WordPressDotComAuthenticator(coreDataStack: contextManager, authenticator: fakeAuthenticator(callback: ["code": "random"]))
+        let authenticator = WordPressDotComAuthenticator(coreDataStack: contextManager, authenticator: fakeAuthenticator(callback: ["code": "random"]), redirectURIScheme: "testapp")
         do {
             let _ = try await authenticator.attemptSignIn(from: UIViewController(), context: .reauthentication(accountEmail: account.email))
             XCTFail("Unexpected successful result")