Never see a 401 error again.

RxSwift and Handling Invalid Tokens

I was recently asked to help with creating a system that would reauthorize a user when a 401 was encountered during a network request. After reauthorizing the user, it should retry the request with the new token. Because the word “retry” was used, I immediately thought of the `retry` operator, but in retrospect it might have been better to use `catchError`… That said, here is what I ended up with and my thought process while developing it.

First I roughed out what I need:

1. Since URLRequests will potentially need to be created with new tokens, I need a function that can create a request when given a token.
2. Since retryWhen provides an Observable<Error> that emits every error the stream throws, I need a function that will refresh the token and emit a trigger event on unauthorized errors while passing all other errors down the line.
3. Since multiple requests could be unauthorized while the service is waiting for a new token, I need a way to notify all the requests once the token is provided.

Lastly, since this is going to be a pretty complex job, I need to make sure I can test it without involving the network.

Because of requirement (3) I know that there will have to be some sort of object that can bind all the requests together and this object is primarally in charge of acquiring new tokens when needed. I’ll name it’s class TokenAcquisitionService. This object will also provide the most recent token on request.

With all of the above information. I can expect a function definition like this:

/**
 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.
 - parameter response: A function that sends requests to the network 
     and emits responses. Can be for example 
     `URLSession.shared.rx.response`
 - parameter tokenAcquisitionService: The object responsible for 
     tracking the auth token. All requests should use the same 
     object.
 - parameter request: A function that can build the request when 
     given a token.
 - returns: response of a guaranteed authorized network request.
 **/
typealias Response = (URLRequst) -> Observable<(response: HTTPURLResponse, data: Data)>
typealias Request = (String) -> URLRequest
func getData(response: @escaping Response, tokenAcquisitionService: TokenAcquisitionService, request: @escaping Request) -> Observable<(response: HTTPURLResponse, data: Data)> {
    return Observable
        .deferred { tokenAcquisitionService.token.take(1) }
        .map { request($0) }
        .flatMap { response($0) }
        .map { response in
            guard response.response.statusCode != 401 else { throw ResponseError.unauthorized }
            return response
        }
        .retryWhen { $0.renewToken(with: tokenAcquisitionService) }
}

As you can see above, code isn’t pretty when it’s in a Medium post. The RetryingTokenNetworkService code can be found here.

An explanation of the above function is as follows:

• .deferred { tokenAcquisitionService.token.take(1) } will subscribe to, and pass on, the token emitted by the service whenever this observable is subscribed to. In effect, when the request is started and whenever it is retried, this line will get the latest token from the service.
• .map { request($0) } will build the request with the token provided.
• .flatMap { response($0) } will send the network request and wait for the response.
• guard response.response.statusCode != 401 else { throw ResponseError.unauthorized } will cause the observable to emit an error if the response was rejected for being unauthorized.
• .retryWhen { $0.renewToken(with: tokenAcquisitionService) } will give the service a chance to handle the .unauthorized error, but pass on any other error.

You can use the above function with any network layer you care to name, whether that’s Alamofire, URLSession or whatever, as long as the layer allows the creation of the (URLRequest) -> Observable<(response: HTTPURLResponse, data: Data)> function. Because I used a function definition here instead of tying my code to a particular service, it can be used with any of them and allows for easier testing.

The TokenAcquisitionService itself consists of two functions and a property (along with some private properties.) The code in the gist above has doc comments so I will just explain how they work here rather than what they do. Be sure to read the doc comments in the gist before you continue with the below.

func trackErrors(for:)

In this method, I start with the source object as an observable.

• guard (error as? ResponseError) == .unauthorized else { throw error} will throw any error other than the one it is watching for. The idea is that if something else goes wrong, the calling code would want to know. An .unauthorized error will cause a next event to emit.
• do(onNext:) The block handed to this operator will send a next event to the private Subject that handles the chain which gets a new token. This is run through a subject because all the network requests that need to retry must be merged together so they can all be notified once the new token is acquired. The recursive lock guards the subject against events from different threads.
• filter { _ in false } eats all the next events emitted by the chain. This is because we don’t want to trigger a retry until after the new token has been acquired, but we still want to pass on any stop events (completed or an error other than an unauthorized network request.)
• Observable.merge(token.skip(1).map { _ in }, error) Remember that subscriptions to token will immediately receive a next event but this is unnecessary because we already know that the current token is invalid. We want the trigger for the next token. We also want all the stop events from the error observable (remember, this observable won’t emit next events.

init(initialToken:getToken:extractToken:)

In here is the code that actually refreshes the token. Keep in mind that this is the code that all unauthorized network requests will be subscribed to, as well as any streams that want to know what the current token is and when it changes.

• relay.flatMapFirst { getToken() } is the line that requests a new token and makes sure that only one request is made for each batch of unauthorized requests that are waiting.
• guard urlResponse.response.statusCode / 100 == 2 else { throw ResponseError.refusedToken(response: urlResponse.response, data: urlResponse.data) } will inform all listening requests that the service was unable to successfully extract a token, otherwise it will extract the token and emit it.
• startWith(initialToken) ensures that the token property is pre-loaded with a value.
• 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.

FAQ:

What if I have to display a login screen in order to get a new auth token?

You could do that inside the getToken function that is passed into the service. It could even make a get token request with the server and if that fails, present a screen to the user and make another get token request. You could put a `retry(3)` in there to give the user a few chances. On an ultimate failure you could either throw an error, or push out a response with a non-200 series status code (even a 401. If a 401 occurs during the getToken attempt the service will emit an error.)