The TokenAcquisitionService automatically retry requests if it receives an unauthorized error. Complete with proof that it works correctly.

TokenAcquisitionService.swift

//
//  TokenAcquisitionService.swift
//
//  Created by Daniel Tartaglia on 16 Jan 2019.
//  Copyright © 2024 Daniel Tartaglia. MIT License.
//

import Foundation
import RxSwift

public typealias Response = (URLRequest) -> Observable<(response: HTTPURLResponse, data: Data)>

/// Builds and makes network requests using the token provided by the service. Will request a new token and retry if
/// the result is an unauthorized (401) error.
///
/// - Parameters:
///   - response: A function that sends requests to the network and emits responses. Can be for example
///               `URLSession.shared.rx.response`
///   - tokenAcquisitionService: The object responsible for tracking the auth token. All requests should use the same
///                              object.
///   - request: A function that can build the request when given a token.
/// - Returns: response of a guaranteed authorized network request.
public func getData<T>(response: @escaping Response,
                       tokenAcquisitionService: TokenAcquisitionService<T>,
                       request: @escaping (T) throws -> URLRequest) -> Observable<(response: HTTPURLResponse, data: Data)> {
    return Observable
        .deferred {
            tokenAcquisitionService.token.take(1).compactMap { result in
                guard case let .success(token) = result else { return nil }
                return token
            }
        }
        .map { try request($0) }
        .flatMap { response($0) }
        .map { response in
            guard response.response.statusCode != 401 else { throw TokenAcquisitionError.unauthorized }
            return response
        }
        .retry(when: { $0.renewToken(with: tokenAcquisitionService) })
}

// MARK: -
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)
    }
}

