Merge "Wifi: Fix connectivity issues with PSK-SHA256+SAE mode APs" into qt-qpr1-dev

2 files changed, 78 insertions, 2 deletions

diff --git a/service/java/com/android/server/wifi/util/ScanResultUtil.java b/service/java/com/android/server/wifi/util/ScanResultUtil.java
index 39e9d2c40..b92483899 100644
--- a/service/java/com/android/server/wifi/util/ScanResultUtil.java
+++ b/service/java/com/android/server/wifi/util/ScanResultUtil.java
@@ -104,10 +104,10 @@ public class ScanResultUtil {
 
     /**
      * Helper method to check if the provided |scanResult| corresponds to PSK-SAE transition
-     * network. This checks if the provided capabilities string contains PSK+SAE or not.
+     * network. This checks if the provided capabilities string contains both PSK and SAE or not.
      */
     public static boolean isScanResultForPskSaeTransitionNetwork(ScanResult scanResult) {
-        return scanResult.capabilities.contains("PSK+SAE");
+        return scanResult.capabilities.contains("PSK") && scanResult.capabilities.contains("SAE");
     }
 
     /**
diff --git a/tests/wifitests/src/com/android/server/wifi/util/ScanResultUtilTest.java b/tests/wifitests/src/com/android/server/wifi/util/ScanResultUtilTest.java
index 45adffd1d..266a2ce23 100644
--- a/tests/wifitests/src/com/android/server/wifi/util/ScanResultUtilTest.java
+++ b/tests/wifitests/src/com/android/server/wifi/util/ScanResultUtilTest.java
@@ -116,6 +116,82 @@ public class ScanResultUtilTest {
         assertTrue(config.allowedKeyManagement.get(WifiConfiguration.KeyMgmt.IEEE8021X));
     }
 
+    /**
+     * Test that a PSK-SHA256+SAE network is detected as transition mode
+     */
+    @Test
+    public void testPskSha256SaeTransitionModeCheck() {
+        final String ssid = "WPA3-Transition";
+        String caps = "[WPA2-FT/PSK-CCMP][RSN-FT/PSK+PSK-SHA256+SAE+FT/SAE-CCMP][ESS][WPS]";
+
+        ScanResult input = new ScanResult(WifiSsid.createFromAsciiEncoded(ssid), ssid,
+                "ab:cd:01:ef:45:89", 1245, 0, caps, -78, 2450, 1025, 22, 33, 20, 0,
+                0, true);
+
+        input.informationElements = new InformationElement[] {
+                createIE(InformationElement.EID_SSID, ssid.getBytes(StandardCharsets.UTF_8))
+        };
+
+        assertTrue(ScanResultUtil.isScanResultForPskSaeTransitionNetwork(input));
+    }
+
+    /**
+     * Test that a PSK+SAE network is detected as transition mode
+     */
+    @Test
+    public void testPskSaeTransitionModeCheck() {
+        final String ssid = "WPA3-Transition";
+        String caps = "[WPA2-FT/PSK+PSK+SAE+FT/SAE-CCMP][ESS][WPS]";
+
+        ScanResult input = new ScanResult(WifiSsid.createFromAsciiEncoded(ssid), ssid,
+                "ab:cd:01:ef:45:89", 1245, 0, caps, -78, 2450, 1025, 22, 33, 20, 0,
+                0, true);
+
+        input.informationElements = new InformationElement[] {
+                createIE(InformationElement.EID_SSID, ssid.getBytes(StandardCharsets.UTF_8))
+        };
+
+        assertTrue(ScanResultUtil.isScanResultForPskSaeTransitionNetwork(input));
+    }
+
+    /**
+     * Test that a PSK network is not detected as transition mode
+     */
+    @Test
+    public void testPskNotInTransitionModeCheck() {
+        final String ssid = "WPA2-Network";
+        String caps = "[WPA2-FT/PSK+PSK][ESS][WPS]";
+
+        ScanResult input = new ScanResult(WifiSsid.createFromAsciiEncoded(ssid), ssid,
+                "ab:cd:01:ef:45:89", 1245, 0, caps, -78, 2450, 1025, 22, 33, 20, 0,
+                0, true);
+
+        input.informationElements = new InformationElement[] {
+                createIE(InformationElement.EID_SSID, ssid.getBytes(StandardCharsets.UTF_8))
+        };
+
+        assertFalse(ScanResultUtil.isScanResultForPskSaeTransitionNetwork(input));
+    }
+
+    /**
+     * Test that an SAE network is not detected as transition mode
+     */
+    @Test
+    public void testSaeNotInTransitionModeCheck() {
+        final String ssid = "WPA3-Network";
+        String caps = "[WPA2-FT/SAE+SAE][ESS][WPS]";
+
+        ScanResult input = new ScanResult(WifiSsid.createFromAsciiEncoded(ssid), ssid,
+                "ab:cd:01:ef:45:89", 1245, 0, caps, -78, 2450, 1025, 22, 33, 20, 0,
+                0, true);
+
+        input.informationElements = new InformationElement[] {
+                createIE(InformationElement.EID_SSID, ssid.getBytes(StandardCharsets.UTF_8))
+        };
+
+        assertFalse(ScanResultUtil.isScanResultForPskSaeTransitionNetwork(input));
+    }
+
     private static InformationElement createIE(int id, byte[] bytes) {
         InformationElement ie = new InformationElement();
         ie.id = id;