OkHttp auth header interceptor with retry logic

AuthHeaderInterceptor.kt

import androidx.annotation.WorkerThread
import kotlinx.coroutines.runBlocking
import okhttp3.Interceptor
import okhttp3.Request
import okhttp3.Response
import timber.log.Timber

object Retried

/**
 * This will block against infinite retries if multiple interceptors of this type are used. This
 * probably isn't necessary however.
 */
fun Request.hasBeenRetried() = tag(Retried::class.java) != null
fun Request.markAsRetry() = newBuilder().tag(Retried::class.java, Retried).build()

/**
 * Base implementation for an authorization interceptor which does the following:
 * - For every request it will call [getToken] and decorate the request with an auth header
 * - If [getToken] returns `null`, then it will skip authorization
 * - If [getToken] returns a token, but the HTTP response code is in [notAuthorizedHttpCodes],
 * then this will call [refreshToken] and if given a token will attempt to retry the call once.
 */
abstract class AuthHeaderInterceptor : Interceptor {

    open val notAuthorizedHttpCodes: Set<Int> = setOf(401)

    open val authHeaderName = "Authorization"
    open val authHeaderPattern: (token: String) -> String = { "Bearer $it" }

    @WorkerThread
    abstract fun getToken(): String?

    @WorkerThread
    abstract suspend fun refreshToken(): String?

    override fun intercept(chain: Interceptor.Chain): Response {
        val request = chain.request()
        if (request.hasBeenRetried()) return chain.proceed(request)

        val token = safely { getToken() } ?: return chain.proceed(request)

        val authorizedRequest = request.signWithToken(token)
        val response = chain.proceed(authorizedRequest)

        return if (notAuthorizedHttpCodes.contains(response.code)) {
            safely { runBlocking { refreshToken() } }?.let { newToken ->
                response.close()
                chain.proceed(request.signWithToken(newToken).markAsRetry())
            } ?: response
        } else response
    }

    private fun Request.signWithToken(token: String): Request {
        return newBuilder()
            .addHeader(authHeaderName, authHeaderPattern(token))
            .build()
    }

    private fun <T> safely(block: () -> T): T? = runCatching(block)
        .onFailure {
            Timber.e(it, "Exception thrown in $this, caught but should not happen.")
        }.getOrNull()
}