Refactor enforceMeteredApnPolicy

Address review comment on aosp/3027487
This CL also adds test for enforceMeteredApnPolicy

Test: CSBlockedReasonsTest
Bug: 344800107
Change-Id: Ib531bf5b34566fc91fec9f977c5257b332189706
diff --git a/service/src/com/android/server/ConnectivityService.java b/service/src/com/android/server/ConnectivityService.java
index be1d3c7..a4983b3 100755
--- a/service/src/com/android/server/ConnectivityService.java
+++ b/service/src/com/android/server/ConnectivityService.java
@@ -8218,21 +8218,13 @@
             // Policy already enforced.
             return;
         }
-        if (mDeps.isAtLeastV()) {
-            if (mBpfNetMaps.isUidRestrictedOnMeteredNetworks(uid)) {
-                // If UID is restricted, don't allow them to bring up metered APNs.
-                networkCapabilities.addCapability(NET_CAPABILITY_NOT_METERED);
-            }
-            return;
-        }
-        final long ident = Binder.clearCallingIdentity();
-        try {
-            if (mPolicyManager.isUidRestrictedOnMeteredNetworks(uid)) {
-                // If UID is restricted, don't allow them to bring up metered APNs.
-                networkCapabilities.addCapability(NET_CAPABILITY_NOT_METERED);
-            }
-        } finally {
-            Binder.restoreCallingIdentity(ident);
+        final boolean isRestrictedOnMeteredNetworks = mDeps.isAtLeastV()
+                ? mBpfNetMaps.isUidRestrictedOnMeteredNetworks(uid)
+                : BinderUtils.withCleanCallingIdentity(() ->
+                        mPolicyManager.isUidRestrictedOnMeteredNetworks(uid));
+        if (isRestrictedOnMeteredNetworks) {
+            // If UID is restricted, don't allow them to bring up metered APNs.
+            networkCapabilities.addCapability(NET_CAPABILITY_NOT_METERED);
         }
     }
 
diff --git a/tests/unit/java/com/android/server/connectivityservice/CSBlockedReasonsTest.kt b/tests/unit/java/com/android/server/connectivityservice/CSBlockedReasonsTest.kt
index 3ad8de8..985d403 100644
--- a/tests/unit/java/com/android/server/connectivityservice/CSBlockedReasonsTest.kt
+++ b/tests/unit/java/com/android/server/connectivityservice/CSBlockedReasonsTest.kt
@@ -419,4 +419,30 @@
         deps.setChangeIdEnabled(true, NETWORK_BLOCKED_WITHOUT_INTERNET_PERMISSION)
         doTestBlockedReasonsNoInternetPermission(blockedByNoInternetPermission = true)
     }
+
+    private fun doTestEnforceMeteredApnPolicy(restricted: Boolean) {
+        doReturn(restricted).`when`(bpfNetMaps).isUidRestrictedOnMeteredNetworks(Process.myUid())
+
+        val cellAgent = Agent(nc = cellNc())
+        cellAgent.connect()
+        val cb = TestableNetworkCallback()
+        cm.requestNetwork(cellRequest(), cb)
+
+        if (restricted) {
+            waitForIdle()
+            cb.assertNoCallback()
+        } else {
+            cb.expectAvailableCallbacks(cellAgent.network, validated = false)
+        }
+    }
+
+    @Test
+    fun testEnforceMeteredApnPolicy_restricted() {
+        doTestEnforceMeteredApnPolicy(restricted = true)
+    }
+
+    @Test
+    fun testEnforceMeteredApnPolicy_notRestricted() {
+        doTestEnforceMeteredApnPolicy(restricted = false)
+    }
 }