From bfbcddd5942d5dcf73a21ca5615cfd669e4eb610 Mon Sep 17 00:00:00 2001
From: Guilherme Souza <grsouza@pm.me>
Date: Mon, 1 Apr 2024 14:11:53 -0300
Subject: [PATCH] refactor(auth): in memory code verifier storage (#296)

---
 Sources/Auth/AuthClient.swift                 | 55 +++++++------------
 .../Auth/Internal/CodeVerifierStorage.swift   | 23 ++------
 Tests/AuthTests/Mocks/Mocks.swift             |  5 +-
 Tests/AuthTests/RequestsTests.swift           |  4 +-
 4 files changed, 30 insertions(+), 57 deletions(-)

diff --git a/Sources/Auth/AuthClient.swift b/Sources/Auth/AuthClient.swift
index 2217b088..03f59a87 100644
--- a/Sources/Auth/AuthClient.swift
+++ b/Sources/Auth/AuthClient.swift
@@ -529,34 +529,31 @@ public actor AuthClient {
 
   /// Log in an existing user by exchanging an Auth Code issued during the PKCE flow.
   public func exchangeCodeForSession(authCode: String) async throws -> Session {
-    guard let codeVerifier = try codeVerifierStorage.getCodeVerifier() else {
+    guard let codeVerifier = codeVerifierStorage.get() else {
       throw AuthError.pkce(.codeVerifierNotFound)
     }
-    do {
-      let session: Session = try await api.execute(
-        .init(
-          path: "/token",
-          method: .post,
-          query: [URLQueryItem(name: "grant_type", value: "pkce")],
-          body: configuration.encoder.encode(
-            [
-              "auth_code": authCode,
-              "code_verifier": codeVerifier,
-            ]
-          )
+
+    let session: Session = try await api.execute(
+      .init(
+        path: "/token",
+        method: .post,
+        query: [URLQueryItem(name: "grant_type", value: "pkce")],
+        body: configuration.encoder.encode(
+          [
+            "auth_code": authCode,
+            "code_verifier": codeVerifier,
+          ]
         )
       )
-      .decoded(decoder: configuration.decoder)
+    )
+    .decoded(decoder: configuration.decoder)
 
-      try codeVerifierStorage.deleteCodeVerifier()
+    codeVerifierStorage.set(nil)
 
-      try await sessionManager.update(session)
-      eventEmitter.emit(.signedIn, session: session)
+    try await sessionManager.update(session)
+    eventEmitter.emit(.signedIn, session: session)
 
-      return session
-    } catch {
-      throw error
-    }
+    return session
   }
 
   /// Log in an existing user via a third-party provider.
@@ -1020,19 +1017,7 @@ public actor AuthClient {
     }
 
     let codeVerifier = PKCE.generateCodeVerifier()
-
-    do {
-      try codeVerifierStorage.storeCodeVerifier(codeVerifier)
-    } catch {
-      assertionFailure(
-        """
-        An error occurred while storing the code verifier,
-        PKCE flow may not work as expected.
-
-        Error: \(error.localizedDescription)
-        """
-      )
-    }
+    codeVerifierStorage.set(codeVerifier)
 
     let codeChallenge = PKCE.generateCodeChallenge(from: codeVerifier)
     let codeChallengeMethod = codeVerifier == codeChallenge ? "plain" : "s256"
@@ -1049,7 +1034,7 @@ public actor AuthClient {
 
   private func isPKCEFlow(url: URL) -> Bool {
     let fragments = extractParams(from: url)
-    let currentCodeVerifier = try? codeVerifierStorage.getCodeVerifier()
+    let currentCodeVerifier = codeVerifierStorage.get()
     return fragments.contains(where: { $0.name == "code" }) && currentCodeVerifier != nil
   }
 
diff --git a/Sources/Auth/Internal/CodeVerifierStorage.swift b/Sources/Auth/Internal/CodeVerifierStorage.swift
index 07d4b326..b7feeb7a 100644
--- a/Sources/Auth/Internal/CodeVerifierStorage.swift
+++ b/Sources/Auth/Internal/CodeVerifierStorage.swift
@@ -1,30 +1,19 @@
+import ConcurrencyExtras
 import Foundation
 @_spi(Internal) import _Helpers
 
 struct CodeVerifierStorage: Sendable {
-  var getCodeVerifier: @Sendable () throws -> String?
-  var storeCodeVerifier: @Sendable (_ code: String) throws -> Void
-  var deleteCodeVerifier: @Sendable () throws -> Void
+  var get: @Sendable () -> String?
+  var set: @Sendable (_ code: String?) -> Void
 }
 
 extension CodeVerifierStorage {
   static let live: Self = {
-    @Dependency(\.configuration.localStorage) var localStorage
-
-    let key = "supabase.code-verifier"
+    let code = LockIsolated(String?.none)
 
     return Self(
-      getCodeVerifier: {
-        try localStorage.retrieve(key: key).flatMap {
-          String(data: $0, encoding: .utf8)
-        }
-      },
-      storeCodeVerifier: { code in
-        try localStorage.store(key: key, value: Data(code.utf8))
-      },
-      deleteCodeVerifier: {
-        try localStorage.remove(key: key)
-      }
+      get: { code.value },
+      set: { code.setValue($0) }
     )
   }()
 }
diff --git a/Tests/AuthTests/Mocks/Mocks.swift b/Tests/AuthTests/Mocks/Mocks.swift
index 9bc7da38..4c73d0e6 100644
--- a/Tests/AuthTests/Mocks/Mocks.swift
+++ b/Tests/AuthTests/Mocks/Mocks.swift
@@ -16,9 +16,8 @@ let clientURL = URL(string: "http://localhost:54321/auth/v1")!
 
 extension CodeVerifierStorage {
   static let mock = Self(
-    getCodeVerifier: unimplemented("CodeVerifierStorage.getCodeVerifier"),
-    storeCodeVerifier: unimplemented("CodeVerifierStorage.storeCodeVerifier"),
-    deleteCodeVerifier: unimplemented("CodeVerifierStorage.deleteCodeVerifier")
+    get: unimplemented("CodeVerifierStorage.get"),
+    set: unimplemented("CodeVerifierStorage.set")
   )
 }
 
diff --git a/Tests/AuthTests/RequestsTests.swift b/Tests/AuthTests/RequestsTests.swift
index 13c7a383..46d1f6dc 100644
--- a/Tests/AuthTests/RequestsTests.swift
+++ b/Tests/AuthTests/RequestsTests.swift
@@ -160,7 +160,7 @@ final class RequestsTests: XCTestCase {
 
       Current.sessionManager = .live
       Current.sessionStorage.storeSession = { _ in }
-      Current.codeVerifierStorage.getCodeVerifier = { nil }
+      Current.codeVerifierStorage.get = { nil }
       Current.currentDate = { currentDate }
 
       let url = URL(
@@ -184,7 +184,7 @@ final class RequestsTests: XCTestCase {
   func testSessionFromURLWithMissingComponent() async {
     let sut = makeSUT()
 
-    Current.codeVerifierStorage.getCodeVerifier = { nil }
+    Current.codeVerifierStorage.get = { nil }
 
     let url = URL(
       string: