Refresh Token Mechanism with Combine framework iOS
Asked Answered
S

1

1

Followed refresh token strategy from this gist - https://gist.github.com/saroar/ca78de9dc798cdbaaa47791380062596

but couldn't get it working.

Ideal flow - run function from NetworkAgent will make a call using the tokenSubject from Authenticator. If that call fails with authentication error, refreshToken function from Authenticator will be called. Now, when this refresh token call gets completed, the tokenSubject will send a value (check the receiveCompletion block of sink operator in Authenticator) which NetworkAgent has already subscribed in run function. Once the NetworkAgent receives this value from tokenSubject it will again execute that request in the run function.

Bug - Once the 401 error is received, the refresh token request is made but after that completes successfully, failed request doesn't get retried. Seems like the value sent by the tokenSubject is not received by the subscriber may be due to some reference issue.

Please find my implementation below: -

This is the Network Agent that makes all the api calls -

class NetworkAgent {
struct Response {
let value: T
let response: URLResponse
}
let authenticator = Authenticator()

func run<T: Decodable>(_ request: URLRequest) -> AnyPublisher<Response<T>, Error> {
    let tokenSubject = authenticator.tokenSubject()
    // a better way to handle refresh token logic would be to use tryCatch operator suceeded by retry operator
    // should be done in one of the upcoming releases
    
    return tokenSubject
        .flatMap({ token -> AnyPublisher<Response<T>, Error> in
            return URLSession.shared
                .dataTaskPublisher(for: request)
                .tryMap { [weak self] result -> Response<T> in
                    self?.logResponse(result.response, request: request, data: result.data)
                    
                    // check for the response status
                    if let httpURLResponse = result.response as? HTTPURLResponse {
                        let networkResponse = APIClient.status(httpURLResponse)
                        switch networkResponse {
                        case .success: break
                            
                        case .failure(let error):
                            if error == HTTPURLResponseError.authenticationError {
                                // refresh token logic
                                if request.url?.absoluteString.contains("/auth/token/refresh") == false {
                                    self?.authenticator.refreshToken(using: tokenSubject)
                                }
                            } else {
                                throw error
                            }
                        }
                    }
                    
                    // map response and return in case of a successful request
                    let value = try JSONDecoder().decode(T.self, from: result.data)
                    return Response(value: value, response: result.response)
                }
                .retry(3)
                .receive(on: DispatchQueue.main)
                .eraseToAnyPublisher()
        })
        .handleEvents(receiveOutput: { _ in
            tokenSubject.send(completion: .finished)
        })
        .eraseToAnyPublisher()
}
}

And this is the Authenticator class implementation -

class Authenticator {
//MARK:- Properties
private var disposables = Set()
private let queue = DispatchQueue(label: "Autenticator.(UUID().uuidString)")

//MARK:- Functions
func refreshToken<S: Subject>(using subject: S) where S.Output == Bool {
    queue.sync {
        URLSession.shared
            .dataTaskPublisher(for: APIRouter.refreshToken.request()!)
            .retry(3)
            .sink { com in
                print(com)
                subject.send(true)
            } receiveValue: { data in
                print(#line, data)
                
                do {
                    let jsonDecoder2 = JSONDecoder()
                    let user = try jsonDecoder2.decode(User.self, from: data.data)
                    
                    LocalData.loggedInUser = user
                    LocalData.token = user.token
                    
                } catch  {
                    print(#line, error)
                }
                
            }
            .store(in: &disposables)
    }
}

func tokenSubject() -> CurrentValueSubject<Bool, Never> {
    return CurrentValueSubject(true)
}
}

Any input would be greatly appreciated.

Skein answered 13/4, 2021 at 13:57 Comment(4)
You initial request failed (finished) with authenticationError you need to call it again yourself.Napolitano
@Napolitano thanks for looking at it but I am already calling the refresh token function of Authenticator when the original request fails in tryMap operator.Skein
I am a little bit confused, what is the original request you mentioned? As far as i understood you use the refresh logic to get a new token but you older request (what ever that is) failed because of 401. You need to retry that one with the new token.Napolitano
@Napolitano I've updated the question details with the program flow. Please look at it now, it may help to understand the issue.Skein
F
1

Take a look here, working solution:

       session
           .dataTaskPublisher(for: request)
           .tryMap { result in
               guard let urlResponse = result.response as? HTTPURLResponse, urlResponse.statusCode == 401 else {
                   return result
               }
               throw ResponseError.sessionExpired
           }
           .tryCatch({ [weak self] error -> AnyPublisher<(data: Data, response: URLResponse), URLError> in
               guard let self = self else { throw error }

               if (error as? ResponseError) == .sessionExpired {
                   log("SESSION RESTORE", "Request failed with 401:", request.url?.absoluteString ?? "")
                   log("SESSION RESTORE", "Start")
                   return try self.sessionRestorePublisher(requestBody: requestBody, passError: error)
               }
               throw error
            })
            ....

    private func sessionRestorePublisher(requestBody: Request, passError: Error) throws -> AnyPublisher<(data: Data, response: URLResponse), URLError> {
        guard let refreshSession = self.refreshSession() else { throw passError }

        return refreshSession
            .tryMap { _ in
                throw ResponseError.sessionExpired
            }
            .tryCatch({ [weak self] _ -> AnyPublisher<(data: Data, response: URLResponse), URLError> in
                guard let self = self else {
                    throw passError
                }

                log("SESSION RESTORE", "Prepare updated request")
                guard let newRequest = self.getURLRequest(requestBody: requestBody) else {
                    throw passError
                }

                log("SESSION RESTORE", "Updated request FIRE:", newRequest.url?.absoluteString ?? "")
                return self.session
                    .dataTaskPublisher(for: newRequest)
                    .eraseToAnyPublisher()
             })
            .mapError { $0 as? URLError ?? URLError(URLError.Code.userAuthenticationRequired) }
            .eraseToAnyPublisher()
    }

Keep in mind, session is NSURLSession, func refreshSession() -> AnyPublisher<Bool, Error>? - here is another dataTaskPublisher that refreshes session and saves somewhere it NSUserDefaults new token so here in tryMap we throw sessionExpired, and then we catch it in tryCatch, and replace our publisher with another publisher that refreshes session. and in sessionResporePublisher we specially raise error in tryMap and catch it again in tryCatch to replace it back with our original publisher with reconstructed request with new token. For me works like a charm, session restored on FLY.

Only one thing more - don't throw ResponseError.sessionExpired but check if session was indeed refreshed with Bool, or whatever check you do here, and then do appropriate action - replace with new publisher or follow to login in app )

Best Regards Guys!, have a fun coding!

Freemason answered 7/12, 2022 at 14:49 Comment(2)
This would only work if all your network calls are serial, what if they are concurrent, then you might end up refreshing the token multiple times each one invalidating the previous.Dayton
yes, it's possible, in this case we should lock refresh call part to refresh only one timeFreemason

© 2022 - 2024 — McMap. All rights reserved.