Token Acquisition Service for Combine

TokenAcquisitionService.swift

//
//  TokenAcquisitionService.swift
//  CombineSandbox
//
//  Created by Daniel Tartaglia on 11/27/19.
//  Copyright © 2019 Daniel Tartaglia. MIT License.
//

import Foundation
import Combine

public func getData<T>(response: @escaping (URLRequest) -> URLSession.DataTaskPublisher, tokenAcquisitionService: TokenAcquisitionService<T>, request: @escaping (T) -> URLRequest) -> AnyPublisher<URLSession.DataTaskPublisher.Output, Error> {
    return Deferred { tokenAcquisitionService.token.first() }
        .map { request($0) }
        .setFailureType(to: Error.self)
        .flatMap { response($0).mapError { $0 as Error } }
        .tryMap { output in
            guard let response = output.response as? HTTPURLResponse, response.statusCode == 200 else {
                throw TokenAcquisitionError.unauthorized
            }
            return output
        }
        .retryWhen { $0.renewToken(with: tokenAcquisitionService) } // found here: https://gist.github.com/danielt1263/17ebe60a1c7d9aa87c8b5393639a079e
        .eraseToAnyPublisher()
}

public class TokenAcquisitionService<T> {

    public var token: AnyPublisher<T, Never> {
        _token.catch { _ in Empty(completeImmediately: true) }.eraseToAnyPublisher()
    }

    public init(initialToken: T, getToken: @escaping (T) -> URLSession.DataTaskPublisher, extractToken: @escaping (Data) throws -> T) {
        _token = CurrentValueSubject(initialToken)

        cancellable = relay
            .setFailureType(to: Error.self)
            .flatMap(maxPublishers: .max(1)) { getToken($0).mapError { $0 as Error } }
            .tryMap { (output) -> T in
                guard (output.response as! HTTPURLResponse).statusCode / 100 == 2 else { throw TokenAcquisitionError.refusedToken(response: output.response, data: output.data) }
                return try extractToken(output.data)
            }
            .prepend(initialToken)
            .subscribe(_token)
    }

    public func setToken(_ token: T) {
        lock.lock()
        _token.send(token)
        lock.unlock()
    }

    func trackErrors<P>(for publisher: P) -> AnyPublisher<Void, Error> where P: Publisher, P.Output: Error {
        let error = publisher
            .tryMap { error in
                guard (error as? TokenAcquisitionError) == .unauthorized else { throw error }
            }
            .flatMap { [unowned self] in self.token.setFailureType(to: Error.self) }
            .handleEvents(receiveOutput: { [unowned self] oldToken in
                self.lock.lock()
                self.relay.send(oldToken)
                self.lock.unlock()
            })
            .map { _ in }

        let updated = token.dropFirst().map { _ in }.setFailureType(to: Error.self)

        return updated.merge(with: error).eraseToAnyPublisher()
    }

    private let _token: CurrentValueSubject<T, Error>
    private let relay = PassthroughSubject<T, Never>()
    private let lock = NSRecursiveLock()
    private let cancellable: Cancellable?
}

public enum TokenAcquisitionError: Error, Equatable {
    case unauthorized
    case refusedToken(response: URLResponse, data: Data)
}

extension Publisher where Output: Error {
    public func renewToken<T>(with service: TokenAcquisitionService<T>) -> AnyPublisher<Void, Error> {
        return service.trackErrors(for: self)
    }
}