From 7f4c9204a0528de42fe480742a11d22926751343 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C5=81ukasz=20Rutkowski?= <lukus178@gmail.com>
Date: Mon, 6 Nov 2023 17:45:06 +0100
Subject: [PATCH] Refresh no longer valid token stored in active task (#14)

---
 Sources/OAuthenticator/Authenticator.swift    |  7 +-
 .../AuthenticatorTests.swift                  | 70 +++++++++++++++++++
 2 files changed, 76 insertions(+), 1 deletion(-)

diff --git a/Sources/OAuthenticator/Authenticator.swift b/Sources/OAuthenticator/Authenticator.swift
index 8d870bb..9876d3d 100644
--- a/Sources/OAuthenticator/Authenticator.swift
+++ b/Sources/OAuthenticator/Authenticator.swift
@@ -199,7 +199,12 @@ extension Authenticator {
 
         var login: Login
         do {
-            login = try await loginFromTask(task: task)
+            do {
+                login = try await loginFromTask(task: task)
+            } catch AuthenticatorError.tokenInvalid {
+                let newTask = makeLoginTask(manual: manual, userAuthenticator: userAuthenticator)
+                login = try await loginFromTask(task: newTask)
+            }
 
             // Inform authenticationResult closure of new login information
             self.config.authenticationStatusHandler?(.success(login))
diff --git a/Tests/OAuthenticatorTests/AuthenticatorTests.swift b/Tests/OAuthenticatorTests/AuthenticatorTests.swift
index 5c0fec9..67bac4a 100644
--- a/Tests/OAuthenticatorTests/AuthenticatorTests.swift
+++ b/Tests/OAuthenticatorTests/AuthenticatorTests.swift
@@ -414,4 +414,74 @@ final class AuthenticatorTests: XCTestCase {
 		XCTAssertEqual(mockLoader.requests[0].allHTTPHeaderFields!["Authorization"], "Bearer EXPIRED")
 		XCTAssertEqual(mockLoader.requests[1].allHTTPHeaderFields!["Authorization"], "Bearer REFRESHED")
 	}
+
+    @MainActor
+    func testTokenExpiredAfterUseRefresh() async throws {
+        var sentRequests: [URLRequest] = []
+        let mockLoader: URLResponseProvider = { request in
+            sentRequests.append(request)
+            return ("hello".data(using: .utf8)!, URLResponse())
+        }
+
+        var refreshedLogins: [Login] = []
+        let refreshProvider: TokenHandling.RefreshProvider = { login, _, _ in
+            refreshedLogins.append(login)
+            return Login(token: "REFRESHED")
+        }
+
+        let tokenHandling = TokenHandling(
+            authorizationURLProvider: Self.disabledAuthorizationURLProvider,
+            loginProvider: Self.disabledLoginProvider,
+            refreshProvider: refreshProvider,
+            responseStatusProvider: TokenHandling.allResponsesValid
+        )
+
+        let storedLogin = Login(
+            accessToken: Token(value: "EXPIRE SOON", expiry: Date().addingTimeInterval(1)),
+            refreshToken: Token(value: "REFRESH")
+        )
+        var loadLoginCount = 0
+        var savedLogins: [Login] = []
+        let storage = LoginStorage {
+            loadLoginCount += 1
+            return storedLogin
+        } storeLogin: { login in
+            savedLogins.append(login)
+        }
+
+        let config = Authenticator.Configuration(appCredentials: Self.mockCredentials,
+                                                 loginStorage: storage,
+                                                 tokenHandling: tokenHandling,
+                                                 userAuthenticator: Self.disabledUserAuthenticator)
+
+        let auth = Authenticator(config: config, urlLoader: mockLoader)
+
+        let (_, _) = try await auth.response(for: URLRequest(url: URL(string: "https://example.com")!))
+        XCTAssertEqual(sentRequests.count, 1, "First request should be sent")
+        XCTAssertEqual(sentRequests.first?.value(forHTTPHeaderField: "Authorization"), "Bearer EXPIRE SOON", "Non expired token should be used for first request")
+        XCTAssertTrue(refreshedLogins.isEmpty, "Token should not be refreshed after first request")
+        XCTAssertEqual(loadLoginCount, 1, "Login should be loaded from storage once")
+        XCTAssertTrue(savedLogins.isEmpty, "Login storage should not be updated after first request")
+
+        // Let the token expire
+        sleep(1)
+
+        let (_, _) = try await auth.response(for: URLRequest(url: URL(string: "https://example.com")!))
+        XCTAssertEqual(refreshedLogins.count, 1, "Token should be refreshed")
+        XCTAssertEqual(refreshedLogins.first?.accessToken.value, "EXPIRE SOON", "Expired token should be passed to refresh call")
+        XCTAssertEqual(refreshedLogins.first?.refreshToken?.value, "REFRESH", "Refresh token should be passed to refresh call")
+        XCTAssertEqual(loadLoginCount, 2, "New login should be loaded from storage")
+        XCTAssertEqual(sentRequests.count, 2, "Second request should be sent")
+        let secondRequest = sentRequests.dropFirst().first
+        XCTAssertEqual(secondRequest?.value(forHTTPHeaderField: "Authorization"), "Bearer REFRESHED", "Refreshed token should be used for second request")
+        XCTAssertEqual(savedLogins.first?.accessToken.value, "REFRESHED", "Refreshed token should be saved to storage")
+
+        let (_, _) = try await auth.response(for: URLRequest(url: URL(string: "https://example.com")!))
+        XCTAssertEqual(refreshedLogins.count, 1, "No additional refreshes should happen")
+        XCTAssertEqual(loadLoginCount, 2, "No additional login loads should happen")
+        XCTAssertEqual(sentRequests.count, 3, "Third request should be sent")
+        let thirdRequest = sentRequests.dropFirst(2).first
+        XCTAssertEqual(thirdRequest?.value(forHTTPHeaderField: "Authorization"), "Bearer REFRESHED", "Refreshed token should be used for third request")
+        XCTAssertEqual(savedLogins.count, 1, "No additional logins should be saved to storage")
+    }
 }