[PM-24195] fix: Pass userId to refreshToken and setTokens

BitwardenShared/Core/Auth/Services/TestHelpers/MockKeychainRepository.swift

@@ -35,6 +35,10 @@ class MockKeychainRepository: KeychainRepository {
     var setServerCommunicationConfigCalledConfig: BitwardenSdk.ServerCommunicationConfig?
     var setServerCommunicationConfigCalledHostname: String? // swiftlint:disable:this identifier_name
 
+    // Track which userId was passed to get/set methods for testing
+    var getAccessTokenUserId: String?
+    var getRefreshTokenUserId: String?
+
     func deleteAllItems() async throws {
         deleteAllItemsCalled = true
         mockStorage.removeAll()
@@ -77,7 +81,8 @@ class MockKeychainRepository: KeychainRepository {
     }
 
     func getAccessToken(userId: String) async throws -> String {
-        try getAccessTokenResult.get()
+        getAccessTokenUserId = userId
+        return try getAccessTokenResult.get()
     }
 
     func getAuthenticatorVaultKey(userId: String) async throws -> String {
@@ -89,7 +94,8 @@ class MockKeychainRepository: KeychainRepository {
     }
 
     func getRefreshToken(userId: String) async throws -> String {
-        try getRefreshTokenResult.get()
+        getRefreshTokenUserId = userId
+        return try getRefreshTokenResult.get()
     }
 
     func getPendingAdminLoginRequest(userId: String) async throws -> String? {

BitwardenShared/Core/Platform/Services/API/APIService.swift

@@ -44,6 +44,7 @@ class APIService {
     ///   - client: The underlying `HTTPClient` that performs the network request. Defaults
     ///     to `URLSession.shared`.
     ///   - environmentService: The service used by the application to retrieve the environment settings.
+    ///   - errorReporter: The service used by the application to report non-fatal errors.
     ///   - flightRecorder: The service used by the application for recording temporary debug logs.
     ///   - serverCommunicationConfigClientSingleton: The service to get the server communication client
     ///   used to break circular dependency.
@@ -55,6 +56,7 @@ class APIService {
         accountTokenProvider: AccountTokenProvider? = nil,
         client: HTTPClient = URLSession.shared,
         environmentService: EnvironmentService,
+        errorReporter: ErrorReporter,
         flightRecorder: FlightRecorder,
         serverCommunicationConfigClientSingleton: @escaping () -> ServerCommunicationConfigClientSingleton?,
         stateService: StateService,
@@ -85,6 +87,7 @@ class APIService {
         self.accountTokenProvider = accountTokenProvider ?? DefaultAccountTokenProvider(
             httpService: httpServiceBuilder.makeService(baseURLGetter: { environmentService.identityURL }),
             tokenService: tokenService,
+            errorReporter: errorReporter,
         )
 
         apiService = httpServiceBuilder.makeService(

BitwardenShared/Core/Platform/Services/API/AccountTokenProvider.swift

@@ -22,6 +22,9 @@ actor DefaultAccountTokenProvider: AccountTokenProvider {
     /// The delegate to use for specific operations on the token provider.
     private weak var accountTokenProviderDelegate: AccountTokenProviderDelegate?
 
+    /// The service used to report non-fatal errors.
+    private let errorReporter: ErrorReporter
+
     /// The `HTTPService` used to make the API call to refresh the access token.
     private let httpService: HTTPService
 
@@ -42,15 +45,18 @@ actor DefaultAccountTokenProvider: AccountTokenProvider {
     ///   - httpService: The service used to make the API call to refresh the access token.
     ///   - timeProvider: The service used to get the present time.
     ///   - tokenService: The service used to get the current tokens from.
+    ///   - errorReporter: The service used to report non-fatal errors.
     ///
     init(
         httpService: HTTPService,
         timeProvider: TimeProvider = CurrentTime(),
         tokenService: TokenService,
+        errorReporter: ErrorReporter,
     ) {
         self.httpService = httpService
         self.timeProvider = timeProvider
         self.tokenService = tokenService
+        self.errorReporter = errorReporter
     }
 
     // MARK: Methods
@@ -81,16 +87,29 @@ actor DefaultAccountTokenProvider: AccountTokenProvider {
             defer { self.refreshTask = nil }
 
             do {
-                let refreshToken = try await tokenService.getRefreshToken()
+                let expectedUserId = try await tokenService.getActiveAccountId()
+
+                let refreshToken = try await tokenService.getRefreshToken(userId: expectedUserId)
                 let response = try await httpService.send(
                     IdentityTokenRefreshRequest(refreshToken: refreshToken),
                 )
                 let expirationDate = timeProvider.presentTime.addingTimeInterval(TimeInterval(response.expiresIn))
 
+                let userIdAfter = try await tokenService.getActiveAccountId()
+                guard expectedUserId == userIdAfter else {
+                    let error = AccountTokenProviderError(
+                        userIdBefore: expectedUserId,
+                        userIdAfter: userIdAfter,
+                    )
+                    errorReporter.log(error: error)
+                    throw error
+                }
+
                 try await tokenService.setTokens(
                     accessToken: response.accessToken,
                     refreshToken: response.refreshToken,
                     expirationDate: expirationDate,
+                    userId: expectedUserId,
                 )
 
                 return response.accessToken

BitwardenShared/Core/Platform/Services/API/AccountTokenProviderError.swift

@@ -0,0 +1,24 @@
+import Foundation
+
+// MARK: - AccountTokenProviderError
+
+/// Error logged when the active account changes during a token refresh operation.
+///
+struct AccountTokenProviderError: Error, CustomStringConvertible {
+    // MARK: Properties
+
+    /// The active user ID before the token refresh operation.
+    let userIdBefore: String
+
+    /// The active user ID after the token refresh operation.
+    let userIdAfter: String
+
+    // MARK: CustomStringConvertible
+
+    var description: String {
+        """
+        Token refresh race condition detected: Active account changed from '\(userIdBefore)' to '\(userIdAfter)' \
+        during token refresh operation. Tokens were not stored.
+        """
+    }
+}

BitwardenShared/Core/Platform/Services/API/AccountTokenProviderTests.swift

@@ -11,6 +11,7 @@ class AccountTokenProviderTests: BitwardenTestCase {
     // MARK: Properties
 
     var client: MockHTTPClient!
+    var errorReporter: MockErrorReporter!
     var subject: DefaultAccountTokenProvider!
     var timeProvider: MockTimeProvider!
     var tokenService: MockTokenService!
@@ -25,20 +26,23 @@ class AccountTokenProviderTests: BitwardenTestCase {
         super.setUp()
 
         client = MockHTTPClient()
+        errorReporter = MockErrorReporter()
         timeProvider = MockTimeProvider(.mockTime(Date(year: 2025, month: 10, day: 2)))
         tokenService = MockTokenService()
 
         subject = DefaultAccountTokenProvider(
             httpService: HTTPService(baseURL: URL(string: "https://example.com")!, client: client),
             timeProvider: timeProvider,
             tokenService: tokenService,
+            errorReporter: errorReporter,
         )
     }
 
     override func tearDown() async throws {
         try await super.tearDown()
 
         client = nil
+        errorReporter = nil
         subject = nil
         timeProvider = nil
         tokenService = nil
@@ -171,4 +175,94 @@ class AccountTokenProviderTests: BitwardenTestCase {
             _ = try await subject.refreshToken()
         }
     }
+
+    /// `refreshToken()` throws and does not store tokens when the active account switches during the HTTP request,
+    /// preventing tokens from being stored under the wrong account.
+    func test_refreshToken_accountSwitchDuringRequest_throwsAndDoesNotStoreTokens() async throws {
+        // Setup: Account 1 is active
+        tokenService.activeAccountId = "1"
+        tokenService.refreshTokenByUserId["1"] = "REFRESH_1"
+
+        client.result = .httpSuccess(testData: .identityTokenRefresh)
+
+        // Simulate account switch during HTTP request
+        client.onRequest = { _ in
+            // Active account switches to Account 2 while HTTP request is in flight
+            self.tokenService.activeAccountId = "2"
+        }
+
+        await assertAsyncThrows {
+            _ = try await subject.refreshToken()
+        }
+
+        // Verify: error logged, setTokens never called, no tokens stored under either account
+        XCTAssertEqual(errorReporter.errors.count, 1)
+        let error = errorReporter.errors[0] as? AccountTokenProviderError
+        XCTAssertNotNil(error)
+        XCTAssertEqual(error?.userIdBefore, "1")
+        XCTAssertEqual(error?.userIdAfter, "2")
+        XCTAssertNil(tokenService.setTokensCalledWithUserId)
+    }
+
+    /// `refreshToken()` logs an error and throws when the active account changes during the token refresh operation,
+    /// preventing tokens from being stored in the wrong account.
+    func test_refreshToken_throwsRaceCondition_whenUserIdChanges() async throws {
+        tokenService.activeAccountId = "user-1"
+        tokenService.accessToken = "🔑"
+        tokenService.refreshToken = "🔒"
+
+        client.result = .httpSuccess(testData: .identityTokenRefresh)
+
+        // Simulate account switch during HTTP request
+        client.onRequest = { _ in
+            self.tokenService.activeAccountId = "user-2"
+        }
+
+        await assertAsyncThrows {
+            _ = try await subject.refreshToken()
+        }
+
+        // Verify error was logged and setTokens was never called
+        XCTAssertEqual(errorReporter.errors.count, 1)
+        let error = errorReporter.errors[0] as? AccountTokenProviderError
+        XCTAssertNotNil(error)
+        XCTAssertEqual(error?.userIdBefore, "user-1")
+        XCTAssertEqual(error?.userIdAfter, "user-2")
+        XCTAssertNil(tokenService.setTokensCalledWithUserId)
+    }
+
+    /// `refreshToken()` does not log an error when the active account remains the same.
+    func test_refreshToken_doesNotLogError_whenUserIdStaysSame() async throws {
+        tokenService.activeAccountId = "user-1"
+        tokenService.accessToken = "🔑"
+        tokenService.refreshToken = "🔒"
+
+        client.result = .httpSuccess(testData: .identityTokenRefresh)
+
+        _ = try await subject.refreshToken()
+
+        // Verify no error was logged
+        XCTAssertEqual(errorReporter.errors.count, 0)
+    }
+
+    /// `refreshToken()` throws when `getActiveAccountId` throws before setting tokens,
+    /// preventing tokens from being stored without verifying the active account.
+    func test_refreshToken_throws_whenGetUserIdAfterThrows() async throws {
+        tokenService.accessToken = "🔑"
+        tokenService.refreshToken = "🔒"
+        client.result = .httpSuccess(testData: .identityTokenRefresh)
+
+        // Clear the active account ID during the HTTP request so the
+        // getActiveAccountId() call (before setTokens) throws noActiveAccount.
+        client.onRequest = { _ in
+            self.tokenService.activeAccountId = ""
+        }
+
+        await assertAsyncThrows(error: StateServiceError.noActiveAccount) {
+            _ = try await subject.refreshToken()
+        }
+
+        // setTokens was never called
+        XCTAssertNil(tokenService.setTokensCalledWithUserId)
+    }
 }

BitwardenShared/Core/Platform/Services/API/RefreshableAPIServiceTests.swift

@@ -21,6 +21,7 @@ class RefreshableAPIServiceTests: BitwardenTestCase {
         subject = APIService(
             accountTokenProvider: accountTokenProvider,
             environmentService: MockEnvironmentService(),
+            errorReporter: MockErrorReporter(),
             flightRecorder: MockFlightRecorder(),
             serverCommunicationConfigClientSingleton: { MockServerCommunicationConfigClientSingleton() },
             stateService: MockStateService(),

BitwardenShared/Core/Platform/Services/API/TestHelpers/APIService+Mocks.swift

@@ -9,6 +9,7 @@ extension APIService {
         accountTokenProvider: AccountTokenProvider? = nil,
         client: HTTPClient,
         environmentService: EnvironmentService = MockEnvironmentService(),
+        errorReporter: ErrorReporter = MockErrorReporter(),
         flightRecorder: FlightRecorder = MockFlightRecorder(),
         // swiftlint:disable:next line_length
         serverCommunicationConfigClientSingleton: ServerCommunicationConfigClientSingleton = MockServerCommunicationConfigClientSingleton(),
@@ -18,6 +19,7 @@ extension APIService {
             accountTokenProvider: accountTokenProvider,
             client: client,
             environmentService: environmentService,
+            errorReporter: errorReporter,
             flightRecorder: flightRecorder,
             serverCommunicationConfigClientSingleton: { serverCommunicationConfigClientSingleton },
             stateService: stateService,

BitwardenShared/Core/Platform/Services/ServiceContainer.swift

@@ -532,6 +532,7 @@ public class ServiceContainer: Services { // swiftlint:disable:this type_body_le
         let apiService = APIService(
             client: noRedirectSession,
             environmentService: environmentService,
+            errorReporter: errorReporter,
             flightRecorder: flightRecorder,
             serverCommunicationConfigClientSingleton: { serverCommConfigClientSingletonHolder },
             stateService: stateService,

BitwardenShared/Core/Platform/Services/TestHelpers/MockTokenService.swift

@@ -10,15 +10,35 @@ class MockTokenService: TokenService {
     var getIsExternalResult: Result<Bool, Error> = .success(false)
     var refreshToken: String? = "REFRESH_TOKEN"
 
+    // Track which userId was used in explicit userId methods
+    var getAccessTokenCalledWithUserId: String?
+    var getRefreshTokenCalledWithUserId: String?
+    var setTokensCalledWithUserId: String?
+    var activeAccountId: String = "1"
+    var accessTokenByUserId: [String: String] = [:]
+    var refreshTokenByUserId: [String: String] = [:]
+
     func getAccessToken() async throws -> String {
         guard let accessToken else { throw StateServiceError.noActiveAccount }
         return accessToken
     }
 
+    func getAccessToken(userId: String) async throws -> String {
+        getAccessTokenCalledWithUserId = userId
+        return accessTokenByUserId[userId] ?? accessToken ?? "ACCESS_TOKEN"
+    }
+
     func getAccessTokenExpirationDate() async throws -> Date? {
         try accessTokenExpirationDateResult.get()
     }
 
+    func getActiveAccountId() async throws -> String {
+        if activeAccountId.isEmpty {
+            throw StateServiceError.noActiveAccount
+        }
+        return activeAccountId
+    }
+
     func getIsExternal() async throws -> Bool {
         try getIsExternalResult.get()
     }
@@ -28,9 +48,24 @@ class MockTokenService: TokenService {
         return refreshToken
     }
 
+    func getRefreshToken(userId: String) async throws -> String {
+        getRefreshTokenCalledWithUserId = userId
+        return refreshTokenByUserId[userId] ?? refreshToken ?? "REFRESH_TOKEN"
+    }
+
     func setTokens(accessToken: String, refreshToken: String, expirationDate: Date) async {
         self.accessToken = accessToken
         self.refreshToken = refreshToken
         self.expirationDate = expirationDate
     }
+
+    func setTokens(accessToken: String, refreshToken: String, expirationDate: Date, userId: String) async {
+        setTokensCalledWithUserId = userId
+        accessTokenByUserId[userId] = accessToken
+        refreshTokenByUserId[userId] = refreshToken
+        self.expirationDate = expirationDate
+        // Also update legacy properties for backward compatibility with existing tests
+        self.accessToken = accessToken
+        self.refreshToken = refreshToken
+    }
 }