From d19e26ae4452cf25e8c29bf08412e1572e41b063 Mon Sep 17 00:00:00 2001
From: changbetty <changbetty@google.com>
Date: Thu, 23 Dec 2021 15:22:37 +0000
Subject: [PATCH] To add restriction check for Wi-Fi Direct

Bug: 203170358
Test: make RunSettingsLibRoboTests ROBOTEST_FILTER=WifiEnterpriseRestrictionUtilsTest
Change-Id: I03991622d92e28b4a804be4682bc4058c5c37098
---
 .../wifi/WifiEnterpriseRestrictionUtils.java  | 18 ++++++++++++-
 .../WifiEnterpriseRestrictionUtilsTest.java   | 26 +++++++++++++++++++
 2 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/packages/SettingsLib/src/com/android/settingslib/wifi/WifiEnterpriseRestrictionUtils.java b/packages/SettingsLib/src/com/android/settingslib/wifi/WifiEnterpriseRestrictionUtils.java
index 24233729ad450..21a4ac6220fb2 100644
--- a/packages/SettingsLib/src/com/android/settingslib/wifi/WifiEnterpriseRestrictionUtils.java
+++ b/packages/SettingsLib/src/com/android/settingslib/wifi/WifiEnterpriseRestrictionUtils.java
@@ -29,7 +29,7 @@ public class WifiEnterpriseRestrictionUtils {
     private static final String TAG = "WifiEntResUtils";
 
     /**
-     * Confirm Wi-Fi tethering is available according to whether user restriction is set
+     * Confirm Wi-Fi tethering is allowed according to whether user restriction is set
      *
      * @param context A context
      * @return whether the device is permitted to use Wi-Fi Tethering
@@ -44,6 +44,22 @@ public class WifiEnterpriseRestrictionUtils {
         return true;
     }
 
+    /**
+     * Confirm Wi-Fi Direct is allowed according to whether user restriction is set
+     *
+     * @param context A context
+     * @return whether the device is permitted to use Wi-Fi Direct
+     */
+    public static boolean isWifiDirectAllowed(Context context) {
+        final UserManager userManager = context.getSystemService(UserManager.class);
+        final Bundle restrictions = userManager.getUserRestrictions();
+        if (isAtLeastT() && restrictions.getBoolean(UserManager.DISALLOW_WIFI_DIRECT)) {
+            Log.i(TAG, "Wi-Fi Direct isn't available due to user restriction.");
+            return false;
+        }
+        return true;
+    }
+
     @ChecksSdkIntAtLeast(api=Build.VERSION_CODES.TIRAMISU)
     private static boolean isAtLeastT() {
         return Build.VERSION.SDK_INT >= Build.VERSION_CODES.TIRAMISU;
diff --git a/packages/SettingsLib/tests/robotests/src/com/android/settingslib/wifi/WifiEnterpriseRestrictionUtilsTest.java b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/wifi/WifiEnterpriseRestrictionUtilsTest.java
index 455dc64fa97d9..3c339defbaf45 100644
--- a/packages/SettingsLib/tests/robotests/src/com/android/settingslib/wifi/WifiEnterpriseRestrictionUtilsTest.java
+++ b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/wifi/WifiEnterpriseRestrictionUtilsTest.java
@@ -77,4 +77,30 @@ public class WifiEnterpriseRestrictionUtilsTest {
 
         assertThat(WifiEnterpriseRestrictionUtils.isWifiTetheringAllowed(mContext)).isTrue();
     }
+
+    @Test
+    public void isWifiDirectAllowed_setSDKForS_shouldReturnTrue() {
+        ReflectionHelpers.setStaticField(Build.VERSION.class, "SDK_INT", Build.VERSION_CODES.S);
+        when(mBundle.getBoolean(UserManager.DISALLOW_WIFI_DIRECT)).thenReturn(true);
+
+        assertThat(WifiEnterpriseRestrictionUtils.isWifiDirectAllowed(mContext)).isTrue();
+    }
+
+    @Test
+    public void isWifiDirectAllowed_setSDKForTAndDisallowForRestriction_shouldReturnFalse() {
+        ReflectionHelpers.setStaticField(
+            Build.VERSION.class, "SDK_INT", Build.VERSION_CODES.TIRAMISU);
+        when(mBundle.getBoolean(UserManager.DISALLOW_WIFI_DIRECT)).thenReturn(true);
+
+        assertThat(WifiEnterpriseRestrictionUtils.isWifiDirectAllowed(mContext)).isFalse();
+    }
+
+    @Test
+    public void isWifiDirectAllowed_setSDKForTAndAllowForRestriction_shouldReturnTrue() {
+        ReflectionHelpers.setStaticField(
+            Build.VERSION.class, "SDK_INT", Build.VERSION_CODES.TIRAMISU);
+        when(mBundle.getBoolean(UserManager.DISALLOW_WIFI_DIRECT)).thenReturn(false);
+
+        assertThat(WifiEnterpriseRestrictionUtils.isWifiDirectAllowed(mContext)).isTrue();
+    }
 }