fix(auth): 0154 — serialize TokenManager.refresh to prevent double-redeem race

Co-authored-by: OpenAI Codex <codex@openai.com>

Author: roninjin10

Shared/Sources/SmithersAuth/TokenManager.swift

@@ -87,51 +87,23 @@ public final class TokenManager {
     /// tokens hit the Keychain before this function resolves.
     @discardableResult
     public func refresh() async throws -> OAuth2Tokens {
-        // De-duplicate concurrent refreshes — two parallel 401s must not
-        // both redeem the same refresh token.
-        let (existingTask, current): (Task<OAuth2Tokens, Error>?, OAuth2Tokens?) = {
-            lock.lock(); defer { lock.unlock() }
-            return (inFlightRefresh, cached)
-        }()
-        if let existing = existingTask {
-            return try await existing.value
-        }
-        guard let current = current else {
-            throw TokenManagerError.notSignedIn
-        }
-        let task = Task { [client, store] () throws -> OAuth2Tokens in
-            let newTokens: OAuth2Tokens
-            do {
-                newTokens = try await client.refresh(refreshToken: current.refreshToken)
-            } catch let err as OAuth2Error {
-                if case .whitelistDenied(let msg) = err {
-                    throw TokenManagerError.whitelistDenied(msg)
-                }
-                throw TokenManagerError.refreshFailed(err)
+        let task = withLock { () -> Task<OAuth2Tokens, Error>? in
+            if let existing = inFlightRefresh {
+                return existing
             }
-            // WRITE-BEFORE-RETURN. If this throws, the caller treats the
-            // user as signed out — see `refreshAndRetry` wiring.
-            do {
-                try store.save(newTokens)
-            } catch let e as TokenStoreError {
-                throw TokenManagerError.persistenceFailed(e)
+            guard let current = cached else {
+                return nil
             }
-            return newTokens
+            let task = makeRefreshTask(current: current)
+            inFlightRefresh = task
+            return task
         }
-        setInflight(task)
 
-        do {
-            let newTokens = try await task.value
-            self.setCached(newTokens)
-            self.clearInflight()
-            return newTokens
-        } catch {
-            self.clearInflight()
-            // A refresh failure (expired/rotated/revoked refresh token OR
-            // a store write failure) locks the user out. Wipe silently.
-            await self.localSignOut()
-            throw error
+        guard let task else {
+            throw TokenManagerError.notSignedIn
         }
+
+        return try await task.value
     }
 
     private func setCached(_ t: OAuth2Tokens) {
@@ -142,8 +114,44 @@ public final class TokenManager {
         lock.lock(); inFlightRefresh = nil; lock.unlock()
     }
 
-    private func setInflight(_ t: Task<OAuth2Tokens, Error>) {
-        lock.lock(); inFlightRefresh = t; lock.unlock()
+    private func withLock<T>(_ body: () throws -> T) rethrows -> T {
+        lock.lock()
+        defer { lock.unlock() }
+        return try body()
+    }
+
+    private func makeRefreshTask(current: OAuth2Tokens) -> Task<OAuth2Tokens, Error> {
+        Task { [self, client, store] in
+            defer { clearInflight() }
+
+            do {
+                let newTokens: OAuth2Tokens
+                do {
+                    newTokens = try await client.refresh(refreshToken: current.refreshToken)
+                } catch let err as OAuth2Error {
+                    if case .whitelistDenied(let msg) = err {
+                        throw TokenManagerError.whitelistDenied(msg)
+                    }
+                    throw TokenManagerError.refreshFailed(err)
+                }
+
+                // WRITE-BEFORE-RETURN. If this throws, the caller treats the
+                // user as signed out — see `refreshAndRetry` wiring.
+                do {
+                    try store.save(newTokens)
+                } catch let e as TokenStoreError {
+                    throw TokenManagerError.persistenceFailed(e)
+                }
+
+                setCached(newTokens)
+                return newTokens
+            } catch {
+                // A refresh failure (expired/rotated/revoked refresh token OR
+                // a store write failure) locks the user out. Wipe silently.
+                await localSignOut()
+                throw error
+            }
+        }
     }
 
     /// 401-retry helper. `perform` receives a bearer access token; return

Shared/Tests/SmithersAuthTests/MockHTTPTransport.swift

@@ -34,19 +34,33 @@ final class MockHTTPTransport: HTTPTransport {
         }
     }
 
+    private let lock = NSLock()
     var responses: [CannedResponse] = []
+    var sendDelayNanoseconds: UInt64 = 0
     private(set) var recorded: [Recorded] = []
 
     func send(_ request: URLRequest) async throws -> (Data, Int, [String: String]) {
-        recorded.append(Recorded(
-            url: request.url!,
-            method: request.httpMethod ?? "GET",
-            body: request.httpBody
-        ))
-        guard !responses.isEmpty else {
-            return (Data(), 500, [:])
+        if sendDelayNanoseconds > 0 {
+            try? await Task.sleep(nanoseconds: sendDelayNanoseconds)
         }
-        let r = responses.removeFirst()
-        return (r.body, r.status, r.headers)
+
+        return withLock {
+            recorded.append(Recorded(
+                url: request.url!,
+                method: request.httpMethod ?? "GET",
+                body: request.httpBody
+            ))
+            guard !responses.isEmpty else {
+                return (Data(), 500, [:])
+            }
+            let r = responses.removeFirst()
+            return (r.body, r.status, r.headers)
+        }
+    }
+
+    private func withLock<T>(_ body: () throws -> T) rethrows -> T {
+        lock.lock()
+        defer { lock.unlock() }
+        return try body()
     }
 }

Shared/Tests/SmithersAuthTests/MockedServerIntegrationTests.swift

@@ -193,23 +193,49 @@ final class MockedServerIntegrationTests: XCTestCase {
 
     // MARK: - Concurrent 401s collapse into one refresh
 
-    @MainActor
     func test_concurrent_refreshes_deduplicate() async throws {
         let transport = MockHTTPTransport()
-        transport.responses.append(.json(payload: [
-            "access_token": "A2", "refresh_token": "R2",
-        ]))
+        transport.sendDelayNanoseconds = 25_000_000
+        for _ in 0..<2 {
+            transport.responses.append(.json(payload: [
+                "access_token": "A2", "refresh_token": "R2",
+            ]))
+        }
         let client = OAuth2Client(config: makeConfig(), transport: transport)
         let store = InMemoryTokenStore(initial: OAuth2Tokens(accessToken: "A1", refreshToken: "R1"))
         let mgr = TokenManager(client: client, store: store)
 
-        async let a = mgr.refresh()
-        async let b = mgr.refresh()
-        let (ta, tb) = try await (a, b)
+        let a = Task { try await mgr.refresh() }
+        let b = Task { try await mgr.refresh() }
+
+        let ta = try await a.value
+        let tb = try await b.value
+
+        XCTAssertEqual(ta, tb)
         XCTAssertEqual(ta.accessToken, "A2")
-        XCTAssertEqual(tb.accessToken, "A2")
-        // Only ONE refresh hit the wire.
-        XCTAssertEqual(transport.recorded.count, 1)
+        XCTAssertEqual(refreshRequests(in: transport).count, 1)
+    }
+
+    func test_concurrent_refreshes_stress_deduplicate() async throws {
+        let transport = MockHTTPTransport()
+        transport.sendDelayNanoseconds = 25_000_000
+        for _ in 0..<20 {
+            transport.responses.append(.json(payload: [
+                "access_token": "A2", "refresh_token": "R2",
+            ]))
+        }
+        let client = OAuth2Client(config: makeConfig(), transport: transport)
+        let store = InMemoryTokenStore(initial: OAuth2Tokens(accessToken: "A1", refreshToken: "R1"))
+        let mgr = TokenManager(client: client, store: store)
+
+        let tasks = (0..<20).map { _ in
+            Task { try await mgr.refresh() }
+        }
+        let results = try await tasks.asyncMap { try await $0.value }
+        let first = try XCTUnwrap(results.first)
+
+        XCTAssertTrue(results.allSatisfy { $0 == first })
+        XCTAssertEqual(refreshRequests(in: transport).count, 1)
     }
 
     // MARK: - CSRF state mismatch is rejected
@@ -231,3 +257,22 @@ private final class CountingWipeHandler: SessionWipeHandler {
     var wipeCount = 0
     func wipeAfterSignOut() { wipeCount += 1 }
 }
+
+private func refreshRequests(in transport: MockHTTPTransport) -> [MockHTTPTransport.Recorded] {
+    transport.recorded.filter {
+        $0.url.absoluteString.hasSuffix("/api/oauth2/token")
+            && ($0.bodyString ?? "").contains("grant_type=refresh_token")
+    }
+}
+
+private extension Array {
+    func asyncMap<T>(_ transform: (Element) async throws -> T) async rethrows -> [T] {
+        var results: [T] = []
+        results.reserveCapacity(count)
+        for element in self {
+            let value = try await transform(element)
+            results.append(value)
+        }
+        return results
+    }
+}