Add report header to failure diagnostics

Add a one line header of JSON info about the failure in failure
diagnostics reports.

This can then be automatically collected to calculate statistics about
particular failures, for example how often they happen when wifi is not
connected, not validated etc.

Test: atest
Change-Id: Ic4c99c53fc2686aaf2f1651522aa2f0d82d0cb86
diff --git a/staticlibs/testutils/devicetests/com/android/testutils/AutoReleaseNetworkCallbackRule.kt b/staticlibs/testutils/devicetests/com/android/testutils/AutoReleaseNetworkCallbackRule.kt
index 93422ad..be6947f 100644
--- a/staticlibs/testutils/devicetests/com/android/testutils/AutoReleaseNetworkCallbackRule.kt
+++ b/staticlibs/testutils/devicetests/com/android/testutils/AutoReleaseNetworkCallbackRule.kt
@@ -98,10 +98,10 @@
         cellRequestCb = null
     }
 
-    private fun addCallback(
-        cb: TestableNetworkCallback,
-        registrar: (TestableNetworkCallback) -> Unit
-    ): TestableNetworkCallback {
+    private fun <T> addCallback(
+        cb: T,
+        registrar: (NetworkCallback) -> Unit
+    ): T where T : NetworkCallback {
         registrar(cb)
         cbToCleanup.add(cb)
         return cb
@@ -142,17 +142,24 @@
     /**
      * File a callback for a NetworkRequest.
      *
-     * This will fail tests (throw) if the cell network cannot be obtained, or if it was already
-     * requested.
-     *
      * Tests may call [unregisterNetworkCallback] once they are done using the returned [Network],
      * otherwise it will be automatically unrequested after the test.
      */
     @JvmOverloads
     fun registerNetworkCallback(
+        request: NetworkRequest
+    ): TestableNetworkCallback = registerNetworkCallback(request, TestableNetworkCallback())
+
+    /**
+     * File a callback for a NetworkRequest.
+     *
+     * Tests may call [unregisterNetworkCallback] once they are done using the returned [Network],
+     * otherwise it will be automatically unrequested after the test.
+     */
+    fun <T> registerNetworkCallback(
         request: NetworkRequest,
-        cb: TestableNetworkCallback = TestableNetworkCallback()
-    ) = addCallback(cb) { cm.registerNetworkCallback(request, it) }
+        cb: T
+    ) where T : NetworkCallback = addCallback(cb) { cm.registerNetworkCallback(request, it) }
 
     /**
      * @see ConnectivityManager.registerDefaultNetworkCallback
diff --git a/staticlibs/testutils/devicetests/com/android/testutils/ConnectivityDiagnosticsCollector.kt b/staticlibs/testutils/devicetests/com/android/testutils/ConnectivityDiagnosticsCollector.kt
index f5a5b4d..ea86281 100644
--- a/staticlibs/testutils/devicetests/com/android/testutils/ConnectivityDiagnosticsCollector.kt
+++ b/staticlibs/testutils/devicetests/com/android/testutils/ConnectivityDiagnosticsCollector.kt
@@ -16,19 +16,46 @@
 
 package com.android.testutils
 
+import android.Manifest.permission.NETWORK_SETTINGS
+import android.Manifest.permission.READ_PRIVILEGED_PHONE_STATE
+import android.content.pm.PackageManager.FEATURE_TELEPHONY
+import android.content.pm.PackageManager.FEATURE_WIFI
 import android.device.collectors.BaseMetricListener
 import android.device.collectors.DataRecord
+import android.net.ConnectivityManager.NetworkCallback
+import android.net.ConnectivityManager.NetworkCallback.FLAG_INCLUDE_LOCATION_INFO
+import android.net.Network
+import android.net.NetworkCapabilities
+import android.net.NetworkCapabilities.NET_CAPABILITY_INTERNET
+import android.net.NetworkCapabilities.NET_CAPABILITY_VALIDATED
+import android.net.NetworkCapabilities.TRANSPORT_CELLULAR
+import android.net.NetworkCapabilities.TRANSPORT_VPN
+import android.net.NetworkCapabilities.TRANSPORT_WIFI
+import android.net.NetworkRequest
+import android.net.wifi.WifiInfo
+import android.net.wifi.WifiManager
 import android.os.Build
 import android.os.ParcelFileDescriptor
+import android.telephony.TelephonyManager
+import android.telephony.TelephonyManager.SIM_STATE_UNKNOWN
 import android.util.Log
+import androidx.annotation.RequiresApi
 import androidx.test.platform.app.InstrumentationRegistry
+import com.android.modules.utils.build.SdkLevel.isAtLeastS
 import java.io.ByteArrayOutputStream
 import java.io.File
+import java.io.FileOutputStream
 import java.io.PrintWriter
 import java.time.ZonedDateTime
+import java.time.format.DateTimeFormatter
+import java.util.concurrent.CompletableFuture
+import java.util.concurrent.TimeUnit
+import java.util.concurrent.TimeoutException
 import kotlin.test.assertNull
+import org.json.JSONObject
 import org.junit.AssumptionViolatedException
 import org.junit.runner.Description
+import org.junit.runner.Result
 import org.junit.runner.notification.Failure
 
 /**
@@ -52,11 +79,35 @@
         var instance: ConnectivityDiagnosticsCollector? = null
     }
 
+    private var failureHeader: String? = null
     private val buffer = ByteArrayOutputStream()
     private val collectorDir: File by lazy {
         createAndEmptyDirectory(COLLECTOR_DIR)
     }
     private val outputFiles = mutableSetOf<String>()
+    private val cbHelper = NetworkCallbackHelper()
+    private val networkCallback = MonitoringNetworkCallback()
+
+    inner class MonitoringNetworkCallback : NetworkCallback() {
+        val currentMobileDataNetworks = mutableMapOf<Network, NetworkCapabilities>()
+        val currentVpnNetworks = mutableMapOf<Network, NetworkCapabilities>()
+        val currentWifiNetworks = mutableMapOf<Network, NetworkCapabilities>()
+
+        override fun onLost(network: Network) {
+            currentWifiNetworks.remove(network)
+            currentMobileDataNetworks.remove(network)
+        }
+
+        override fun onCapabilitiesChanged(network: Network, nc: NetworkCapabilities) {
+            if (nc.hasTransport(TRANSPORT_VPN)) {
+                currentVpnNetworks[network] = nc
+            } else if (nc.hasTransport(TRANSPORT_WIFI)) {
+                currentWifiNetworks[network] = nc
+            } else if (nc.hasTransport(TRANSPORT_CELLULAR)) {
+                currentMobileDataNetworks[network] = nc
+            }
+        }
+    }
 
     override fun onSetUp() {
         assertNull(instance, "ConnectivityDiagnosticsCollectors were set up multiple times")
@@ -72,6 +123,24 @@
         instance = null
     }
 
+    override fun onTestRunStart(runData: DataRecord?, description: Description?) {
+        runAsShell(NETWORK_SETTINGS) {
+            cbHelper.registerNetworkCallback(
+                NetworkRequest.Builder()
+                    .addCapability(NET_CAPABILITY_INTERNET)
+                    .addTransportType(TRANSPORT_WIFI)
+                    .addTransportType(TRANSPORT_CELLULAR)
+                    .build(), networkCallback
+            )
+        }
+    }
+
+    override fun onTestRunEnd(runData: DataRecord?, result: Result?) {
+        // onTestRunEnd is called regardless of success/failure, and the Result contains summary of
+        // run/failed/ignored... tests.
+        cbHelper.unregisterAll()
+    }
+
     override fun onTestFail(testData: DataRecord, description: Description, failure: Failure) {
         // TODO: find a way to disable this behavior only on local runs, to avoid slowing them down
         // when iterating on failing tests.
@@ -119,12 +188,116 @@
         }
         val outFile = File(collectorDir, filename + FILENAME_SUFFIX)
         outputFiles.add(filename)
-        outFile.writeBytes(buffer.toByteArray())
+        FileOutputStream(outFile).use { fos ->
+            failureHeader?.let {
+                fos.write(it.toByteArray())
+                fos.write("\n".toByteArray())
+            }
+            fos.write(buffer.toByteArray())
+        }
+        failureHeader = null
         buffer.reset()
         val fileKey = "${ConnectivityDiagnosticsCollector::class.qualifiedName}_$filename"
         testData.addFileMetric(fileKey, outFile)
     }
 
+    private fun maybeCollectFailureHeader() {
+        if (failureHeader != null) {
+            Log.i(TAG, "Connectivity diagnostics failure header already collected, skipping")
+            return
+        }
+
+        val instr = InstrumentationRegistry.getInstrumentation()
+        val ctx = instr.context
+        val pm = ctx.packageManager
+        val hasWifi = pm.hasSystemFeature(FEATURE_WIFI)
+        val hasMobileData = pm.hasSystemFeature(FEATURE_TELEPHONY)
+        val tm = if (hasMobileData) ctx.getSystemService(TelephonyManager::class.java) else null
+        // getAdoptedShellPermissions is S+. Optimistically assume that tests are not holding on
+        // shell permissions during failure/cleanup on R.
+        val canUseShell = !isAtLeastS() ||
+                instr.uiAutomation.getAdoptedShellPermissions().isNullOrEmpty()
+        val headerObj = JSONObject()
+        if (canUseShell) {
+            runAsShell(READ_PRIVILEGED_PHONE_STATE, NETWORK_SETTINGS) {
+                headerObj.apply {
+                    put("deviceSerial", Build.getSerial())
+                    // The network callback filed on start cannot get the WifiInfo as it would need
+                    // to keep NETWORK_SETTINGS permission throughout the test run. Try to
+                    // obtain it while holding the permission at the end of the test.
+                    val wifiInfo = networkCallback.currentWifiNetworks.keys.firstOrNull()?.let {
+                        getWifiInfo(it)
+                    }
+                    put("ssid", wifiInfo?.ssid)
+                    put("bssid", wifiInfo?.bssid)
+                    put("simState", tm?.simState ?: SIM_STATE_UNKNOWN)
+                    put("mccMnc", tm?.simOperator)
+                }
+            }
+        } else {
+            Log.w(TAG, "The test is still holding shell permissions, cannot collect privileged " +
+                    "device info")
+            headerObj.put("shellPermissionsUnavailable", true)
+        }
+        failureHeader = headerObj.apply {
+            put("time", DateTimeFormatter.ISO_OFFSET_DATE_TIME.format(ZonedDateTime.now()))
+            put(
+                "wifiEnabled",
+                hasWifi && ctx.getSystemService(WifiManager::class.java).isWifiEnabled
+            )
+            put("connectedWifiCount", networkCallback.currentWifiNetworks.size)
+            put("validatedWifiCount", networkCallback.currentWifiNetworks.filterValues {
+                it.hasCapability(NET_CAPABILITY_VALIDATED)
+            }.size)
+            put("mobileDataConnectivityPossible", tm?.isDataConnectivityPossible ?: false)
+            put("connectedMobileDataCount", networkCallback.currentMobileDataNetworks.size)
+            put("validatedMobileDataCount",
+                networkCallback.currentMobileDataNetworks.filterValues {
+                    it.hasCapability(NET_CAPABILITY_VALIDATED)
+                }.size
+            )
+        }.toString()
+    }
+
+    private class WifiInfoCallback : NetworkCallback {
+        private val network: Network
+        val wifiInfoFuture = CompletableFuture<WifiInfo?>()
+        constructor(network: Network) : super() {
+            this.network = network
+        }
+        @RequiresApi(Build.VERSION_CODES.S)
+        constructor(network: Network, flags: Int) : super(flags) {
+            this.network = network
+        }
+        override fun onCapabilitiesChanged(net: Network, nc: NetworkCapabilities) {
+            if (network == net) {
+                wifiInfoFuture.complete(nc.transportInfo as? WifiInfo)
+            }
+        }
+    }
+
+    private fun getWifiInfo(network: Network): WifiInfo? {
+        // Get the SSID via network callbacks, as the Networks are obtained via callbacks, and
+        // synchronous calls (CM#getNetworkCapabilities) and callbacks should not be mixed.
+        // A new callback needs to be filed and received while holding NETWORK_SETTINGS permission.
+        val cb = if (isAtLeastS()) {
+            WifiInfoCallback(network, FLAG_INCLUDE_LOCATION_INFO)
+        } else {
+            WifiInfoCallback(network)
+        }
+        cbHelper.registerNetworkCallback(
+            NetworkRequest.Builder()
+                .addTransportType(TRANSPORT_WIFI)
+                .addCapability(NET_CAPABILITY_INTERNET).build(), cb)
+        return try {
+            cb.wifiInfoFuture.get(1L, TimeUnit.SECONDS)
+        } catch (e: TimeoutException) {
+            null
+        } finally {
+            cbHelper.unregisterNetworkCallback(cb)
+        }
+    }
+
     /**
      * Add connectivity diagnostics to the test data dump.
      *
@@ -134,6 +307,7 @@
      * @param exceptionContext An exception to write a stacktrace to the dump for context.
      */
     fun collectTestFailureDiagnostics(exceptionContext: Throwable? = null) {
+        maybeCollectFailureHeader()
         collectDumpsysConnectivity(exceptionContext)
     }