From 610bfe28a04b1de7688f16aedf22a5b9ced1f2f7 Mon Sep 17 00:00:00 2001
From: itaihanski <itaihanski@gmail.com>
Date: Thu, 1 Feb 2024 17:43:09 +0200
Subject: [PATCH] Add support for authenticated flows

---
 src/internal/http/DescopeClient.swift |  7 +++++++
 src/internal/routes/Flow.swift        | 10 +++++++---
 src/types/Flows.swift                 | 22 ++++++++++++++++++++++
 3 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/src/internal/http/DescopeClient.swift b/src/internal/http/DescopeClient.swift
index 4ec3aa3..2098335 100644
--- a/src/internal/http/DescopeClient.swift
+++ b/src/internal/http/DescopeClient.swift
@@ -378,6 +378,13 @@ class DescopeClient: HTTPClient {
         ])
     }
     
+    func flowPrime(codeChallenge: String, flowId: String, refreshJwt: String) async throws {
+        try await post("flow/prime", headers: authorization(with: refreshJwt), body: [
+            "codeChallenge": codeChallenge,
+            "flowId": flowId,
+        ])
+    }
+    
     // MARK: - Others
     
     func me(refreshJwt: String) async throws -> UserResponse {
diff --git a/src/internal/routes/Flow.swift b/src/internal/routes/Flow.swift
index bde04c7..53c9e87 100644
--- a/src/internal/routes/Flow.swift
+++ b/src/internal/routes/Flow.swift
@@ -18,7 +18,11 @@ class Flow: Route, DescopeFlow {
     func start(runner: DescopeFlowRunner) async throws -> AuthenticationResponse {
         // adds some required query parameters to the flow URL to facilitate PKCE and
         // redirection at the end of the flow
-        let (initialURL, codeVerifier) = try prepareInitialRequest(for: runner)
+        let (initialURL, codeVerifier, codeChallenge) = try prepareInitialRequest(for: runner)
+        // prime the flow if the runner has flow authentication info
+        if let flowAuthentication = runner.flowAuthentication {
+            try await client.flowPrime(codeChallenge: codeChallenge, flowId: flowAuthentication.flowId, refreshJwt: flowAuthentication.refreshJwt)
+        }
         log(.info, "Starting flow authentication", initialURL)
         
         // sets the flow we're about to present as the current flow
@@ -188,7 +192,7 @@ private extension Data {
     }
 }
 
-private func prepareInitialRequest(for runner: DescopeFlowRunner) throws -> (url: URL, codeVerifier: String) {
+private func prepareInitialRequest(for runner: DescopeFlowRunner) throws -> (url: URL, codeVerifier: String, codeChallenge: String) {
     guard let randomBytes = Data(randomBytesCount: 32) else { throw DescopeError.flowFailed.with(message: "Error generating random bytes") }
     let hashedBytes = Data(SHA256.hash(data: randomBytes))
     
@@ -208,7 +212,7 @@ private func prepareInitialRequest(for runner: DescopeFlowRunner) throws -> (url
 
     guard let initialURL = components.url else { throw DescopeError.flowFailed.with(message: "Failed to create flow URL") }
     
-    return (initialURL, codeVerifier)
+    return (initialURL, codeVerifier, codeChallenge)
 }
 
 private func prepareRedirectRequest(for runner: DescopeFlowRunner, redirectURL: URL) -> URL? {
diff --git a/src/types/Flows.swift b/src/types/Flows.swift
index ceca1c3..950543c 100644
--- a/src/types/Flows.swift
+++ b/src/types/Flows.swift
@@ -30,9 +30,31 @@ import AuthenticationServices
 ///     email. See the documentation for ``resume(with:)`` for more details.
 @MainActor
 public class DescopeFlowRunner {
+    /// Provide authentication info if the flow is being run by an already
+    /// authenticated user.
+    public struct Authentication {
+        /// The flow ID about to be run
+        public var flowId: String
+        /// The refresh JWT from and active descope session
+        public var refreshJwt: String
+        
+        /// Creates a new ``DescopeFlowRunner.Authentication`` object that encapsulates the
+        /// information required to run a flow for an authenticated user.
+        ///
+        /// - Parameter flowId: The flow ID about to be run.
+        /// - Parameter refreshJwt: The refresh JWT from and active descope session
+        public init(flowId: String, refreshJwt: String) {
+            self.flowId = flowId
+            self.refreshJwt = refreshJwt
+        }
+    }
+    
     /// The URL where the flow is hosted.
     public let flowURL: String
     
+    /// Optional authentication info to allow running flows for authenticated users
+    public var flowAuthentication: Authentication?
+    
     /// Whether the authentication view is allowed to access shared user data.
     ///
     /// Setting this to `true` allows the sandboxed browser in the authentication view