From e944af5bd6c562bdaaa2ee69a25cf541c3f29cff Mon Sep 17 00:00:00 2001
From: Patrick Rohr <prohr@google.com>
Date: Mon, 7 Dec 2020 22:45:48 +0100
Subject: [PATCH] Refactor updating rules for all apps in
 NetworkPolicyManagerService

This gives more flexibility to the user (e.g. will allow restricted
networking mode to update all firewall rules at once) and removes some
unnecessary code.

Bug: 170322816
Bug: 157505406
Test: atest NeworkPolicyManagerServiceTest
Change-Id: Id31a60ad37c280bfb4f5bacf63aafe9de28c5e62
Merged-In: Id31a60ad37c280bfb4f5bacf63aafe9de28c5e62
---
 .../net/NetworkPolicyManagerService.java      | 35 +++++--------------
 1 file changed, 8 insertions(+), 27 deletions(-)

diff --git a/services/core/java/com/android/server/net/NetworkPolicyManagerService.java b/services/core/java/com/android/server/net/NetworkPolicyManagerService.java
index bd80befe84be4..a039215487f82 100644
--- a/services/core/java/com/android/server/net/NetworkPolicyManagerService.java
+++ b/services/core/java/com/android/server/net/NetworkPolicyManagerService.java
@@ -270,6 +270,7 @@ import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
+import java.util.function.IntConsumer;
 
 /**
  * Service that maintains low-level network policy rules, using
@@ -4050,7 +4051,8 @@ public class NetworkPolicyManagerService extends INetworkPolicyManager.Stub {
         try {
             updateRulesForDeviceIdleUL();
             updateRulesForPowerSaveUL();
-            updateRulesForAllAppsUL(TYPE_RESTRICT_POWER);
+            forEachUid("updateRulesForRestrictPower",
+                    uid -> updateRulesForPowerRestrictionsUL(uid));
         } finally {
             Trace.traceEnd(Trace.TRACE_TAG_NETWORK);
         }
@@ -4060,31 +4062,19 @@ public class NetworkPolicyManagerService extends INetworkPolicyManager.Stub {
     private void updateRulesForRestrictBackgroundUL() {
         Trace.traceBegin(Trace.TRACE_TAG_NETWORK, "updateRulesForRestrictBackgroundUL");
         try {
-            updateRulesForAllAppsUL(TYPE_RESTRICT_BACKGROUND);
+            forEachUid("updateRulesForRestrictBackground",
+                    uid -> updateRulesForDataUsageRestrictionsUL(uid));
         } finally {
             Trace.traceEnd(Trace.TRACE_TAG_NETWORK);
         }
     }
 
-    private static final int TYPE_RESTRICT_BACKGROUND = 1;
-    private static final int TYPE_RESTRICT_POWER = 2;
-    @Retention(RetentionPolicy.SOURCE)
-    @IntDef(flag = false, value = {
-            TYPE_RESTRICT_BACKGROUND,
-            TYPE_RESTRICT_POWER,
-    })
-    public @interface RestrictType {
-    }
-
-    // TODO: refactor / consolidate all those updateXyz methods, there are way too many of them...
-    @GuardedBy("mUidRulesFirstLock")
-    private void updateRulesForAllAppsUL(@RestrictType int type) {
+    private void forEachUid(String tag, IntConsumer consumer) {
         if (Trace.isTagEnabled(Trace.TRACE_TAG_NETWORK)) {
-            Trace.traceBegin(Trace.TRACE_TAG_NETWORK, "updateRulesForRestrictPowerUL-" + type);
+            Trace.traceBegin(Trace.TRACE_TAG_NETWORK, "forEachUid-" + tag);
         }
         try {
             // update rules for all installed applications
-
             final PackageManager pm = mContext.getPackageManager();
             final List<UserInfo> users;
             final List<ApplicationInfo> apps;
@@ -4112,16 +4102,7 @@ public class NetworkPolicyManagerService extends INetworkPolicyManager.Stub {
                 for (int j = 0; j < appsSize; j++) {
                     final ApplicationInfo app = apps.get(j);
                     final int uid = UserHandle.getUid(user.id, app.uid);
-                    switch (type) {
-                        case TYPE_RESTRICT_BACKGROUND:
-                            updateRulesForDataUsageRestrictionsUL(uid);
-                            break;
-                        case TYPE_RESTRICT_POWER:
-                            updateRulesForPowerRestrictionsUL(uid);
-                            break;
-                        default:
-                            Slog.w(TAG, "Invalid type for updateRulesForAllApps: " + type);
-                    }
+                    consumer.accept(uid);
                 }
             }
         } finally {