Merge "Add new network process capability." am: ae0fa6d359 am: 028e394c23

Original change: https://android-review.googlesource.com/c/platform/frameworks/base/+/1612822

MUST ONLY BE SUBMITTED BY AUTOMERGER

Change-Id: I6894d30d9cc47c3d14596f878b0b91084a4d50e8

core/api/test-current.txt

@@ -82,7 +82,7 @@ package android.app {
     method public static boolean isHighEndGfx();
     method public static void resumeAppSwitches() throws android.os.RemoteException;
     method @RequiresPermission(android.Manifest.permission.CHANGE_CONFIGURATION) public void scheduleApplicationInfoChanged(java.util.List<java.lang.String>, int);
-    field public static final int PROCESS_CAPABILITY_ALL = 7; // 0x7
+    field public static final int PROCESS_CAPABILITY_ALL = 15; // 0xf
     field public static final int PROCESS_CAPABILITY_ALL_EXPLICIT = 1; // 0x1
     field public static final int PROCESS_CAPABILITY_ALL_IMPLICIT = 6; // 0x6
     field public static final int PROCESS_CAPABILITY_FOREGROUND_CAMERA = 2; // 0x2

core/java/android/app/ActivityManager.java

@@ -601,11 +601,15 @@ public class ActivityManager {
     @TestApi
     public static final int PROCESS_CAPABILITY_FOREGROUND_MICROPHONE = 1 << 2;
 
+    /** @hide Process can access network despite any power saving resrictions */
+    public static final int PROCESS_CAPABILITY_NETWORK = 1 << 3;
+
     /** @hide all capabilities, the ORing of all flags in {@link ProcessCapability}*/
     @TestApi
     public static final int PROCESS_CAPABILITY_ALL = PROCESS_CAPABILITY_FOREGROUND_LOCATION
             | PROCESS_CAPABILITY_FOREGROUND_CAMERA
-            | PROCESS_CAPABILITY_FOREGROUND_MICROPHONE;
+            | PROCESS_CAPABILITY_FOREGROUND_MICROPHONE
+            | PROCESS_CAPABILITY_NETWORK;
     /**
      * All explicit capabilities. These are capabilities that need to be specified from manifest
      * file.
@@ -631,6 +635,15 @@ public class ActivityManager {
         pw.print((caps & PROCESS_CAPABILITY_FOREGROUND_LOCATION) != 0 ? 'L' : '-');
         pw.print((caps & PROCESS_CAPABILITY_FOREGROUND_CAMERA) != 0 ? 'C' : '-');
         pw.print((caps & PROCESS_CAPABILITY_FOREGROUND_MICROPHONE) != 0 ? 'M' : '-');
+        pw.print((caps & PROCESS_CAPABILITY_NETWORK) != 0 ? 'N' : '-');
+    }
+
+    /** @hide */
+    public static void printCapabilitiesSummary(StringBuilder sb, @ProcessCapability int caps) {
+        sb.append((caps & PROCESS_CAPABILITY_FOREGROUND_LOCATION) != 0 ? 'L' : '-');
+        sb.append((caps & PROCESS_CAPABILITY_FOREGROUND_CAMERA) != 0 ? 'C' : '-');
+        sb.append((caps & PROCESS_CAPABILITY_FOREGROUND_MICROPHONE) != 0 ? 'M' : '-');
+        sb.append((caps & PROCESS_CAPABILITY_NETWORK) != 0 ? 'N' : '-');
     }
 
     /**
@@ -641,13 +654,21 @@ public class ActivityManager {
         printCapabilitiesSummary(pw, caps);
         final int remain = caps & ~(PROCESS_CAPABILITY_FOREGROUND_LOCATION
                 | PROCESS_CAPABILITY_FOREGROUND_CAMERA
-                | PROCESS_CAPABILITY_FOREGROUND_MICROPHONE);
+                | PROCESS_CAPABILITY_FOREGROUND_MICROPHONE
+                | PROCESS_CAPABILITY_NETWORK);
         if (remain != 0) {
             pw.print('+');
             pw.print(remain);
         }
     }
 
+    /** @hide */
+    public static String getCapabilitiesSummary(@ProcessCapability int caps) {
+        final StringBuilder sb = new StringBuilder();
+        printCapabilitiesSummary(sb, caps);
+        return sb.toString();
+    }
+
     // NOTE: If PROCESS_STATEs are added, then new fields must be added
     // to frameworks/base/core/proto/android/app/enums.proto and the following method must
     // be updated to correctly map between them.
@@ -4777,6 +4798,80 @@ public class ActivityManager {
         }
     }
 
+    /** @hide */
+    public static String procStateToString(int procState) {
+        final String procStateStr;
+        switch (procState) {
+            case ActivityManager.PROCESS_STATE_PERSISTENT:
+                procStateStr = "PER ";
+                break;
+            case ActivityManager.PROCESS_STATE_PERSISTENT_UI:
+                procStateStr = "PERU";
+                break;
+            case ActivityManager.PROCESS_STATE_TOP:
+                procStateStr = "TOP ";
+                break;
+            case ActivityManager.PROCESS_STATE_BOUND_TOP:
+                procStateStr = "BTOP";
+                break;
+            case ActivityManager.PROCESS_STATE_FOREGROUND_SERVICE:
+                procStateStr = "FGS ";
+                break;
+            case ActivityManager.PROCESS_STATE_BOUND_FOREGROUND_SERVICE:
+                procStateStr = "BFGS";
+                break;
+            case ActivityManager.PROCESS_STATE_IMPORTANT_FOREGROUND:
+                procStateStr = "IMPF";
+                break;
+            case ActivityManager.PROCESS_STATE_IMPORTANT_BACKGROUND:
+                procStateStr = "IMPB";
+                break;
+            case ActivityManager.PROCESS_STATE_TRANSIENT_BACKGROUND:
+                procStateStr = "TRNB";
+                break;
+            case ActivityManager.PROCESS_STATE_BACKUP:
+                procStateStr = "BKUP";
+                break;
+            case ActivityManager.PROCESS_STATE_SERVICE:
+                procStateStr = "SVC ";
+                break;
+            case ActivityManager.PROCESS_STATE_RECEIVER:
+                procStateStr = "RCVR";
+                break;
+            case ActivityManager.PROCESS_STATE_TOP_SLEEPING:
+                procStateStr = "TPSL";
+                break;
+            case ActivityManager.PROCESS_STATE_HEAVY_WEIGHT:
+                procStateStr = "HVY ";
+                break;
+            case ActivityManager.PROCESS_STATE_HOME:
+                procStateStr = "HOME";
+                break;
+            case ActivityManager.PROCESS_STATE_LAST_ACTIVITY:
+                procStateStr = "LAST";
+                break;
+            case ActivityManager.PROCESS_STATE_CACHED_ACTIVITY:
+                procStateStr = "CAC ";
+                break;
+            case ActivityManager.PROCESS_STATE_CACHED_ACTIVITY_CLIENT:
+                procStateStr = "CACC";
+                break;
+            case ActivityManager.PROCESS_STATE_CACHED_RECENT:
+                procStateStr = "CRE ";
+                break;
+            case ActivityManager.PROCESS_STATE_CACHED_EMPTY:
+                procStateStr = "CEM ";
+                break;
+            case ActivityManager.PROCESS_STATE_NONEXISTENT:
+                procStateStr = "NONE";
+                break;
+            default:
+                procStateStr = "??";
+                break;
+        }
+        return procStateStr;
+    }
+
     /**
      * The AppTask allows you to manage your own application's tasks.
      * See {@link android.app.ActivityManager#getAppTasks()}

core/java/android/net/NetworkPolicyManager.java

@@ -16,14 +16,17 @@
 
 package android.net;
 
+import static android.app.ActivityManager.procStateToString;
 import static android.content.pm.PackageManager.GET_SIGNATURES;
 
 import android.annotation.IntDef;
 import android.annotation.NonNull;
+import android.annotation.Nullable;
 import android.annotation.RequiresPermission;
 import android.annotation.SystemService;
 import android.annotation.TestApi;
 import android.app.ActivityManager;
+import android.app.ActivityManager.ProcessCapability;
 import android.compat.annotation.UnsupportedAppUsage;
 import android.content.Context;
 import android.content.Intent;
@@ -617,8 +620,18 @@ public class NetworkPolicyManager {
      * to access network when the device is idle or in battery saver mode. Otherwise, false.
      * @hide
      */
-    public static boolean isProcStateAllowedWhileIdleOrPowerSaveMode(int procState) {
+    public static boolean isProcStateAllowedWhileIdleOrPowerSaveMode(@Nullable UidState uidState) {
-        return procState <= FOREGROUND_THRESHOLD_STATE;
+        if (uidState == null) {
+            return false;
+        }
+        return isProcStateAllowedWhileIdleOrPowerSaveMode(uidState.procState, uidState.capability);
+    }
+
+    /** @hide */
+    public static boolean isProcStateAllowedWhileIdleOrPowerSaveMode(
+            int procState, @ProcessCapability int capability) {
+        return procState <= FOREGROUND_THRESHOLD_STATE
+                || (capability & ActivityManager.PROCESS_CAPABILITY_NETWORK) != 0;
     }
 
     /**
@@ -626,10 +639,43 @@ public class NetworkPolicyManager {
      * to access network when the device is in data saver mode. Otherwise, false.
      * @hide
      */
+    public static boolean isProcStateAllowedWhileOnRestrictBackground(@Nullable UidState uidState) {
+        if (uidState == null) {
+            return false;
+        }
+        return isProcStateAllowedWhileOnRestrictBackground(uidState.procState);
+    }
+
+    /** @hide */
     public static boolean isProcStateAllowedWhileOnRestrictBackground(int procState) {
+        // Data saver and bg policy restrictions will only take procstate into account.
         return procState <= FOREGROUND_THRESHOLD_STATE;
     }
 
+    /** @hide */
+    public static final class UidState {
+        public int uid;
+        public int procState;
+        public int capability;
+
+        public UidState(int uid, int procState, int capability) {
+            this.uid = uid;
+            this.procState = procState;
+            this.capability = capability;
+        }
+
+        @Override
+        public String toString() {
+            final StringBuilder sb = new StringBuilder();
+            sb.append("{procState=");
+            sb.append(procStateToString(procState));
+            sb.append(",cap=");
+            ActivityManager.printCapabilitiesSummary(sb, capability);
+            sb.append("}");
+            return sb.toString();
+        }
+    }
+
     /** @hide */
     @TestApi
     @NonNull

services/core/java/com/android/server/am/ProcessList.java

@@ -1016,76 +1016,7 @@ public final class ProcessList {
     }
 
     public static String makeProcStateString(int curProcState) {
-        String procState;
+        return ActivityManager.procStateToString(curProcState);
-        switch (curProcState) {
-            case ActivityManager.PROCESS_STATE_PERSISTENT:
-                procState = "PER ";
-                break;
-            case ActivityManager.PROCESS_STATE_PERSISTENT_UI:
-                procState = "PERU";
-                break;
-            case ActivityManager.PROCESS_STATE_TOP:
-                procState = "TOP ";
-                break;
-            case ActivityManager.PROCESS_STATE_BOUND_TOP:
-                procState = "BTOP";
-                break;
-            case ActivityManager.PROCESS_STATE_FOREGROUND_SERVICE:
-                procState = "FGS ";
-                break;
-            case ActivityManager.PROCESS_STATE_BOUND_FOREGROUND_SERVICE:
-                procState = "BFGS";
-                break;
-            case ActivityManager.PROCESS_STATE_IMPORTANT_FOREGROUND:
-                procState = "IMPF";
-                break;
-            case ActivityManager.PROCESS_STATE_IMPORTANT_BACKGROUND:
-                procState = "IMPB";
-                break;
-            case ActivityManager.PROCESS_STATE_TRANSIENT_BACKGROUND:
-                procState = "TRNB";
-                break;
-            case ActivityManager.PROCESS_STATE_BACKUP:
-                procState = "BKUP";
-                break;
-            case ActivityManager.PROCESS_STATE_SERVICE:
-                procState = "SVC ";
-                break;
-            case ActivityManager.PROCESS_STATE_RECEIVER:
-                procState = "RCVR";
-                break;
-            case ActivityManager.PROCESS_STATE_TOP_SLEEPING:
-                procState = "TPSL";
-                break;
-            case ActivityManager.PROCESS_STATE_HEAVY_WEIGHT:
-                procState = "HVY ";
-                break;
-            case ActivityManager.PROCESS_STATE_HOME:
-                procState = "HOME";
-                break;
-            case ActivityManager.PROCESS_STATE_LAST_ACTIVITY:
-                procState = "LAST";
-                break;
-            case ActivityManager.PROCESS_STATE_CACHED_ACTIVITY:
-                procState = "CAC ";
-                break;
-            case ActivityManager.PROCESS_STATE_CACHED_ACTIVITY_CLIENT:
-                procState = "CACC";
-                break;
-            case ActivityManager.PROCESS_STATE_CACHED_RECENT:
-                procState = "CRE ";
-                break;
-            case ActivityManager.PROCESS_STATE_CACHED_EMPTY:
-                procState = "CEM ";
-                break;
-            case ActivityManager.PROCESS_STATE_NONEXISTENT:
-                procState = "NONE";
-                break;
-            default:
-                procState = "??";
-                break;
-        }
-        return procState;
     }
 
     public static int makeProcStateProtoEnum(int curProcState) {
@@ -3833,10 +3764,13 @@ public final class ProcessList {
     int getBlockStateForUid(UidRecord uidRec) {
         // Denotes whether uid's process state is currently allowed network access.
         final boolean isAllowed =
-                isProcStateAllowedWhileIdleOrPowerSaveMode(uidRec.getCurProcState())
+                isProcStateAllowedWhileIdleOrPowerSaveMode(uidRec.getCurProcState(),
+                        uidRec.curCapability)
                 || isProcStateAllowedWhileOnRestrictBackground(uidRec.getCurProcState());
         // Denotes whether uid's process state was previously allowed network access.
-        final boolean wasAllowed = isProcStateAllowedWhileIdleOrPowerSaveMode(uidRec.setProcState)
+        final boolean wasAllowed =
+                isProcStateAllowedWhileIdleOrPowerSaveMode(uidRec.setProcState,
+                        uidRec.setCapability)
                 || isProcStateAllowedWhileOnRestrictBackground(uidRec.setProcState);
 
         // When the uid is coming to foreground, AMS should inform the app thread that it should

services/core/java/com/android/server/net/NetworkPolicyLogger.java

@@ -29,6 +29,7 @@ import static android.net.NetworkPolicyManager.FIREWALL_RULE_DEFAULT;
 import static android.os.Process.INVALID_UID;
 
 import android.app.ActivityManager;
+import android.app.ActivityManager.ProcessCapability;
 import android.net.NetworkPolicyManager;
 import android.os.UserHandle;
 import android.util.Log;
@@ -97,13 +98,15 @@ public class NetworkPolicyLogger {
         }
     }
 
-    void uidStateChanged(int uid, int procState, long procStateSeq) {
+    void uidStateChanged(int uid, int procState, long procStateSeq,
+            @ProcessCapability int capability) {
         synchronized (mLock) {
             if (LOGV || uid == mDebugUid) {
                 Slog.v(TAG, uid + " state changed to "
-                        + ProcessList.makeProcStateString(procState) + " with seq=" + procStateSeq);
+                        + ProcessList.makeProcStateString(procState) + ",seq=" + procStateSeq
+                        + ",cap=" + ActivityManager.getCapabilitiesSummary(capability));
             }
-            mUidStateChangeBuffer.uidStateChanged(uid, procState, procStateSeq);
+            mUidStateChangeBuffer.uidStateChanged(uid, procState, procStateSeq, capability);
         }
     }
 
@@ -373,7 +376,8 @@ public class NetworkPolicyLogger {
             super(Data.class, capacity);
         }
 
-        public void uidStateChanged(int uid, int procState, long procStateSeq) {
+        public void uidStateChanged(int uid, int procState, long procStateSeq,
+                @ProcessCapability int capability) {
             final Data data = getNextSlot();
             if (data == null) return;
 
@@ -381,6 +385,7 @@ public class NetworkPolicyLogger {
             data.type = EVENT_UID_STATE_CHANGED;
             data.ifield1 = uid;
             data.ifield2 = procState;
+            data.ifield3 = capability;
             data.lfield1 = procStateSeq;
             data.timeStamp = System.currentTimeMillis();
         }
@@ -546,8 +551,9 @@ public class NetworkPolicyLogger {
                 case EVENT_NETWORK_BLOCKED:
                     return data.ifield1 + "-" + getBlockedReason(data.ifield2);
                 case EVENT_UID_STATE_CHANGED:
-                    return data.ifield1 + "-" + ProcessList.makeProcStateString(data.ifield2)
+                    return data.ifield1 + ":" + ProcessList.makeProcStateString(data.ifield2)
-                            + "-" + data.lfield1;
+                            + ":" + ActivityManager.getCapabilitiesSummary(data.ifield3)
+                            + ":" + data.lfield1;
                 case EVENT_POLICIES_CHANGED:
                     return getPolicyChangedLog(data.ifield1, data.ifield2, data.ifield3);
                 case EVENT_METEREDNESS_CHANGED:

services/core/java/com/android/server/net/NetworkPolicyManagerService.java

@@ -130,6 +130,7 @@ import android.annotation.IntDef;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.app.ActivityManager;
+import android.app.ActivityManager.ProcessCapability;
 import android.app.ActivityManagerInternal;
 import android.app.AppGlobals;
 import android.app.AppOpsManager;
@@ -165,6 +166,7 @@ import android.net.NetworkCapabilities;
 import android.net.NetworkIdentity;
 import android.net.NetworkPolicy;
 import android.net.NetworkPolicyManager;
+import android.net.NetworkPolicyManager.UidState;
 import android.net.NetworkQuotaInfo;
 import android.net.NetworkRequest;
 import android.net.NetworkSpecifier;
@@ -559,7 +561,7 @@ public class NetworkPolicyManagerService extends INetworkPolicyManager.Stub {
 
     /** Foreground at UID granularity. */
     @GuardedBy("mUidRulesFirstLock")
-    final SparseIntArray mUidState = new SparseIntArray();
+    final SparseArray<UidState> mUidState = new SparseArray<UidState>();
 
     /** Map from network ID to last observed meteredness state */
     @GuardedBy("mNetworkPoliciesSecondLock")
@@ -990,9 +992,12 @@ public class NetworkPolicyManagerService extends INetworkPolicyManager.Stub {
 
     final private IUidObserver mUidObserver = new IUidObserver.Stub() {
         @Override public void onUidStateChanged(int uid, int procState, long procStateSeq,
-                int capability) {
+                @ProcessCapability int capability) {
+            // TODO: Avoid creating a new UidStateCallbackInfo object every time
+            // we get a callback for an uid
             mUidEventHandler.obtainMessage(UID_MSG_STATE_CHANGED,
-                    uid, procState, procStateSeq).sendToTarget();
+                    new UidStateCallbackInfo(uid, procState, procStateSeq, capability))
+                            .sendToTarget();
         }
 
         @Override public void onUidGone(int uid, boolean disabled) {
@@ -1009,6 +1014,22 @@ public class NetworkPolicyManagerService extends INetworkPolicyManager.Stub {
         }
     };
 
+    private static final class UidStateCallbackInfo {
+        public final int uid;
+        public final int procState;
+        public final long procStateSeq;
+        @ProcessCapability
+        public final int capability;
+
+        UidStateCallbackInfo(int uid, int procState, long procStateSeq,
+                @ProcessCapability int capability) {
+            this.uid = uid;
+            this.procState = procState;
+            this.procStateSeq = procStateSeq;
+            this.capability = capability;
+        }
+    }
+
     final private BroadcastReceiver mPowerSaveWhitelistReceiver = new BroadcastReceiver() {
         @Override
         public void onReceive(Context context, Intent intent) {
@@ -3728,14 +3749,12 @@ public class NetworkPolicyManagerService extends INetworkPolicyManager.Stub {
                     fout.print("UID=");
                     fout.print(uid);
 
-                    final int state = mUidState.get(uid, ActivityManager.PROCESS_STATE_CACHED_EMPTY);
+                    final UidState uidState = mUidState.get(uid);
-                    fout.print(" state=");
+                    if (uidState == null) {
-                    fout.print(state);
+                        fout.print(" state={null}");
-                    if (state <= ActivityManager.PROCESS_STATE_TOP) {
-                        fout.print(" (fg)");
                     } else {
-                        fout.print(state <= ActivityManager.PROCESS_STATE_BOUND_FOREGROUND_SERVICE
+                        fout.print(" state=");
-                                ? " (fg svc)" : " (bg)");
+                        fout.print(uidState.toString());
                     }
 
                     final int uidRules = mUidRules.get(uid, RULE_NONE);
@@ -3792,26 +3811,20 @@ public class NetworkPolicyManagerService extends INetworkPolicyManager.Stub {
     @VisibleForTesting
     boolean isUidForeground(int uid) {
         synchronized (mUidRulesFirstLock) {
-            return isUidStateForeground(
+            return isProcStateAllowedWhileIdleOrPowerSaveMode(mUidState.get(uid));
-                    mUidState.get(uid, ActivityManager.PROCESS_STATE_CACHED_EMPTY));
         }
     }
 
     @GuardedBy("mUidRulesFirstLock")
     private boolean isUidForegroundOnRestrictBackgroundUL(int uid) {
-        final int procState = mUidState.get(uid, ActivityManager.PROCESS_STATE_CACHED_EMPTY);
+        final UidState uidState = mUidState.get(uid);
-        return isProcStateAllowedWhileOnRestrictBackground(procState);
+        return isProcStateAllowedWhileOnRestrictBackground(uidState);
     }
 
     @GuardedBy("mUidRulesFirstLock")
     private boolean isUidForegroundOnRestrictPowerUL(int uid) {
-        final int procState = mUidState.get(uid, ActivityManager.PROCESS_STATE_CACHED_EMPTY);
+        final UidState uidState = mUidState.get(uid);
-        return isProcStateAllowedWhileIdleOrPowerSaveMode(procState);
+        return isProcStateAllowedWhileIdleOrPowerSaveMode(uidState);
-    }
-
-    private boolean isUidStateForeground(int state) {
-        // only really in foreground when screen is also on
-        return state <= NetworkPolicyManager.FOREGROUND_THRESHOLD_STATE;
     }
 
     /**
@@ -3820,16 +3833,18 @@ public class NetworkPolicyManagerService extends INetworkPolicyManager.Stub {
      * {@link #updateRulesForPowerRestrictionsUL(int)}. Returns true if the state was updated.
      */
     @GuardedBy("mUidRulesFirstLock")
-    private boolean updateUidStateUL(int uid, int uidState) {
+    private boolean updateUidStateUL(int uid, int procState, @ProcessCapability int capability) {
         Trace.traceBegin(Trace.TRACE_TAG_NETWORK, "updateUidStateUL");
         try {
-            final int oldUidState = mUidState.get(uid, ActivityManager.PROCESS_STATE_CACHED_EMPTY);
+            final UidState oldUidState = mUidState.get(uid);
-            if (oldUidState != uidState) {
+            if (oldUidState == null || oldUidState.procState != procState
+                    || oldUidState.capability != capability) {
+                final UidState newUidState = new UidState(uid, procState, capability);
                 // state changed, push updated rules
-                mUidState.put(uid, uidState);
+                mUidState.put(uid, newUidState);
-                updateRestrictBackgroundRulesOnUidStatusChangedUL(uid, oldUidState, uidState);
+                updateRestrictBackgroundRulesOnUidStatusChangedUL(uid, oldUidState, newUidState);
                 if (isProcStateAllowedWhileIdleOrPowerSaveMode(oldUidState)
-                        != isProcStateAllowedWhileIdleOrPowerSaveMode(uidState) ) {
+                        != isProcStateAllowedWhileIdleOrPowerSaveMode(newUidState)) {
                     updateRuleForAppIdleUL(uid);
                     if (mDeviceIdleMode) {
                         updateRuleForDeviceIdleUL(uid);
@@ -3851,11 +3866,10 @@ public class NetworkPolicyManagerService extends INetworkPolicyManager.Stub {
     private boolean removeUidStateUL(int uid) {
         final int index = mUidState.indexOfKey(uid);
         if (index >= 0) {
-            final int oldUidState = mUidState.valueAt(index);
+            final UidState oldUidState = mUidState.valueAt(index);
             mUidState.removeAt(index);
-            if (oldUidState != ActivityManager.PROCESS_STATE_CACHED_EMPTY) {
+            if (oldUidState != null) {
-                updateRestrictBackgroundRulesOnUidStatusChangedUL(uid, oldUidState,
+                updateRestrictBackgroundRulesOnUidStatusChangedUL(uid, oldUidState, null);
-                        ActivityManager.PROCESS_STATE_CACHED_EMPTY);
                 if (mDeviceIdleMode) {
                     updateRuleForDeviceIdleUL(uid);
                 }
@@ -3882,8 +3896,8 @@ public class NetworkPolicyManagerService extends INetworkPolicyManager.Stub {
         }
     }
 
-    private void updateRestrictBackgroundRulesOnUidStatusChangedUL(int uid, int oldUidState,
+    private void updateRestrictBackgroundRulesOnUidStatusChangedUL(int uid,
-            int newUidState) {
+            @Nullable UidState oldUidState, @Nullable UidState newUidState) {
         final boolean oldForeground =
                 isProcStateAllowedWhileOnRestrictBackground(oldUidState);
         final boolean newForeground =
@@ -4988,11 +5002,14 @@ public class NetworkPolicyManagerService extends INetworkPolicyManager.Stub {
         public boolean handleMessage(Message msg) {
             switch (msg.what) {
                 case UID_MSG_STATE_CHANGED: {
-                    final int uid = msg.arg1;
+                    final UidStateCallbackInfo uidStateCallbackInfo =
-                    final int procState = msg.arg2;
+                            (UidStateCallbackInfo) msg.obj;
-                    final long procStateSeq = (Long) msg.obj;
+                    final int uid = uidStateCallbackInfo.uid;
+                    final int procState = uidStateCallbackInfo.procState;
+                    final long procStateSeq = uidStateCallbackInfo.procStateSeq;
+                    final int capability = uidStateCallbackInfo.capability;
 
-                    handleUidChanged(uid, procState, procStateSeq);
+                    handleUidChanged(uid, procState, procStateSeq, capability);
                     return true;
                 }
                 case UID_MSG_GONE: {
@@ -5007,23 +5024,24 @@ public class NetworkPolicyManagerService extends INetworkPolicyManager.Stub {
         }
     };
 
-    void handleUidChanged(int uid, int procState, long procStateSeq) {
+    void handleUidChanged(int uid, int procState, long procStateSeq,
+            @ProcessCapability int capability) {
         Trace.traceBegin(Trace.TRACE_TAG_NETWORK, "onUidStateChanged");
         try {
             boolean updated;
             synchronized (mUidRulesFirstLock) {
                 // We received a uid state change callback, add it to the history so that it
                 // will be useful for debugging.
-                mLogger.uidStateChanged(uid, procState, procStateSeq);
+                mLogger.uidStateChanged(uid, procState, procStateSeq, capability);
                 // Now update the network policy rules as per the updated uid state.
-                updated = updateUidStateUL(uid, procState);
+                updated = updateUidStateUL(uid, procState, capability);
                 // Updating the network rules is done, so notify AMS about this.
                 mActivityManagerInternal.notifyNetworkPolicyRulesUpdated(uid, procStateSeq);
             }
             // Do this without the lock held. handleUidChanged() and handleUidGone() are
             // called from the handler, so there's no multi-threading issue.
             if (updated) {
-                updateNetworkStats(uid, isUidStateForeground(procState));
+                updateNetworkStats(uid, isProcStateAllowedWhileOnRestrictBackground(procState));
             }
         } finally {
             Trace.traceEnd(Trace.TRACE_TAG_NETWORK);
@@ -5358,6 +5376,13 @@ public class NetworkPolicyManagerService extends INetworkPolicyManager.Stub {
         }
     }
 
+    private static void collectKeys(SparseArray<UidState> source, SparseBooleanArray target) {
+        final int size = source.size();
+        for (int i = 0; i < size; i++) {
+            target.put(source.keyAt(i), true);
+        }
+    }
+
     @Override
     public void factoryReset(String subscriber) {
         mContext.enforceCallingOrSelfPermission(NETWORK_SETTINGS, TAG);