OkHttp Interceptor that ensures requests happen on mobile data even if WiFi is on

MobileDataOnlyInterceptor.kt

import android.Manifest
import android.content.Context
import android.content.pm.PackageManager
import android.net.ConnectivityManager
import android.net.Network
import android.net.NetworkCapabilities
import android.net.NetworkRequest
import android.os.Build
import android.support.annotation.RequiresApi
import android.support.v4.content.ContextCompat
import okhttp3.Interceptor
import okhttp3.OkHttpClient
import okhttp3.Response
import java.io.IOException
import java.net.SocketException
import java.util.concurrent.CountDownLatch
import java.util.concurrent.TimeUnit

private const val MOBILE_NETWORK_REQUEST_TIMEOUT_MS = 5000L

/**
 * Blocks requests not made on a cellular network.
 * May short-circuit the chain if it succeeds getting access to an available cellular network.
 */
class MobileDataOnlyInterceptor(
    private val context: Context,
    private val okHttpClient: OkHttpClient
) : Interceptor {

    @Throws(IOException::class)
    override fun intercept(chain: Interceptor.Chain): Response {
        val originalRequest = chain.request()

        if (context.mobileDataPreferenceSupported) {
            var response: Response? = null
            requestMobileNetwork {
                response = try {
                    okHttpClient.newBuilder()
                        .socketFactory(it.socketFactory)
                        .build()
                        .newCall(originalRequest)
                        .execute()
                        .apply {
                            val byteLimit = 1024L * 1024L // arbitrary, any response should fit
                            // request entire response body while connected to the cellular network
                            body()?.source()?.request(byteLimit)
                        }
                } catch (e: SocketException) {
                    // we may fail to bind sockets to the mobile network; attempt with default network
                    // https://trafiapp.atlassian.net/browse/TB-461
                    null
                }
            }
            response?.let { return it }
        }

        if (context.activeNetworkTypeIsNotMobile) {
            throw IOException("Not on mobile data")
        }

        return chain.proceed(originalRequest)
    }

    @RequiresApi(Build.VERSION_CODES.LOLLIPOP)
    private fun requestMobileNetwork(onAvailable: (Network) -> Unit) {
        val manager = context.getSystemService(Context.CONNECTIVITY_SERVICE) as ConnectivityManager

        val request = NetworkRequest.Builder()
            .addTransportType(NetworkCapabilities.TRANSPORT_CELLULAR)
            .addCapability(NetworkCapabilities.NET_CAPABILITY_INTERNET)
            .build()

        var network: Network? = null
        val latch = CountDownLatch(1)

        val callback = object : ConnectivityManager.NetworkCallback() {
            override fun onAvailable(n: Network) {
                super.onAvailable(n)
                if (latch.count == 0L) return

                network = n
                latch.countDown()
            }
        }

        manager.requestNetwork(request, callback)

        try {
            latch.await(MOBILE_NETWORK_REQUEST_TIMEOUT_MS, TimeUnit.MILLISECONDS)
        } catch (e: InterruptedException) {
        }

        try {
            network?.let { onAvailable(it) }
        } finally {
            manager.unregisterNetworkCallback(callback)
        }
    }
}

private val Context.mobileDataPreferenceSupported: Boolean
    get() = Build.VERSION.SDK_INT >= Build.VERSION_CODES.LOLLIPOP &&
            PackageManager.PERMISSION_GRANTED ==
            ContextCompat.checkSelfPermission(this, Manifest.permission.CHANGE_NETWORK_STATE)

@Suppress("DEPRECATION")
private val Context.activeNetworkTypeIsNotMobile: Boolean
    get() {
        val manager = getSystemService(Context.CONNECTIVITY_SERVICE) as ConnectivityManager
        val isMobile =
            manager.activeNetworkInfo?.let {
                it.isConnectedOrConnecting && it.type == ConnectivityManager.TYPE_MOBILE
            } ?: false
        return !isMobile
    }