system_server: add two child chains to firewall

This is an attempt to speed up getting out of device idle.  It groups
uid firewall rules in these child chains so we can attach/detach a whole
chain instead of individual uid rules.

BUG:21446713
Change-Id: Ie8f392da2deabe7cc86a9ecf4ed080163861d41e

5 files changed, 284 insertions, 28 deletions

diff --git a/core/java/android/net/NetworkPolicyManager.java b/core/java/android/net/NetworkPolicyManager.java
index ecc3fb4..3f40484 100644
--- a/core/java/android/net/NetworkPolicyManager.java
+++ b/core/java/android/net/NetworkPolicyManager.java
@@ -61,6 +61,17 @@ public class NetworkPolicyManager {
     public static final int FIREWALL_RULE_ALLOW = 1;
     public static final int FIREWALL_RULE_DENY = 2;
 
+    public static final int FIREWALL_TYPE_WHITELIST = 0;
+    public static final int FIREWALL_TYPE_BLACKLIST = 1;
+
+    public static final int FIREWALL_CHAIN_NONE = 0;
+    public static final int FIREWALL_CHAIN_DOZABLE = 1;
+    public static final int FIREWALL_CHAIN_STANDBY = 2;
+
+    public static final String FIREWALL_CHAIN_NAME_NONE = "none";
+    public static final String FIREWALL_CHAIN_NAME_DOZABLE = "dozable";
+    public static final String FIREWALL_CHAIN_NAME_STANDBY = "standby";
+
     private static final boolean ALLOW_PLATFORM_APP_POLICY = true;
 
     /**
diff --git a/core/java/android/os/INetworkManagementService.aidl b/core/java/android/os/INetworkManagementService.aidl
index b29e8d0..8114155 100644
--- a/core/java/android/os/INetworkManagementService.aidl
+++ b/core/java/android/os/INetworkManagementService.aidl
@@ -342,7 +342,9 @@ interface INetworkManagementService
     void setFirewallInterfaceRule(String iface, boolean allow);
     void setFirewallEgressSourceRule(String addr, boolean allow);
     void setFirewallEgressDestRule(String addr, int port, boolean allow);
-    void setFirewallUidRule(int uid, int rule);
+    void setFirewallUidRule(int chain, int uid, int rule);
+    void setFirewallUidRules(int chain, in int[] uids, in int[] rules);
+    void setFirewallChainEnabled(int chain, boolean enable);
 
     /**
      * Set all packets from users in ranges to go through VPN specified by netId.
diff --git a/services/core/java/com/android/server/NetworkManagementService.java b/services/core/java/com/android/server/NetworkManagementService.java
index a1e2a54..baa55e7 100644
--- a/services/core/java/com/android/server/NetworkManagementService.java
+++ b/services/core/java/com/android/server/NetworkManagementService.java
@@ -19,6 +19,15 @@ package com.android.server;
 import static android.Manifest.permission.CONNECTIVITY_INTERNAL;
 import static android.Manifest.permission.DUMP;
 import static android.Manifest.permission.SHUTDOWN;
+import static android.net.NetworkPolicyManager.FIREWALL_CHAIN_DOZABLE;
+import static android.net.NetworkPolicyManager.FIREWALL_CHAIN_NAME_DOZABLE;
+import static android.net.NetworkPolicyManager.FIREWALL_CHAIN_NAME_NONE;
+import static android.net.NetworkPolicyManager.FIREWALL_CHAIN_NAME_STANDBY;
+import static android.net.NetworkPolicyManager.FIREWALL_CHAIN_NONE;
+import static android.net.NetworkPolicyManager.FIREWALL_CHAIN_STANDBY;
+import static android.net.NetworkPolicyManager.FIREWALL_RULE_DEFAULT;
+import static android.net.NetworkPolicyManager.FIREWALL_TYPE_BLACKLIST;
+import static android.net.NetworkPolicyManager.FIREWALL_TYPE_WHITELIST;
 import static android.net.NetworkStats.SET_DEFAULT;
 import static android.net.NetworkStats.TAG_ALL;
 import static android.net.NetworkStats.TAG_NONE;
@@ -35,6 +44,7 @@ import static com.android.server.NetworkManagementService.NetdResponseCode.Tethe
 import static com.android.server.NetworkManagementService.NetdResponseCode.TtyListResult;
 import static com.android.server.NetworkManagementSocketTagger.PROP_QTAGUID_ENABLED;
 
+import android.annotation.NonNull;
 import android.app.ActivityManagerNative;
 import android.content.Context;
 import android.net.ConnectivityManager;
@@ -192,6 +202,21 @@ public class NetworkManagementService extends INetworkManagementService.Stub
     /** Set of UIDs that are to be blocked/allowed by firewall controller. */
     @GuardedBy("mQuotaLock")
     private SparseIntArray mUidFirewallRules = new SparseIntArray();
+    /**
+     * Set of UIDs that are to be blocked/allowed by firewall controller.  This set of Ids matches
+     * to application idles.
+     */
+    @GuardedBy("mQuotaLock")
+    private SparseIntArray mUidFirewallStandbyRules = new SparseIntArray();
+    /**
+     * Set of UIDs that are to be blocked/allowed by firewall controller.  This set of Ids matches
+     * to device idles.
+     */
+    @GuardedBy("mQuotaLock")
+    private SparseIntArray mUidFirewallDozableRules = new SparseIntArray();
+
+    private boolean mStandbyChainEnabled = false;
+    private boolean mDozableChainEnabled = false;
 
     private Object mIdleTimerLock = new Object();
     /** Set of interfaces with active idle timers. */
@@ -282,6 +307,9 @@ public class NetworkManagementService extends INetworkManagementService.Stub
     }
 
     public void systemReady() {
+        // init firewall states
+        mDozableChainEnabled = false;
+        mStandbyChainEnabled = true;
         prepareNativeDaemon();
         if (DBG) Slog.d(TAG, "Prepared");
     }
