How to synchronously refresh an access token using Alamofire + RxSwift

Question

I have this generic fetchData() function in my NetworkManager class that is able to request make a authorised request to the network and if it fail (after a number of retries) emits an error that will restart my app (requesting a new login). I need that this retry token be called synchronously, I mean, if multiple requests failed, only one should be requesting the refresh token at once. And if that one fail, and the other one requests must be discarded. I already tried some approached using DispatchGroup / NSRecursiveLock / and also with calling the function cancelRequests describing bellow (in this case, the tasks count is always 0). How can I make this behaviour works in this scenario?

• My NetworkManager class:

    public func fetchData<Type: Decodable>(fromApi api: TargetType,
                                           decodeFromKeyPath keyPath: String? = nil) -> Single<Response> {
        
        let request = MultiTarget(api)

        return provider.rx.request(request)
                .asRetriableAuthenticated(target: request)
    }

    func cancelAllRequests(){
        if #available(iOS 9.0, *) {
            DefaultAlamofireManager
                .sharedManager
                .session
                .getAllTasks { (tasks) in
                tasks.forEach{ $0.cancel() }
            }
        } else {
            DefaultAlamofireManager
                .sharedManager
                .session
                .getTasksWithCompletionHandler { (sessionDataTask, uploadData, downloadData) in
                    
                sessionDataTask.forEach { $0.cancel() }
                uploadData.forEach { $0.cancel() }
                downloadData.forEach { $0.cancel() }
            }
        }
    }

• The Single extension that make the retry works:

public extension PrimitiveSequence where TraitType == SingleTrait, ElementType == Response {
    
    private var refreshTokenParameters: TokenParameters {
        TokenParameters(clientId: "pdappclient",
                grantType: "refresh_token",
                refreshToken: KeychainManager.shared.refreshToken)
    }

    func retryWithToken(target: MultiTarget) -> Single<E> {
        self.catchError { error -> Single<Response> in
                    if case Moya.MoyaError.statusCode(let response) = error {
                        if self.isTokenExpiredError(error) {
                            return Single.error(error)
                        } else {
                            return self.parseError(response: response)
                        }
                    }
                    return Single.error(error)
                }
                .retryToken(target: target)
                .catchError { error -> Single<Response> in
                    if case Moya.MoyaError.statusCode(let response) = error {
                        return self.parseError(response: response)
                    }
                    return Single.error(InvalidGrantException())
                }
    }

    private func retryToken(target: MultiTarget) -> Single<E> {
        let maxRetries = 1
        return self.retryWhen({ error in
            error
                    .enumerated()
                    .flatMap { (attempt, error) -> Observable<Int> in
                        if attempt >= maxRetries {
                            return Observable.error(error)
                        }
                        if self.isTokenExpiredError(error) {
                            return Observable<Int>.just(attempt + 1)
                        }
                        return Observable.error(error)
                    }
                    .flatMap { _ -> Single<TokenResponse> in
                        self.refreshTokenRequest()
                    }
                    .share()
                    .asObservable()
        })
    }
    
    private func refreshTokenRequest() -> Single<TokenResponse> {
        return NetworkManager.shared.fetchData(fromApi: IdentityServerAPI
            .token(parameters: self.refreshTokenParameters)).do(onSuccess: { tokenResponse in
                    
            KeychainManager.shared.accessToken = tokenResponse.accessToken
            KeychainManager.shared.refreshToken = tokenResponse.refreshToken
        }, onError: { error in
            NetworkManager.shared.cancelAllRequests()
        })
    }

    func parseError<E>(response: Response) -> Single<E> {
        if response.statusCode == 401 {
            // TODO
        }

        let decoder = JSONDecoder()
        if let errors = try? response.map([BaseResponseError].self, atKeyPath: "errors", using: decoder,
                failsOnEmptyData: true) {
            return Single.error(BaseAPIErrorResponse(errors: errors))
        }

        return Single.error(APIError2.unknown)
    }

    func isTokenExpiredError(_ error: Error) -> Bool {
        if let moyaError = error as? MoyaError {
            switch moyaError {
            case .statusCode(let response):
                if response.statusCode != 401 {
                    return false
                } else if response.data.count == 0 {
                    return true
                }
            default:
                break
            }
        }
        return false
    }

    func filterUnauthorized() -> Single<E> {
        flatMap { (response) -> Single<E> in
            if 200...299 ~= response.statusCode {
                return Single.just(response)
            } else if response.statusCode == 404 {
                return Single.just(response)
            } else {
                return Single.error(MoyaError.statusCode(response))
            }
        }
    }

