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() } }
}

@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)
		}
}