@@ -568,9 +596,38 @@ public class NetworkManagementService extends INetworkManagementService.Stub
                 final SparseIntArray uidFirewallRules = mUidFirewallRules;
                 mUidFirewallRules = new SparseIntArray();
                 for (int i = 0; i < uidFirewallRules.size(); i++) {
-                    setFirewallUidRule(uidFirewallRules.keyAt(i), uidFirewallRules.valueAt(i));
+                    setFirewallUidRuleInternal(FIREWALL_CHAIN_NONE, uidFirewallRules.keyAt(i),
+                            uidFirewallRules.valueAt(i));
+                }
+            }
+
+            size = mUidFirewallStandbyRules.size();
+            if (size > 0) {
+                Slog.d(TAG, "Pushing " + size + " active firewall standby UID rules");
+                final SparseIntArray uidFirewallRules = mUidFirewallStandbyRules;
+                mUidFirewallStandbyRules = new SparseIntArray();
+                for (int i = 0; i < uidFirewallRules.size(); i++) {
+                    setFirewallUidRuleInternal(FIREWALL_CHAIN_STANDBY, uidFirewallRules.keyAt(i),
+                            uidFirewallRules.valueAt(i));
+                }
+            }
+            if (mStandbyChainEnabled) {
+                setFirewallChainEnabled(FIREWALL_CHAIN_STANDBY, true);
+            }
+
+            size = mUidFirewallDozableRules.size();
+            if (size > 0) {
+                Slog.d(TAG, "Pushing " + size + " active firewall dozable UID rules");
+                final SparseIntArray uidFirewallRules = mUidFirewallDozableRules;
+                mUidFirewallDozableRules = new SparseIntArray();
+                for (int i = 0; i < uidFirewallRules.size(); i++) {
+                    setFirewallUidRuleInternal(FIREWALL_CHAIN_DOZABLE, uidFirewallRules.keyAt(i),
+                            uidFirewallRules.valueAt(i));
                 }
             }
+            if (mDozableChainEnabled) {
+                setFirewallChainEnabled(FIREWALL_CHAIN_DOZABLE, true);
+            }
         }
     }
 