    func asRetriableAuthenticated(target: MultiTarget) -> Single<Element> {
        filterUnauthorized()
                .retryWithToken(target: target)
                .filterStatusCode()
    }

    func filterStatusCode() -> Single<E> {
        flatMap { (response) -> Single<E> in
            if 200...299 ~= response.statusCode {
                return Single.just(response)
            } else {
                return self.parseError(response: response)
            }
        }
    }
}

Answer 1

I found a solution to my problem using DispatchWorkItem and controlling the entrance on my function with a boolean: isTokenRefreshing. Maybe that's not the most elegant solution, but it works.

So, in my NetworkManager class I added this two new properties:

public var savedRequests: [DispatchWorkItem] = []
public var isTokenRefreshing = false

Now in my SingleTrait extension, whenever I enter in the token refresh method I set the boolean isTokenRefreshing to true. So, if it's true, instead of starting another request, I simply throw a RefreshTokenProcessInProgressException and save the current request in my savedRequests array.

private func saveRequest(_ block: @escaping () -> Void) {
    // Save request to DispatchWorkItem array
    NetworkManager.shared.savedRequests.append( DispatchWorkItem {
        block()
    })
}

(Of course, that, if the token refresh succeeds you have to remember to continue all the savedRequests that are saved inside the array, it's not described inside the code down below yet).

Well, my SingleTrait extension is now something like this:

import Foundation
import Moya
import RxSwift
import Domain

public extension PrimitiveSequence where TraitType == SingleTrait, ElementType == Response {
    
    private var refreshTokenParameters: TokenParameters {
        TokenParameters(clientId: "pdappclient",
                grantType: "refresh_token",
                refreshToken: KeychainManager.shared.refreshToken)
    }

    func retryWithToken(target: MultiTarget) -> Single<E> {
        return self.catchError { error -> Single<Response> in
                    if case Moya.MoyaError.statusCode(let response) = error {
                        if self.isTokenExpiredError(error) {
                            return Single.error(error)
                        } else {
                            return self.parseError(response: response)
                        }
                    }
                    return Single.error(error)
                }
                .retryToken(target: target)
                .catchError { error -> Single<Response> in
                    if case Moya.MoyaError.statusCode(let response) = error {
                        return self.parseError(response: response)
                    }
                    return Single.error(error)
                }
    }

    private func retryToken(target: MultiTarget) -> Single<E> {
        let maxRetries = 1
        
        return self.retryWhen({ error in
            error
                    .enumerated()
                    .flatMap { (attempt, error) -> Observable<Int> in
                        if attempt >= maxRetries {
                            return Observable.error(error)
                        }
                        if self.isTokenExpiredError(error) {
                            return Observable<Int>.just(attempt + 1)
                        }
                        return Observable.error(error)
                    }
                    .flatMapFirst { _ -> Single<TokenResponse> in
                        if NetworkManager.shared.isTokenRefreshing {
                            self.saveRequest {
                                self.retryToken(target: target)
                            }
                            return Single.error(RefreshTokenProcessInProgressException())
                        } else {
                            return self.refreshTokenRequest()
                        }
                    }
                    .share()
                    .asObservable()
        })
    }
    
    private func refreshTokenRequest() -> Single<TokenResponse> {
        NetworkManager.shared.isTokenRefreshing = true
        
        return NetworkManager.shared.fetchData(fromApi: IdentityServerAPI
            .token(parameters: self.refreshTokenParameters))
            .do(onSuccess: { tokenResponse in
                KeychainManager.shared.accessToken = tokenResponse.accessToken
                KeychainManager.shared.refreshToken = tokenResponse.refreshToken
            }).catchError { error -> Single<TokenResponse> in
                return Single.error(InvalidGrantException())
        }
    }
    
    private func saveRequest(_ block: @escaping () -> Void) {
        // Save request to DispatchWorkItem array
        NetworkManager.shared.savedRequests.append( DispatchWorkItem {
            block()
        })
    }

    func parseError<E>(response: Response) -> Single<E> {
        if response.statusCode == 401 {
            // TODO
        }

        let decoder = JSONDecoder()
        if let errors = try? response.map([BaseResponseError].self, atKeyPath: "errors", using: decoder,
                failsOnEmptyData: true) {
            return Single.error(BaseAPIErrorResponse(errors: errors))
        }

        return Single.error(APIError2.unknown)
    }

    func isTokenExpiredError(_ error: Error) -> Bool {
        if let moyaError = error as? MoyaError {
            switch moyaError {
            case .statusCode(let response):
                if response.statusCode != 401 {
                    return false
                } else if response.data.count == 0 {
                    return true
                }
            default:
                break
            }
        }
        return false
    }

    func filterUnauthorized() -> Single<E> {
        flatMap { (response) -> Single<E> in
            if 200...299 ~= response.statusCode {
                return Single.just(response)
            } else if response.statusCode == 404 {
                return Single.just(response)
            } else {
                return Single.error(MoyaError.statusCode(response))
            }
        }
    }

    func asRetriableAuthenticated(target: MultiTarget) -> Single<Element> {
        filterUnauthorized()
                .retryWithToken(target: target)
                .filterStatusCode()
    }

    func filterStatusCode() -> Single<E> {
        flatMap { (response) -> Single<E> in
            if 200...299 ~= response.statusCode {
                return Single.just(response)
            } else {
                return self.parseError(response: response)
            }
        }
    }
}

In my case, if the token refresh fails, after a N number of retries, I restart the app. And so, whenever a restart the application I'm setting the isTokenRefreshing to false again.

This is the way I found to solve this problem. If you have another approach, please let me know.

Answer 2

Here is an RxSwift solution: RxSwift and Handling Invalid Tokens

Just posting the link isn't the best, so I will post the core of the solution as well:

The key is to make a class that is much like the ActivityMonitor class but handles token refreshing...

public final class TokenAcquisitionService<T> {

    /// responds with the current token immediatly and emits a new token whenver a new one is aquired. You can, for example, subscribe to it in order to save the token as it's updated.
    public var token: Observable<T> {
        return _token.asObservable()
    }

    public typealias GetToken = (T) -> Observable<(response: HTTPURLResponse, data: Data)>

    /// Creates a `TokenAcquisitionService` object that will store the most recent authorization token acquired and will acquire new ones as needed.
    ///
    /// - Parameters:
    ///   - initialToken: The token the service should start with. Provide a token from storage or an empty string (object represting a missing token) if one has not been aquired yet.
    ///   - getToken: A function responsable for aquiring new tokens when needed.
    ///   - extractToken: A function that can extract a token from the data returned by `getToken`.
    public init(initialToken: T, getToken: @escaping GetToken, extractToken: @escaping (Data) throws -> T) {
        relay
            .flatMapFirst { getToken($0) }
            .map { (urlResponse) -> T in
                guard urlResponse.response.statusCode / 100 == 2 else { throw TokenAcquisitionError.refusedToken(response: urlResponse.response, data: urlResponse.data) }
                return try extractToken(urlResponse.data)
            }
            .startWith(initialToken)
            .subscribe(_token)
            .disposed(by: disposeBag)
    }

    /// Allows the token to be set imperativly if necessary.
    /// - Parameter token: The new token the service should use. It will immediatly be emitted to any subscribers to the service.
    func setToken(_ token: T) {
        lock.lock()
        _token.onNext(token)
        lock.unlock()
    }

    /// Monitors the source for `.unauthorized` error events and passes all other errors on. When an `.unauthorized` error is seen, `self` will get a new token and emit a signal that it's safe to retry the request.
    ///
    /// - Parameter source: An `Observable` (or like type) that emits errors.
    /// - Returns: A trigger that will emit when it's safe to retry the request.
    func trackErrors<O: ObservableConvertibleType>(for source: O) -> Observable<Void> where O.Element == Error {
        let lock = self.lock
        let relay = self.relay
        let error = source
            .asObservable()
            .map { error in
                guard (error as? TokenAcquisitionError) == .unauthorized else { throw error }
            }
            .flatMap { [unowned self] in  self.token }
            .do(onNext: {
                lock.lock()
                relay.onNext($0)
                lock.unlock()
            })
            .filter { _ in false }
            .map { _ in }

        return Observable.merge(token.skip(1).map { _ in }, error)
    }

    private let _token = ReplaySubject<T>.create(bufferSize: 1)
    private let relay = PublishSubject<T>()
    private let lock = NSRecursiveLock()
    private let disposeBag = DisposeBag()
}

extension ObservableConvertibleType where Element == Error {

    /// Monitors self for `.unauthorized` error events and passes all other errors on. When an `.unauthorized` error is seen, the `service` will get a new token and emit a signal that it's safe to retry the request.
    ///
    /// - Parameter service: A `TokenAcquisitionService` object that is being used to store the auth token for the request.
    /// - Returns: A trigger that will emit when it's safe to retry the request.
    public func renewToken<T>(with service: TokenAcquisitionService<T>) -> Observable<Void> {
        return service.trackErrors(for: self)
    }
}

Once you put the above in your app, you can just add a .retryWhen { $0.renewToken(with: tokenAcquisitionService) } to the end of your request. Make sure your request emits a ResponseError.unauthorized if the token is unauthorized and the service will handle the retry.