Merge changes I69a2970d,Ibfb3ae48,Ie595cf1f am: 9ab777b5d5

Original change: https://android-review.googlesource.com/c/platform/packages/modules/Connectivity/+/2138295

Change-Id: If58e9b76e8b722d471597fb8a7007300fdfe4f35
Signed-off-by: Automerger Merge Worker <android-build-automerger-merge-worker@system.gserviceaccount.com>
diff --git a/service/src/com/android/server/BpfNetMaps.java b/service/src/com/android/server/BpfNetMaps.java
index 5650e31..be0eb5f 100644
--- a/service/src/com/android/server/BpfNetMaps.java
+++ b/service/src/com/android/server/BpfNetMaps.java
@@ -25,6 +25,7 @@
 import static android.net.ConnectivityManager.FIREWALL_CHAIN_RESTRICTED;
 import static android.net.ConnectivityManager.FIREWALL_CHAIN_STANDBY;
 import static android.system.OsConstants.EINVAL;
+import static android.system.OsConstants.ENOENT;
 import static android.system.OsConstants.EOPNOTSUPP;
 
 import android.net.INetd;
@@ -35,6 +36,7 @@
 import android.util.Log;
 import android.util.SparseLongArray;
 
+import com.android.internal.annotations.GuardedBy;
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.modules.utils.build.SdkLevel;
 import com.android.net.module.util.BpfMap;
@@ -68,23 +70,27 @@
 
     private static final String CONFIGURATION_MAP_PATH =
             "/sys/fs/bpf/netd_shared/map_netd_configuration_map";
+    private static final String UID_OWNER_MAP_PATH =
+            "/sys/fs/bpf/netd_shared/map_netd_uid_owner_map";
     private static final U32 UID_RULES_CONFIGURATION_KEY = new U32(0);
     private static BpfMap<U32, U32> sConfigurationMap = null;
+    // BpfMap for UID_OWNER_MAP_PATH. This map is not accessed by others.
+    private static BpfMap<U32, UidOwnerValue> sUidOwnerMap = null;
 
     // LINT.IfChange(match_type)
-    private static final long NO_MATCH = 0;
-    private static final long HAPPY_BOX_MATCH = (1 << 0);
-    private static final long PENALTY_BOX_MATCH = (1 << 1);
-    private static final long DOZABLE_MATCH = (1 << 2);
-    private static final long STANDBY_MATCH = (1 << 3);
-    private static final long POWERSAVE_MATCH = (1 << 4);
-    private static final long RESTRICTED_MATCH = (1 << 5);
-    private static final long LOW_POWER_STANDBY_MATCH = (1 << 6);
-    private static final long IIF_MATCH = (1 << 7);
-    private static final long LOCKDOWN_VPN_MATCH = (1 << 8);
-    private static final long OEM_DENY_1_MATCH = (1 << 9);
-    private static final long OEM_DENY_2_MATCH = (1 << 10);
-    private static final long OEM_DENY_3_MATCH = (1 << 11);
+    @VisibleForTesting public static final long NO_MATCH = 0;
+    @VisibleForTesting public static final long HAPPY_BOX_MATCH = (1 << 0);
+    @VisibleForTesting public static final long PENALTY_BOX_MATCH = (1 << 1);
+    @VisibleForTesting public static final long DOZABLE_MATCH = (1 << 2);
+    @VisibleForTesting public static final long STANDBY_MATCH = (1 << 3);
+    @VisibleForTesting public static final long POWERSAVE_MATCH = (1 << 4);
+    @VisibleForTesting public static final long RESTRICTED_MATCH = (1 << 5);
+    @VisibleForTesting public static final long LOW_POWER_STANDBY_MATCH = (1 << 6);
+    @VisibleForTesting public static final long IIF_MATCH = (1 << 7);
+    @VisibleForTesting public static final long LOCKDOWN_VPN_MATCH = (1 << 8);
+    @VisibleForTesting public static final long OEM_DENY_1_MATCH = (1 << 9);
+    @VisibleForTesting public static final long OEM_DENY_2_MATCH = (1 << 10);
+    @VisibleForTesting public static final long OEM_DENY_3_MATCH = (1 << 11);
     // LINT.ThenChange(packages/modules/Connectivity/bpf_progs/bpf_shared.h)
 
     // TODO: Use Java BpfMap instead of JNI code (TrafficController) for map update.