@@ -1954,13 +2011,78 @@ public class NetworkManagementService extends INetworkManagementService.Stub
     }
 
     @Override
-    public void setFirewallUidRule(int uid, int rule) {
+    public void setFirewallChainEnabled(int chain, boolean enable) {
         enforceSystemUid();
-        if (rule == NetworkPolicyManager.FIREWALL_RULE_ALLOW) {
-            Preconditions.checkState(mFirewallEnabled);
+        final String operation = enable ? "enable_chain" : "disable_chain";
+        try {
+            String chainName;
+            switch(chain) {
+                case FIREWALL_CHAIN_STANDBY:
+                    chainName = FIREWALL_CHAIN_NAME_STANDBY;
+                    mStandbyChainEnabled = enable;
+                    break;
+                case FIREWALL_CHAIN_DOZABLE:
+                    chainName = FIREWALL_CHAIN_NAME_DOZABLE;
+                    mDozableChainEnabled = enable;
+                    break;
+                default:
+                    throw new IllegalArgumentException("Bad child chain: " + chain);
+            }
+            mConnector.execute("firewall", operation, chainName);
+        } catch (NativeDaemonConnectorException e) {
+            throw e.rethrowAsParcelableException();
         }
+    }
+
+    private int getFirewallType(int chain) {
+        switch (chain) {
+            case FIREWALL_CHAIN_STANDBY:
+                return FIREWALL_TYPE_BLACKLIST;
+            case FIREWALL_CHAIN_DOZABLE:
+                return FIREWALL_TYPE_WHITELIST;
+            default:
+                return isFirewallEnabled() ? FIREWALL_TYPE_WHITELIST : FIREWALL_TYPE_BLACKLIST;
+        }
+    }
+
+    @Override
+    public void setFirewallUidRules(int chain, int[] uids, int[] rules) {
+        enforceSystemUid();
+        SparseIntArray uidFirewallRules = getUidFirewallRules(chain);
+        SparseIntArray newRules = new SparseIntArray();
+        // apply new set of rules
+        for (int index = uids.length - 1; index >= 0; --index) {
+            int uid = uids[index];
+            int rule = rules[index];
+            setFirewallUidRule(chain, uid, rule);
+            newRules.put(uid, rule);
+        }
+        // collect the rules to remove.
+        SparseIntArray rulesToRemove = new SparseIntArray();
+        for (int index = uidFirewallRules.size() - 1; index >= 0; --index) {
+            int uid = uidFirewallRules.keyAt(index);
+            if (newRules.indexOfKey(uid) < 0) {
+                rulesToRemove.put(uid, FIREWALL_RULE_DEFAULT);
+            }
+        }
+        // remove dead rules
+        for (int index = rulesToRemove.size() - 1; index >= 0; --index) {
+            int uid = rulesToRemove.keyAt(index);
+            setFirewallUidRuleInternal(chain, uid, FIREWALL_RULE_DEFAULT);
+        }
+    }
+
+    @Override
+    public void setFirewallUidRule(int chain, int uid, int rule) {
+        enforceSystemUid();
+        setFirewallUidRuleInternal(chain, uid, rule);
+    }
+
+    private void setFirewallUidRuleInternal(int chain, int uid, int rule) {
         synchronized (mQuotaLock) {
-            final int oldUidFirewallRule = mUidFirewallRules.get(uid);
+            SparseIntArray uidFirewallRules = getUidFirewallRules(chain);
+
+            final int oldUidFirewallRule = uidFirewallRules.get(uid, FIREWALL_RULE_DEFAULT);
             if (DBG) {
                 Slog.d(TAG, "oldRule = " + oldUidFirewallRule
                         + ", newRule=" + rule + " for uid=" + uid);
@@ -1973,7 +2095,7 @@ public class NetworkManagementService extends INetworkManagementService.Stub
 
             try {
                 String ruleName;
-                if (isFirewallEnabled()) { // Whitelist mode
+                if (getFirewallType(chain) == FIREWALL_TYPE_WHITELIST) {
                     if (rule == NetworkPolicyManager.FIREWALL_RULE_ALLOW) {
                         ruleName = "allow";
                     } else {
@@ -1988,17 +2110,44 @@ public class NetworkManagementService extends INetworkManagementService.Stub
                 }
 
                 if (rule == NetworkPolicyManager.FIREWALL_RULE_DEFAULT) {
-                    mUidFirewallRules.delete(uid);
+                    uidFirewallRules.delete(uid);
                 } else {
-                    mUidFirewallRules.put(uid, rule);
+                    uidFirewallRules.put(uid, rule);
                 }
-                mConnector.execute("firewall", "set_uid_rule", uid, ruleName);
+                mConnector.execute("firewall", "set_uid_rule", getFirewallChainName(chain), uid,
+                        ruleName);
             } catch (NativeDaemonConnectorException e) {
                 throw e.rethrowAsParcelableException();
             }
         }
     }
 
+    private @NonNull SparseIntArray getUidFirewallRules(int chain) {
+        switch (chain) {
+            case FIREWALL_CHAIN_STANDBY:
+                return mUidFirewallStandbyRules;
+            case FIREWALL_CHAIN_DOZABLE:
+                return mUidFirewallDozableRules;
+            case FIREWALL_CHAIN_NONE:
+                return mUidFirewallRules;
+            default:
+                throw new IllegalArgumentException("Unknown chain:" + chain);
+        }
+    }
+
+    public @NonNull String getFirewallChainName(int chain) {
+        switch (chain) {
+            case FIREWALL_CHAIN_STANDBY:
+                return FIREWALL_CHAIN_NAME_STANDBY;
+            case FIREWALL_CHAIN_DOZABLE:
+                return FIREWALL_CHAIN_NAME_DOZABLE;
+            case FIREWALL_CHAIN_NONE:
+                return FIREWALL_CHAIN_NAME_NONE;
+            default:
+                throw new IllegalArgumentException("Unknown chain:" + chain);
+        }
+    }
+
     private static void enforceSystemUid() {
         final int uid = Binder.getCallingUid();
         if (uid != Process.SYSTEM_UID) {
@@ -2123,6 +2272,32 @@ public class NetworkManagementService extends INetworkManagementService.Stub
             pw.println("]");
         }
 
+        pw.println("UID firewall standby chain enabled: " + mStandbyChainEnabled);
+        synchronized (mUidFirewallStandbyRules) {
+            pw.print("UID firewall standby rule: [");
+            final int size = mUidFirewallStandbyRules.size();
+            for (int i = 0; i < size; i++) {
+                pw.print(mUidFirewallStandbyRules.keyAt(i));
+                pw.print(":");
+                pw.print(mUidFirewallStandbyRules.valueAt(i));
+                if (i < size - 1) pw.print(",");
+            }
+            pw.println("]");
+        }
+
+        pw.println("UID firewall dozable chain enabled: " + mDozableChainEnabled);
+        synchronized (mUidFirewallDozableRules) {
+            pw.print("UID firewall dozable rule: [");
+            final int size = mUidFirewallDozableRules.size();
+            for (int i = 0; i < size; i++) {
+                pw.print(mUidFirewallDozableRules.keyAt(i));
+                pw.print(":");
+                pw.print(mUidFirewallDozableRules.valueAt(i));
+                if (i < size - 1) pw.print(",");
+            }
+            pw.println("]");
+        }
+
         synchronized (mIdleTimerLock) {
             pw.println("Idle timers:");
             for (HashMap.Entry<String, IdleTimerParams> ent : mActiveIdleTimers.entrySet()) {
diff --git a/services/core/java/com/android/server/net/LockdownVpnTracker.java b/services/core/java/com/android/server/net/LockdownVpnTracker.java
index 0f88883..9db6a06 100644
--- a/services/core/java/com/android/server/net/LockdownVpnTracker.java
+++ b/services/core/java/com/android/server/net/LockdownVpnTracker.java
@@ -17,6 +17,7 @@
 package com.android.server.net;
 
 import static android.Manifest.permission.CONNECTIVITY_INTERNAL;
+import static android.net.NetworkPolicyManager.FIREWALL_CHAIN_NONE;
 import static android.net.NetworkPolicyManager.FIREWALL_RULE_ALLOW;
 import static android.net.NetworkPolicyManager.FIREWALL_RULE_DEFAULT;
 
@@ -201,8 +202,8 @@ public class LockdownVpnTracker {
                     setFirewallEgressSourceRule(addr, true);
                 }
 
-                mNetService.setFirewallUidRule(ROOT_UID, FIREWALL_RULE_ALLOW);
-                mNetService.setFirewallUidRule(Os.getuid(), FIREWALL_RULE_ALLOW);
+                mNetService.setFirewallUidRule(FIREWALL_CHAIN_NONE, ROOT_UID, FIREWALL_RULE_ALLOW);
+                mNetService.setFirewallUidRule(FIREWALL_CHAIN_NONE, Os.getuid(), FIREWALL_RULE_ALLOW);
 
                 mErrorCount = 0;
                 mAcceptedIface = iface;
@@ -291,8 +292,8 @@ public class LockdownVpnTracker {
                     setFirewallEgressSourceRule(addr, false);
                 }
 
-                mNetService.setFirewallUidRule(ROOT_UID, FIREWALL_RULE_DEFAULT);
-                mNetService.setFirewallUidRule(Os.getuid(), FIREWALL_RULE_DEFAULT);
+                mNetService.setFirewallUidRule(FIREWALL_CHAIN_NONE, ROOT_UID, FIREWALL_RULE_DEFAULT);
+                mNetService.setFirewallUidRule(FIREWALL_CHAIN_NONE,Os.getuid(), FIREWALL_RULE_DEFAULT);
 
                 mAcceptedSourceAddr = null;
             }
diff --git a/services/core/java/com/android/server/net/NetworkPolicyManagerService.java b/services/core/java/com/android/server/net/NetworkPolicyManagerService.java
index 7673af4..3d01a5f 100644
--- a/services/core/java/com/android/server/net/NetworkPolicyManagerService.java
+++ b/services/core/java/com/android/server/net/NetworkPolicyManagerService.java
@@ -36,8 +36,10 @@ import static android.net.NetworkPolicy.LIMIT_DISABLED;
 import static android.net.NetworkPolicy.SNOOZE_NEVER;
 import static android.net.NetworkPolicy.WARNING_DISABLED;
 import static android.net.NetworkPolicyManager.EXTRA_NETWORK_TEMPLATE;
+import static android.net.NetworkPolicyManager.FIREWALL_CHAIN_DOZABLE;
+import static android.net.NetworkPolicyManager.FIREWALL_CHAIN_STANDBY;
+import static android.net.NetworkPolicyManager.FIREWALL_RULE_ALLOW;
 import static android.net.NetworkPolicyManager.FIREWALL_RULE_DENY;
-import static android.net.NetworkPolicyManager.FIREWALL_RULE_DEFAULT;
 import static android.net.NetworkPolicyManager.POLICY_ALLOW_BACKGROUND_BATTERY_SAVE;
 import static android.net.NetworkPolicyManager.POLICY_NONE;
 import static android.net.NetworkPolicyManager.POLICY_REJECT_METERED_BACKGROUND;
@@ -80,7 +82,6 @@ import android.app.AppGlobals;
 import android.app.AppOpsManager;
 import android.app.IActivityManager;
 import android.app.INotificationManager;
-import android.app.IProcessObserver;
 import android.app.IUidObserver;
 import android.app.Notification;
 import android.app.PendingIntent;
@@ -141,7 +142,6 @@ import android.util.Log;
 import android.util.NtpTrustedTime;
 import android.util.Pair;
 import android.util.Slog;
-import android.util.SparseArray;
 import android.util.SparseBooleanArray;
 import android.util.SparseIntArray;
 import android.util.TrustedTime;
@@ -279,6 +279,7 @@ public class NetworkPolicyManagerService extends INetworkPolicyManager.Stub
     final SparseIntArray mUidPolicy = new SparseIntArray();
     /** Currently derived rules for each UID. */
     final SparseIntArray mUidRules = new SparseIntArray();
+    final SparseBooleanArray mFirewallChainStates = new SparseBooleanArray();
 
     /**
      * UIDs that have been white-listed to always be able to have network access
@@ -395,8 +396,6 @@ public class NetworkPolicyManagerService extends INetworkPolicyManager.Stub
 
         mUsageStats = LocalServices.getService(UsageStatsManagerInternal.class);
 
-        final PackageManager pm = mContext.getPackageManager();
-
         synchronized (mRulesLock) {
             updatePowerSaveWhitelistLocked();
             mPowerManagerInternal = LocalServices.getService(PowerManagerInternal.class);
@@ -1085,7 +1084,7 @@ public class NetworkPolicyManagerService extends INetworkPolicyManager.Stub
         // will not have a bandwidth limit.  Also only do this if restrict
         // background data use is *not* enabled, since that takes precendence
         // use over those networks can have a cost associated with it).
-        final boolean powerSave = (mRestrictPower || mDeviceIdleMode) && !mRestrictBackground;
+        final boolean powerSave = mRestrictPower && !mRestrictBackground;
 
         // First, generate identities of all connected networks so we can
         // quickly compare them against all defined policies below.
@@ -2006,6 +2005,29 @@ public class NetworkPolicyManagerService extends INetworkPolicyManager.Stub
         }
     }
 
+    void updateRulesForDeviceIdleLocked() {
+        if (mDeviceIdleMode) {
+            // sync the whitelists before enable dozable chain.  We don't care about the rules if
+            // we are disabling the chain.
+            SparseIntArray uidRules = new SparseIntArray();
+            final List<UserInfo> users = mUserManager.getUsers();
+            for (UserInfo user : users) {
+                for (int i = mPowerSaveTempWhitelistAppIds.size() - 1; i >= 0; i--) {
+                    int appId = mPowerSaveTempWhitelistAppIds.keyAt(i);
+                    int uid = UserHandle.getUid(user.id, appId);
+                    uidRules.put(uid, FIREWALL_RULE_ALLOW);
+                }
+                for (int i = mPowerSaveWhitelistAppIds.size() - 1; i >= 0; i--) {
+                    int appId = mPowerSaveWhitelistAppIds.keyAt(i);
+                    int uid = UserHandle.getUid(user.id, appId);
+                    uidRules.put(uid, FIREWALL_RULE_ALLOW);
+                }
+            }
+            setUidFirewallRules(FIREWALL_CHAIN_DOZABLE, uidRules);
+        }
+        enableFirewallChain(FIREWALL_CHAIN_DOZABLE, mDeviceIdleMode);
+    }
+
     /**
      * Update rules that might be changed by {@link #mRestrictBackground},
      * {@link #mRestrictPower}, or {@link #mDeviceIdleMode} value.
@@ -2016,10 +2038,12 @@ public class NetworkPolicyManagerService extends INetworkPolicyManager.Stub
         // If we are in restrict power mode, we allow all important apps
         // to have data access.  Otherwise, we restrict data access to only
         // the top apps.
-        mCurForegroundState = (!mRestrictBackground && (mRestrictPower || mDeviceIdleMode))
+        mCurForegroundState = (!mRestrictBackground && mRestrictPower)
                 ? ActivityManager.PROCESS_STATE_FOREGROUND_SERVICE
                 : ActivityManager.PROCESS_STATE_TOP;
 
+        updateRulesForDeviceIdleLocked();
+
         // update rules for all installed applications
         final List<UserInfo> users = mUserManager.getUsers();
         final List<ApplicationInfo> apps = pm.getInstalledApplications(
@@ -2113,7 +2137,7 @@ public class NetworkPolicyManagerService extends INetworkPolicyManager.Stub
                 // uid in background, and global background disabled
                 uidRules = RULE_REJECT_METERED;
             }
-        } else if (mRestrictPower || mDeviceIdleMode) {
+        } else if (mRestrictPower) {
             final boolean whitelisted = mPowerSaveWhitelistAppIds.get(appId)
                     || mPowerSaveTempWhitelistAppIds.get(appId);
             if (!whitelisted && !uidForeground
@@ -2144,7 +2168,12 @@ public class NetworkPolicyManagerService extends INetworkPolicyManager.Stub
         final boolean oldFirewallReject = (oldRules & RULE_REJECT_ALL) != 0;
         final boolean firewallReject = (uidRules & RULE_REJECT_ALL) != 0;
         if (oldFirewallReject != firewallReject) {
-            setUidFirewallRules(uid, firewallReject);
+            setUidFirewallRule(FIREWALL_CHAIN_STANDBY, uid, firewallReject);
+            if (mDeviceIdleMode && !firewallReject) {
+                // if we are in device idle mode, and we decide to allow this uid.  we need to punch
+                // a hole in the device idle chain.
+                setUidFirewallRule(FIREWALL_CHAIN_DOZABLE, uid, false);
+            }
         }
 
         // dispatch changed rule to existing listeners
@@ -2296,14 +2325,34 @@ public class NetworkPolicyManagerService extends INetworkPolicyManager.Stub
     }
 
     /**
+     * Set uid rules on a particular firewall chain. This is going to synchronize the rules given
+     * here to netd.  It will clean up dead rules and make sure the target chain only contains rules
+     * specified here.
+     */
+    private void setUidFirewallRules(int chain, SparseIntArray uidRules) {
+        try {
+            int size = uidRules.size();
+            int[] uids = new int[size];
+            int[] rules = new int[size];
+            for(int index = size - 1; index >= 0; --index) {
+                uids[index] = uidRules.keyAt(index);
+                rules[index] = uidRules.valueAt(index);
+            }
+            mNetworkManager.setFirewallUidRules(chain, uids, rules);
+        } catch (IllegalStateException e) {
+            Log.wtf(TAG, "problem setting firewall uid rules", e);
+        } catch (RemoteException e) {
+            // ignored; service lives in system_server
+        }
+    }
+
+    /**
      * Add or remove a uid to the firewall blacklist for all network ifaces.
-     * @param uid
-     * @param rejectOnAll
      */
-    private void setUidFirewallRules(int uid, boolean rejectOnAll) {
+    private void setUidFirewallRule(int chain, int uid, boolean rejectOnAll) {
         try {
-            mNetworkManager.setFirewallUidRule(uid,
-                    rejectOnAll ? FIREWALL_RULE_DENY : FIREWALL_RULE_DEFAULT);
+            mNetworkManager.setFirewallUidRule(chain, uid,
+                    rejectOnAll ? FIREWALL_RULE_DENY : FIREWALL_RULE_ALLOW);
         } catch (IllegalStateException e) {
             Log.wtf(TAG, "problem setting firewall uid rules", e);
         } catch (RemoteException e) {
@@ -2311,6 +2360,24 @@ public class NetworkPolicyManagerService extends INetworkPolicyManager.Stub
         }
     }
 
+    /**
+     * Add or remove a uid to the firewall blacklist for all network ifaces.
+     */
+    private void enableFirewallChain(int chain, boolean enable) {
+        if (mFirewallChainStates.indexOfKey(chain) >= 0 &&
+                mFirewallChainStates.get(chain) == enable) {
+            // All is the same, nothing to do.
+            return;
+        }
+        try {
+            mNetworkManager.setFirewallChainEnabled(chain, enable);
+        } catch (IllegalStateException e) {
+            Log.wtf(TAG, "problem enable firewall chain", e);
+        } catch (RemoteException e) {
+            // ignored; service lives in system_server
+        }
+    }
+
     private long getTotalBytes(NetworkTemplate template, long start, long end) {
         try {
             return mNetworkStats.getNetworkTotalBytes(template, start, end);