/// Errors recognized by the `TokenAcquisitionService`.
///
/// - unauthorized: It listens for and activates when it receives an `.unauthorized` error.
/// - refusedToken: It emits a `.refusedToken` error if the `getToken` request fails.
public enum TokenAcquisitionError: Error, Equatable {
    case unauthorized
    case refusedToken(response: HTTPURLResponse, data: Data)
}

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. If token acquisition fails, this will emit a
    /// `.next(.failure)` event.
    public var token: Observable<Result<T, Error>> {
        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/nil 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 { token in
                getToken(token)
                    .map { (urlResponse) -> Result<T, Error> in
                        guard urlResponse.response.statusCode / 100 == 2 else {
                            return .failure(TokenAcquisitionError.refusedToken(response: urlResponse.response, data: urlResponse.data))
                        }
                        return Result(catching: { try extractToken(urlResponse.data) })
                    }
                    .catch { Observable.just(Result.failure($0)) }
            }
            .startWith(.success(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.
    public func setToken(_ token: T) {
        lock.lock()
        _token.onNext(.success(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.take(1) }
            .do(onNext: {
                guard case let .success(token) = $0 else { return }
                lock.lock()
                relay.onNext(token)
                lock.unlock()
            })
            .filter { _ in false }
            .map { _ in }

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

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

TokenAcquisitionServiceTests.swift

//
//  TokenAcquisitionServiceTests.swift
//
//  Created by Daniel Tartaglia on 16 Jan 2019.
//  Copyright © 2024 Daniel Tartaglia. MIT License.
//

import RxSwift
import RxTest
import XCTest

class TokenAcquisitionServiceTests: XCTestCase {

    var scheduler: TestScheduler!
    var tokenResult: TestableObserver<Result<String, Error>>!
    var triggerResult: TestableObserver<Void>!
    var bag: DisposeBag!

    override func setUp() {
        super.setUp()
        scheduler = TestScheduler(initialClock: 0)
        tokenResult = scheduler.createObserver(Result<String, Error>.self)
        triggerResult = scheduler.createObserver(Void.self)
        bag = DisposeBag()
    }

    func testInitial() {
        // given
        func getToken(_ old: String) -> Observable<(response: HTTPURLResponse, data: Data)> {
            XCTFail()
            return .empty()
        }
        func extractToken(_ data: Data) -> String {
            XCTFail()
            return ""
        }
        let service = TokenAcquisitionService(initialToken: "first", getToken: getToken, extractToken: extractToken)

        // when
        service.token
            .bind(to: tokenResult)
            .disposed(by: bag)

        scheduler.start()

        // then
        XCTAssertEqual(tokenResult.events.map(extractEvent), [.next(0, "first")])
    }

    func testUpdate() {
        // given
        func getToken(_ old: String) -> Observable<(response: HTTPURLResponse, data: Data)> {
            XCTFail()
            return .empty()
        }

        let service = TokenAcquisitionService<String>(initialToken: "first", getToken: getToken, extractToken: extractToken)

        // when
        scheduler.scheduleAt(10) {
            service.setToken("second")
        }

        service.token
            .bind(to: tokenResult)
            .disposed(by: bag)

        scheduler.start()

        // then
        XCTAssertEqual(tokenResult.events.map { $0.map { $0.map { try $0.get() } } }, [.next(0, "first"), .next(10, "second")])

    }

    func testUnauthorized() {
        // given
        let trigger = scheduler.createColdObservable([.next(10, TokenAcquisitionError.unauthorized as Error)])
        func getToken(_ old: String) -> Observable<(response: HTTPURLResponse, data: Data)> {
            XCTAssertEqual(old, "first")
            let response = HTTPURLResponse(url: URL(fileURLWithPath: ""), statusCode: 200, httpVersion: nil, headerFields: nil)!
            let data = "second".data(using: .utf8)!
            return Observable.just((response: response, data: data))
        }

        let service = TokenAcquisitionService<String>(initialToken: "first", getToken: getToken, extractToken: extractToken)

        // when
        bag.insert(
            service.token.bind(to: tokenResult),
            trigger.renewToken(with: service).bind(to: triggerResult)
        )

        scheduler.start()

        // then
        XCTAssertEqual(tokenResult.events.map(extractEvent), [.next(0, "first"), .next(10, "second")])
        XCTAssertEqual(triggerResult.events.map { $0.time }, [10])
    }

    func testBadTokenRequest() {
        // given
        let trigger = scheduler.createColdObservable([.next(10, TokenAcquisitionError.unauthorized as Error)])
        let response = HTTPURLResponse(url: URL(fileURLWithPath: ""), statusCode: 500, httpVersion: nil, headerFields: nil)!
        let data = "second".data(using: .utf8)!
        func getToken(_ old: String) -> Observable<(response: HTTPURLResponse, data: Data)> {
            XCTAssertEqual(old, "first")
            return Observable.just((response: response, data: data))
        }

        let service = TokenAcquisitionService<String>(initialToken: "first", getToken: getToken, extractToken: extractToken)

        // when
        bag.insert(
            service.token.bind(to: tokenResult),
            trigger.renewToken(with: service).bind(to: triggerResult)
        )

        scheduler.start()

        // then
        XCTAssertEqual(
            tokenResult.events.map { $0.map { $0.map { $0.mapError { $0 as! TokenAcquisitionError } } } },
            [.next(0, .success("first")), .next(10, .failure(TokenAcquisitionError.refusedToken(response: response, data: data)))]
        )
        XCTAssertEqual(triggerResult.events.map { $0.time }, [10])
    }

    func testOtherErrorsFallThrough() {
        // given
        let trigger = scheduler.createColdObservable([.next(10, RxError.unknown as Error)])
        func getToken(_ old: String) -> Observable<(response: HTTPURLResponse, data: Data)> {
            XCTAssertEqual(old, "first")
            let response = HTTPURLResponse(url: URL(fileURLWithPath: ""), statusCode: 200, httpVersion: nil, headerFields: nil)!
            let data = "second".data(using: .utf8)!
            return Observable.just((response: response, data: data))
        }

        let service = TokenAcquisitionService<String>(initialToken: "first", getToken: getToken, extractToken: extractToken)

        // when
        bag.insert(
            service.token.bind(to: tokenResult),
            trigger.renewToken(with: service).bind(to: triggerResult)
        )

        scheduler.start()

        // then
        XCTAssertEqual(tokenResult.events.map(extractEvent), [.next(0, "first")])
        XCTAssertEqual(triggerResult.events.map { $0.time }, [10])
    }

    func testMultipleUnauthsOnlyCauseOneTokenRequest() {
        // given
        let trigger1 = scheduler.createColdObservable([.next(10, TokenAcquisitionError.unauthorized as Error)])
        let trigger2 = scheduler.createColdObservable([.next(30, TokenAcquisitionError.unauthorized as Error)])
        let triggerResult2 = scheduler.createObserver(Void.self)
        var requestCount = 0
        func getToken(_ old: String) -> Observable<(response: HTTPURLResponse, data: Data)> {
            XCTAssertEqual(old, "first")
            requestCount += 1
            let response = HTTPURLResponse(url: URL(fileURLWithPath: ""), statusCode: 200, httpVersion: nil, headerFields: nil)!
            let data = "second".data(using: .utf8)!
            return Observable.just((response: response, data: data)).delay(.seconds(20), scheduler: scheduler)
        }

        let service = TokenAcquisitionService<String>(initialToken: "first", getToken: getToken, extractToken: extractToken)

        // when
        bag.insert(
            service.token.bind(to: tokenResult),
            trigger1.renewToken(with: service).bind(to: triggerResult),
            trigger2.renewToken(with: service).bind(to: triggerResult2)
        )

        scheduler.start()

        // then
        XCTAssertEqual(tokenResult.events.map(extractEvent), [.next(0, "first"), .next(30, "second")])
        XCTAssertEqual(triggerResult.events.map { $0.time }, [30])
        XCTAssertEqual(triggerResult2.events.map { $0.time }, [30])
        XCTAssertEqual(requestCount, 1)
    }
}

func extractToken(_ data: Data) -> String {
    return String(data: data, encoding: .utf8) ?? ""
}

extension Recorded {
    func map<T>(_ transform: (Value) throws -> T) rethrows -> Recorded<T> {
        Recorded<T>(time: time, value: try transform(value))
    }
}

func extractEvent(_ event: Recorded<Event<Result<String, any Error>>>) -> Recorded<Event<String>> {
    event.map { $0.map { try $0.get() } }
}

DataFetch it's my protocol for Api Calls Classes

I can’t quite understand how to create a token entity for TokenAcquisitionService (

I have no idea why your network API is so complex. URLSession.rx.response is all you should need and any function with the signature (URLRequest) -> Observable<(response: HTTPURLResponse, data: Data)> can serve as a proxy/mock.

As for getting a new token, any function with the signature (T) -> Observable<(response: HTTPURLResponse, data: Data)> will do. This function receives the current token T and just needs to make a network request to get a new token and return an Observable witch will emit the results. The extractToken function gets the results in order to find and return the actual token.

You can change the token service however you want to accommodate your api service, but if you don't know how to make a network request to get a token, then I suggest you talk to whoever is writing your server or read its documentation.

I use Alamofire for the network, I never used URLRequest.rx.response, I have a huge social network, there are a lot of parameters for the request, so it’s complicated) you have an example of how you use the token refresh with your any requests, I would really wanted to reconsider my work with api, I understand that this is not a good option, but I don’t know another, and I don’t know where to get it

Alamofire isn't something I use, but I'm sure it has some way to let you know when there is a 401 error. Find out what way that is and convert it to a TokenAcquisitionError.unauthorized error. I showed that in the custom code sample I gave you.

Thanks for your time, I will work on it

Hi @danielt1263 great article and gist, thank you! Just found this in your post in medium:

share(replay: 1) ensures that any observers that subscribe to the token property will immediately get the most recently acquired token value so they can attempt the first request.

But reviewing your code, you aren't using share(replay: 1) anywhere, so it makes me wonder how actually the same token response is shared across multiple observers when an unauthorized event is fired for multiple requests?

Good catch @guseducampos.

I replaced the share(replay: 1) with a ReplaySubject in order to add the mutable setToken at someone's request and didn't update the article.

I tried to Combine-ify this, but it's missing a few key operators. Namely, retryWhen is a no go here. If anyone has already done the brutal work and is willing to share—please do!

I tried to Combine-ify this, but it's missing a few key operators. Namely, retryWhen is a no go here. If anyone has already done the brutal work and is willing to share—please do!

I'm working on one... Care to help? https://gist.github.com/danielt1263/17ebe60a1c7d9aa87c8b5393639a079e

@danielt1263 hi , i use moya , my network request method:

init(endpointClosure: @escaping MoyaProvider<YunguApi>.EndpointClosure = MoyaProvider<YunguApi>.defaultEndpointMapping,
      requestClosure: @escaping MoyaProvider<YunguApi>.RequestClosure = MoyaProvider<YunguApi>.defaultRequestMapping,
      stubClosure: @escaping MoyaProvider<YunguApi>.StubClosure = MoyaProvider<YunguApi>.neverStub,
      session: Session = MoyaProvider<YunguApi>.defaultAlamofireSession(),
      plugins: [PluginType] = [],
      trackInflights: Bool = false,
      online: Observable<Bool> = connectedToInternet()) {
   
   self.online = online
   self.provider = MoyaProvider(endpointClosure: endpointClosure, requestClosure: requestClosure, stubClosure: stubClosure, session: session, plugins: plugins, trackInflights: trackInflights) 
 }
 
 func request(_ token: YunguApi, isSecondTryAfterAuth: Bool = false) -> Observable<Moya.Response> {
   let actualRequest = provider.rx.request(token)
   return online
     .ignore(value: false)  // Wait until we're online
     .take(1)        // Take 1 to make sure we only invoke the API once.
     .flatMap { _ in // Turn the online state into a network request
       return actualRequest
         
   }
 }

I tried using it but it doesn't work
How can I transform it?
thank you

@marksands

I tried to Combine-ify this, but it's missing a few key operators. Namely, retryWhen is a no go here. If anyone has already done the brutal work and is willing to share—please do!

I was finally able to create a retryWhen operator: https://gist.github.com/danielt1263/17ebe60a1c7d9aa87c8b5393639a079e

Am I right in saying the line:
.flatMapFirst { getToken($0) }
Will ignore any subsequent incoming .unauthorized requests while the first one is getting the new token?

To be a bit clearer:
Two unauthorised api calls are made in quick succession, the flapMapFirst gets a new token and as it ignores the second api call, only the first one will get reattempted with the new token.

Is this correct? I would hope the second api call would get reattempted with the new token too.

from the flatMapFirst docs:

If element is received while there is some projected observable sequence being merged it will simply be ignored.

Thank you for both this gist & the article by the way, I'm learning a lot from trying to implement it myself.

@korovyev Yes, the subsequent incoming .unauthorized events will be ignored while the system is waiting for the request to return a new token; however, the subsequent observable chains will sit waiting for an event to be emitted by the renewToken observable and the observable will emit to all of them once the new token arrives. Therefore, the second api call in your example does get reattempted with the new token.

Yep, my implementation works as you describe. thanks again!

is there any Combine version for this gist ?
got stuck at private let _token = ReplaySubject<T>.create(bufferSize: 1) the closest equivalent in Combine is CurrentValueSubject which requires an initial value to be set in its initializer which can't be accomplished with T (because we don't know what is the type T going to be).

@mesheilah That's easy to take care of, but Combine doesn't have a .retryWhen operator which is required for this solution. I've been working on creating the operator for Combine but I'm not satisfied with what I have so far.

@mesheilah That's easy to take care of, but Combine doesn't have a .retryWhen operator which is required for this solution. I've been working on creating the operator for Combine but I'm not satisfied with what I have so far.

I've checked your custom retryWhen publisher but didn't take it for a spin yet

Thanks for your time, I will work on it

@Terens777 have you found a solution, I'm using alamofire too and I'm having difficulties implement it with daniel code

Thanks for your time, I will work on it

@Terens777 have you found a solution, I'm using alamofire too and I'm having difficulties implement it with daniel code

I have implemented it with Alamofire without problems, you have to create the loadData method that takes pageIndex and returns an Observable of array of result objects, this is the key point, regarding refreshing the token using Alamofire it's easy and it has nothing to do with the code above, you have to create a subclass of RequestInterceptor and override both functions (adapt and retry), there are lots of examples on how to override them.

@danielt1263 hi , i use moya , my network request method:

init(endpointClosure: @escaping MoyaProvider<YunguApi>.EndpointClosure = MoyaProvider<YunguApi>.defaultEndpointMapping,
      requestClosure: @escaping MoyaProvider<YunguApi>.RequestClosure = MoyaProvider<YunguApi>.defaultRequestMapping,
      stubClosure: @escaping MoyaProvider<YunguApi>.StubClosure = MoyaProvider<YunguApi>.neverStub,
      session: Session = MoyaProvider<YunguApi>.defaultAlamofireSession(),
      plugins: [PluginType] = [],
      trackInflights: Bool = false,
      online: Observable<Bool> = connectedToInternet()) {
   
   self.online = online
   self.provider = MoyaProvider(endpointClosure: endpointClosure, requestClosure: requestClosure, stubClosure: stubClosure, session: session, plugins: plugins, trackInflights: trackInflights) 
 }
 
 func request(_ token: YunguApi, isSecondTryAfterAuth: Bool = false) -> Observable<Moya.Response> {
   let actualRequest = provider.rx.request(token)
   return online
     .ignore(value: false)  // Wait until we're online
     .take(1)        // Take 1 to make sure we only invoke the API once.
     .flatMap { _ in // Turn the online state into a network request
       return actualRequest
         
   }
 }

I tried using it but it doesn't work How can I transform it? thank you

@danielt1263 Hello, do you have any suggestion how we can use it in RxMoya? :)

@supervtb I don't know Moya, but you will need something that looks kind of like:

func myRequest(_ token: YunguApi) -> Observable<(response: HTTPURLResponse, data: Data)> {
	YourClass<YunguApi>().request(token)
		.map { moyaResponse in
			return (moyaResponse.httpURLResponse, moyaResponse.data)
		}
}