From b99caadfab9a2f0e9587601c1b9440b2cd3c5a86 Mon Sep 17 00:00:00 2001
From: Junyu Lai <junyulai@google.com>
Date: Mon, 22 Apr 2024 16:22:36 +0800
Subject: [PATCH] Fix NetworkCallback#onAvailable does not match SSID

The test checks whether the client is connected to the hotspot
by register network callback with specified hotspot SSID in the
WifiNetworkSpecifier. However, SSID is not matched when
WifiNetworkAgentSpecifier#canBeSatisfiedBy due to privacy
consideraions.

Thus, manually match SSID is needed to make sure the client
is connected to the hotspot to ensure follow-up tests
can be performed correctly.

Test: m connectivity_multi_devices_snippet && \
      atest CtsConnectivityMultiDevicesTestCases:ConnectivityMultiDevicesTest \
      --rerun-until-failure 100
Bug: 335368434
Change-Id: Ic8dcb7e5c1cf346e4b149367048b6e914894b6d4
---
 .../ConnectivityMultiDevicesSnippet.kt        | 50 ++++++++++---------
 1 file changed, 27 insertions(+), 23 deletions(-)

diff --git a/tests/cts/multidevices/snippet/ConnectivityMultiDevicesSnippet.kt b/tests/cts/multidevices/snippet/ConnectivityMultiDevicesSnippet.kt
index 258648f6b1..543dc03790 100644
--- a/tests/cts/multidevices/snippet/ConnectivityMultiDevicesSnippet.kt
+++ b/tests/cts/multidevices/snippet/ConnectivityMultiDevicesSnippet.kt
@@ -16,6 +16,7 @@
 
 package com.google.snippet.connectivity
 
+import android.Manifest.permission.NETWORK_SETTINGS
 import android.Manifest.permission.OVERRIDE_WIFI_CONFIG
 import android.content.pm.PackageManager.FEATURE_TELEPHONY
 import android.content.pm.PackageManager.FEATURE_WIFI
@@ -30,20 +31,24 @@ import android.net.wifi.ScanResult
 import android.net.wifi.SoftApConfiguration
 import android.net.wifi.SoftApConfiguration.SECURITY_TYPE_WPA2_PSK
 import android.net.wifi.WifiConfiguration
+import android.net.wifi.WifiInfo
 import android.net.wifi.WifiManager
 import android.net.wifi.WifiNetworkSpecifier
 import android.net.wifi.WifiSsid
 import androidx.test.platform.app.InstrumentationRegistry
+import com.android.testutils.AutoReleaseNetworkCallbackRule
 import com.android.testutils.ConnectUtil
 import com.android.testutils.NetworkCallbackHelper
-import com.android.testutils.RecorderCallback.CallbackEntry.Available
 import com.android.testutils.RecorderCallback.CallbackEntry.CapabilitiesChanged
 import com.android.testutils.TestableNetworkCallback
 import com.android.testutils.runAsShell
 import com.google.android.mobly.snippet.Snippet
 import com.google.android.mobly.snippet.rpc.Rpc
+import org.junit.Rule
 
 class ConnectivityMultiDevicesSnippet : Snippet {
+    @get:Rule
+    val networkCallbackRule = AutoReleaseNetworkCallbackRule()
     private val context = InstrumentationRegistry.getInstrumentation().getTargetContext()
     private val wifiManager = context.getSystemService(WifiManager::class.java)!!
     private val cm = context.getSystemService(ConnectivityManager::class.java)!!
@@ -90,8 +95,6 @@ class ConnectivityMultiDevicesSnippet : Snippet {
     @Suppress("DEPRECATION")
     fun connectToWifi(ssid: String, passphrase: String, requireValidation: Boolean): Network {
         val specifier = WifiNetworkSpecifier.Builder()
-            .setSsid(ssid)
-            .setWpa2Passphrase(passphrase)
             .setBand(ScanResult.WIFI_BAND_24_GHZ)
             .build()
         val wifiConfig = WifiConfiguration()
@@ -102,27 +105,28 @@ class ConnectivityMultiDevicesSnippet : Snippet {
         wifiConfig.allowedPairwiseCiphers.set(WifiConfiguration.PairwiseCipher.TKIP)
         wifiConfig.allowedPairwiseCiphers.set(WifiConfiguration.PairwiseCipher.CCMP)
 
-        // Register network callback for the specific wifi.
+        // Add the test configuration and connect to it.
+        val connectUtil = ConnectUtil(context)
+        connectUtil.connectToWifiConfig(wifiConfig)
+
+        // Implement manual SSID matching. Specifying the SSID in
+        // NetworkSpecifier is ineffective
+        // (see WifiNetworkAgentSpecifier#canBeSatisfiedBy for details).
+        // Note that holding permission is necessary when waiting for
+        // the callbacks. The handler thread checks permission; if
+        // it's not present, the SSID will be redacted.
         val networkCallback = TestableNetworkCallback()
-        val wifiRequest = NetworkRequest.Builder().addTransportType(TRANSPORT_WIFI)
-            .setNetworkSpecifier(specifier)
-            .build()
-        cm.registerNetworkCallback(wifiRequest, networkCallback)
-
-        try {
-            // Add the test configuration and connect to it.
-            val connectUtil = ConnectUtil(context)
-            connectUtil.connectToWifiConfig(wifiConfig)
-
-            val event = networkCallback.expect<Available>()
-            if (requireValidation) {
-                networkCallback.eventuallyExpect<CapabilitiesChanged> {
-                    it.caps.hasCapability(NET_CAPABILITY_VALIDATED)
-                }
-            }
-            return event.network
-        } finally {
-            cm.unregisterNetworkCallback(networkCallback)
+        val wifiRequest = NetworkRequest.Builder().addTransportType(TRANSPORT_WIFI).build()
+        return runAsShell(NETWORK_SETTINGS) {
+            // Register the network callback is needed here.
+            // This is to avoid the race condition where callback is fired before
+            // acquiring permission.
+            networkCallbackRule.registerNetworkCallback(wifiRequest, networkCallback)
+            return@runAsShell networkCallback.eventuallyExpect<CapabilitiesChanged> {
+                // Remove double quotes.
+                val ssidFromCaps = (WifiInfo::sanitizeSsid)(it.caps.ssid)
+                ssidFromCaps == ssid && it.caps.hasCapability(NET_CAPABILITY_VALIDATED)
+            }.network
         }
     }