@@ -111,6 +117,14 @@
         sConfigurationMap = configurationMap;
     }
 
+    /**
+     * Set uidOwnerMap for test.
+     */
+    @VisibleForTesting
+    public static void setUidOwnerMapForTest(BpfMap<U32, UidOwnerValue> uidOwnerMap) {
+        sUidOwnerMap = uidOwnerMap;
+    }
+
     private static BpfMap<U32, U32> getConfigurationMap() {
         try {
             return new BpfMap<>(
@@ -120,10 +134,22 @@
         }
     }
 
+    private static BpfMap<U32, UidOwnerValue> getUidOwnerMap() {
+        try {
+            return new BpfMap<>(
+                    UID_OWNER_MAP_PATH, BpfMap.BPF_F_RDWR, U32.class, UidOwnerValue.class);
+        } catch (ErrnoException e) {
+            throw new IllegalStateException("Cannot open uid owner map", e);
+        }
+    }
+
     private static void setBpfMaps() {
         if (sConfigurationMap == null) {
             sConfigurationMap = getConfigurationMap();
         }
+        if (sUidOwnerMap == null) {
+            sUidOwnerMap = getUidOwnerMap();
+        }
     }
 
     /**
@@ -175,6 +201,67 @@
         }
     }
 
+    private void removeRule(final int uid, final long match, final String caller) {
+        try {
+            synchronized (sUidOwnerMap) {
+                final UidOwnerValue oldMatch = sUidOwnerMap.getValue(new U32(uid));
+
+                if (oldMatch == null) {
+                    throw new ServiceSpecificException(ENOENT,
+                            "sUidOwnerMap does not have entry for uid: " + uid);
+                }
+
+                final UidOwnerValue newMatch = new UidOwnerValue(
+                        (match == IIF_MATCH) ? 0 : oldMatch.iif,
+                        oldMatch.rule & ~match
+                );
+
+                if (newMatch.rule == 0) {
+                    sUidOwnerMap.deleteEntry(new U32(uid));
+                } else {
+                    sUidOwnerMap.updateEntry(new U32(uid), newMatch);
+                }
+            }
+        } catch (ErrnoException e) {
+            throw new ServiceSpecificException(e.errno,
+                    caller + " failed to remove rule: " + Os.strerror(e.errno));
+        }
+    }
+
+    private void addRule(final int uid, final long match, final long iif, final String caller) {
+        if (match != IIF_MATCH && iif != 0) {
+            throw new ServiceSpecificException(EINVAL,
+                    "Non-interface match must have zero interface index");
+        }
+
+        try {
+            synchronized (sUidOwnerMap) {
+                final UidOwnerValue oldMatch = sUidOwnerMap.getValue(new U32(uid));
+
+                final UidOwnerValue newMatch;
+                if (oldMatch != null) {
+                    newMatch = new UidOwnerValue(
+                            (match == IIF_MATCH) ? iif : oldMatch.iif,
+                            oldMatch.rule | match
+                    );
+                } else {
+                    newMatch = new UidOwnerValue(
+                            iif,
+                            match
+                    );
+                }
+                sUidOwnerMap.updateEntry(new U32(uid), newMatch);
+            }
+        } catch (ErrnoException e) {
+            throw new ServiceSpecificException(e.errno,
+                    caller + " failed to add rule: " + Os.strerror(e.errno));
+        }
+    }
+
+    private void addRule(final int uid, final long match, final String caller) {
+        addRule(uid, match, 0 /* iif */, caller);
+    }
+
     /**
      * Add naughty app bandwidth rule for specific app
      *
@@ -183,8 +270,8 @@
      *                                  cause of the failure.
      */
     public void addNaughtyApp(final int uid) {
-        final int err = native_addNaughtyApp(uid);
-        maybeThrow(err, "Unable to add naughty app");
+        throwIfPreT("addNaughtyApp is not available on pre-T devices");
+        addRule(uid, PENALTY_BOX_MATCH, "addNaughtyApp");
     }
 
     /**
@@ -195,8 +282,8 @@
      *                                  cause of the failure.
      */
     public void removeNaughtyApp(final int uid) {
-        final int err = native_removeNaughtyApp(uid);
-        maybeThrow(err, "Unable to remove naughty app");
+        throwIfPreT("removeNaughtyApp is not available on pre-T devices");
+        removeRule(uid, PENALTY_BOX_MATCH, "removeNaughtyApp");
     }
 
     /**
@@ -207,8 +294,10 @@
      *                                  cause of the failure.
      */
     public void addNiceApp(final int uid) {
-        final int err = native_addNiceApp(uid);
-        maybeThrow(err, "Unable to add nice app");
+        synchronized (sUidOwnerMap) {
+            final int err = native_addNiceApp(uid);
+            maybeThrow(err, "Unable to add nice app");
+        }
     }
 
     /**
@@ -219,8 +308,10 @@
      *                                  cause of the failure.
      */
     public void removeNiceApp(final int uid) {
-        final int err = native_removeNiceApp(uid);
-        maybeThrow(err, "Unable to remove nice app");
+        synchronized (sUidOwnerMap) {
+            final int err = native_removeNiceApp(uid);
+            maybeThrow(err, "Unable to remove nice app");
+        }
     }
 
     /**
@@ -285,11 +376,13 @@
      */
     public int replaceUidChain(final String chainName, final boolean isAllowlist,
             final int[] uids) {
-        final int err = native_replaceUidChain(chainName, isAllowlist, uids);
-        if (err != 0) {
-            Log.e(TAG, "replaceUidChain failed: " + Os.strerror(-err));
+        synchronized (sUidOwnerMap) {
+            final int err = native_replaceUidChain(chainName, isAllowlist, uids);
+            if (err != 0) {
+                Log.e(TAG, "replaceUidChain failed: " + Os.strerror(-err));
+            }
+            return -err;
         }
-        return -err;
     }
 
     /**
@@ -302,8 +395,10 @@
      *                                  cause of the failure.
      */
     public void setUidRule(final int childChain, final int uid, final int firewallRule) {
-        final int err = native_setUidRule(childChain, uid, firewallRule);
-        maybeThrow(err, "Unable to set uid rule");
+        synchronized (sUidOwnerMap) {
+            final int err = native_setUidRule(childChain, uid, firewallRule);
+            maybeThrow(err, "Unable to set uid rule");
+        }
     }
 
     /**
@@ -328,8 +423,10 @@
             mNetd.firewallAddUidInterfaceRules(ifName, uids);
             return;
         }
-        final int err = native_addUidInterfaceRules(ifName, uids);
-        maybeThrow(err, "Unable to add uid interface rules");
+        synchronized (sUidOwnerMap) {
+            final int err = native_addUidInterfaceRules(ifName, uids);
+            maybeThrow(err, "Unable to add uid interface rules");
+        }
     }
 
     /**
@@ -348,8 +445,10 @@
             mNetd.firewallRemoveUidInterfaceRules(uids);
             return;
         }
-        final int err = native_removeUidInterfaceRules(uids);
-        maybeThrow(err, "Unable to remove uid interface rules");
+        synchronized (sUidOwnerMap) {
+            final int err = native_removeUidInterfaceRules(uids);
+            maybeThrow(err, "Unable to remove uid interface rules");
+        }
     }
 
     /**
@@ -361,8 +460,10 @@
      *                                  cause of the failure.
      */
     public void updateUidLockdownRule(final int uid, final boolean add) {
-        final int err = native_updateUidLockdownRule(uid, add);
-        maybeThrow(err, "Unable to update lockdown rule");
+        synchronized (sUidOwnerMap) {
+            final int err = native_updateUidLockdownRule(uid, add);
+            maybeThrow(err, "Unable to update lockdown rule");
+        }
     }
 
     /**
@@ -412,14 +513,23 @@
     }
 
     private static native void native_init();
+    @GuardedBy("sUidOwnerMap")
     private native int native_addNaughtyApp(int uid);
+    @GuardedBy("sUidOwnerMap")
     private native int native_removeNaughtyApp(int uid);
+    @GuardedBy("sUidOwnerMap")
     private native int native_addNiceApp(int uid);
+    @GuardedBy("sUidOwnerMap")
     private native int native_removeNiceApp(int uid);
+    @GuardedBy("sUidOwnerMap")
     private native int native_replaceUidChain(String name, boolean isAllowlist, int[] uids);
+    @GuardedBy("sUidOwnerMap")
     private native int native_setUidRule(int childChain, int uid, int firewallRule);
+    @GuardedBy("sUidOwnerMap")
     private native int native_addUidInterfaceRules(String ifName, int[] uids);
+    @GuardedBy("sUidOwnerMap")
     private native int native_removeUidInterfaceRules(int[] uids);
+    @GuardedBy("sUidOwnerMap")
     private native int native_updateUidLockdownRule(int uid, boolean add);
     private native int native_swapActiveStatsMap();
     private native void native_setPermissionForUids(int permissions, int[] uids);
diff --git a/service/src/com/android/server/UidOwnerValue.java b/service/src/com/android/server/UidOwnerValue.java
new file mode 100644
index 0000000..f89e354
--- /dev/null
+++ b/service/src/com/android/server/UidOwnerValue.java
@@ -0,0 +1,35 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server;
+
+import com.android.net.module.util.Struct;
+
+/** Value type for per uid traffic control configuration map  */
+public class UidOwnerValue extends Struct {
+    // Allowed interface index. Only applicable if IIF_MATCH is set in the rule bitmask below.
+    @Field(order = 0, type = Type.U32)
+    public final long iif;
+
+    // A bitmask of match type.
+    @Field(order = 1, type = Type.U32)
+    public final long rule;
+
+    public UidOwnerValue(final long iif, final long rule) {
+        this.iif = iif;
+        this.rule = rule;
+    }
+}
diff --git a/tests/unit/java/com/android/server/BpfNetMapsTest.java b/tests/unit/java/com/android/server/BpfNetMapsTest.java
index 689c148..e59fbc2 100644
--- a/tests/unit/java/com/android/server/BpfNetMapsTest.java
+++ b/tests/unit/java/com/android/server/BpfNetMapsTest.java
@@ -26,8 +26,16 @@
 import static android.net.ConnectivityManager.FIREWALL_CHAIN_STANDBY;
 import static android.net.INetd.PERMISSION_INTERNET;
 
+import static com.android.server.BpfNetMaps.DOZABLE_MATCH;
+import static com.android.server.BpfNetMaps.IIF_MATCH;
+import static com.android.server.BpfNetMaps.NO_MATCH;
+import static com.android.server.BpfNetMaps.PENALTY_BOX_MATCH;
+import static com.android.server.BpfNetMaps.POWERSAVE_MATCH;
+import static com.android.server.BpfNetMaps.RESTRICTED_MATCH;
+
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assume.assumeFalse;
@@ -68,7 +76,9 @@
 
     private static final int TEST_UID = 10086;
     private static final int[] TEST_UIDS = {10002, 10003};
-    private static final String IFNAME = "wlan0";
+    private static final String TEST_IF_NAME = "wlan0";
+    private static final int TEST_IF_INDEX = 7;
+    private static final int NO_IIF = 0;
     private static final String CHAINNAME = "fw_dozable";
     private static final U32 UID_RULES_CONFIGURATION_KEY = new U32(0);
     private static final List<Integer> FIREWALL_CHAINS = List.of(
@@ -86,19 +96,22 @@
 
     @Mock INetd mNetd;
     private final BpfMap<U32, U32> mConfigurationMap = new TestBpfMap<>(U32.class, U32.class);
+    private final BpfMap<U32, UidOwnerValue> mUidOwnerMap =
+            new TestBpfMap<>(U32.class, UidOwnerValue.class);
 
     @Before
     public void setUp() throws Exception {
         MockitoAnnotations.initMocks(this);
         BpfNetMaps.setConfigurationMapForTest(mConfigurationMap);
+        BpfNetMaps.setUidOwnerMapForTest(mUidOwnerMap);
         mBpfNetMaps = new BpfNetMaps(mNetd);
     }
 
     @Test
     public void testBpfNetMapsBeforeT() throws Exception {
         assumeFalse(SdkLevel.isAtLeastT());
-        mBpfNetMaps.addUidInterfaceRules(IFNAME, TEST_UIDS);
-        verify(mNetd).firewallAddUidInterfaceRules(IFNAME, TEST_UIDS);
+        mBpfNetMaps.addUidInterfaceRules(TEST_IF_NAME, TEST_UIDS);
+        verify(mNetd).firewallAddUidInterfaceRules(TEST_IF_NAME, TEST_UIDS);
         mBpfNetMaps.removeUidInterfaceRules(TEST_UIDS);
         verify(mNetd).firewallRemoveUidInterfaceRules(TEST_UIDS);
         mBpfNetMaps.setNetPermForUids(PERMISSION_INTERNET, TEST_UIDS);
@@ -238,4 +251,85 @@
         assertThrows(UnsupportedOperationException.class,
                 () -> mBpfNetMaps.setChildChain(FIREWALL_CHAIN_DOZABLE, true /* enable */));
     }
+
+    private void checkUidOwnerValue(final long uid, final long expectedIif,
+            final long expectedMatch) throws Exception {
+        final UidOwnerValue config = mUidOwnerMap.getValue(new U32(uid));
+        if (expectedMatch == 0) {
+            assertNull(config);
+        } else {
+            assertEquals(expectedIif, config.iif);
+            assertEquals(expectedMatch, config.rule);
+        }
+    }
+
+    private void doTestRemoveNaughtyApp(final long iif, final long match) throws Exception {
+        mUidOwnerMap.updateEntry(new U32(TEST_UID), new UidOwnerValue(iif, match));
+
+        mBpfNetMaps.removeNaughtyApp(TEST_UID);
+
+        checkUidOwnerValue(TEST_UID, iif, match & ~PENALTY_BOX_MATCH);
+    }
+
+    @Test
+    @IgnoreUpTo(Build.VERSION_CODES.S_V2)
+    public void testRemoveNaughtyApp() throws Exception {
+        doTestRemoveNaughtyApp(NO_IIF, PENALTY_BOX_MATCH);
+
+        // PENALTY_BOX_MATCH with other matches
+        doTestRemoveNaughtyApp(NO_IIF, PENALTY_BOX_MATCH | DOZABLE_MATCH | POWERSAVE_MATCH);
+
+        // PENALTY_BOX_MATCH with IIF_MATCH
+        doTestRemoveNaughtyApp(TEST_IF_INDEX, PENALTY_BOX_MATCH | IIF_MATCH);
+
+        // PENALTY_BOX_MATCH is not enabled
+        doTestRemoveNaughtyApp(NO_IIF, DOZABLE_MATCH | POWERSAVE_MATCH | RESTRICTED_MATCH);
+    }
+
+    @Test
+    @IgnoreUpTo(Build.VERSION_CODES.S_V2)
+    public void testRemoveNaughtyAppMissingUid() {
+        // UidOwnerMap does not have entry for TEST_UID
+        assertThrows(ServiceSpecificException.class,
+                () -> mBpfNetMaps.removeNaughtyApp(TEST_UID));
+    }
+
+    @Test
+    @IgnoreAfter(Build.VERSION_CODES.S_V2)
+    public void testRemoveNaughtyAppBeforeT() {
+        assertThrows(UnsupportedOperationException.class,
+                () -> mBpfNetMaps.removeNaughtyApp(TEST_UID));
+    }
+
+    private void doTestAddNaughtyApp(final long iif, final long match) throws Exception {
+        if (match != NO_MATCH) {
+            mUidOwnerMap.updateEntry(new U32(TEST_UID), new UidOwnerValue(iif, match));
+        }
+
+        mBpfNetMaps.addNaughtyApp(TEST_UID);
+
+        checkUidOwnerValue(TEST_UID, iif, match | PENALTY_BOX_MATCH);
+    }
+
+    @Test
+    @IgnoreUpTo(Build.VERSION_CODES.S_V2)
+    public void testAddNaughtyApp() throws Exception {
+        doTestAddNaughtyApp(NO_IIF, NO_MATCH);
+
+        // Other matches are enabled
+        doTestAddNaughtyApp(NO_IIF, DOZABLE_MATCH | POWERSAVE_MATCH | RESTRICTED_MATCH);
+
+        // IIF_MATCH is enabled
+        doTestAddNaughtyApp(TEST_IF_INDEX, IIF_MATCH);
+
+        // PENALTY_BOX_MATCH is already enabled
+        doTestAddNaughtyApp(NO_IIF, PENALTY_BOX_MATCH | DOZABLE_MATCH);
+    }
+
+    @Test
+    @IgnoreAfter(Build.VERSION_CODES.S_V2)
+    public void testAddNaughtyAppBeforeT() {
+        assertThrows(UnsupportedOperationException.class,
+                () -> mBpfNetMaps.addNaughtyApp(TEST_UID));
+    }
 }