diff --git a/service/src/com/android/server/BpfNetMaps.java b/service/src/com/android/server/BpfNetMaps.java
index bedd0e8..594223c 100644
--- a/service/src/com/android/server/BpfNetMaps.java
+++ b/service/src/com/android/server/BpfNetMaps.java
@@ -40,7 +40,6 @@
 import android.system.Os;
 import android.util.Log;
 
-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;
@@ -365,7 +364,13 @@
      */
     public void addNaughtyApp(final int uid) {
         throwIfPreT("addNaughtyApp is not available on pre-T devices");
-        addRule(uid, PENALTY_BOX_MATCH, "addNaughtyApp");
+
+        if (sEnableJavaBpfMap) {
+            addRule(uid, PENALTY_BOX_MATCH, "addNaughtyApp");
+        } else {
+            final int err = native_addNaughtyApp(uid);
+            maybeThrow(err, "Unable to add naughty app");
+        }
     }
 
     /**
@@ -377,7 +382,13 @@
      */
     public void removeNaughtyApp(final int uid) {
         throwIfPreT("removeNaughtyApp is not available on pre-T devices");
-        removeRule(uid, PENALTY_BOX_MATCH, "removeNaughtyApp");
+
+        if (sEnableJavaBpfMap) {
+            removeRule(uid, PENALTY_BOX_MATCH, "removeNaughtyApp");
+        } else {
+            final int err = native_removeNaughtyApp(uid);
+            maybeThrow(err, "Unable to remove naughty app");
+        }
     }
 
     /**
@@ -389,7 +400,13 @@
      */
     public void addNiceApp(final int uid) {
         throwIfPreT("addNiceApp is not available on pre-T devices");
-        addRule(uid, HAPPY_BOX_MATCH, "addNiceApp");
+
+        if (sEnableJavaBpfMap) {
+            addRule(uid, HAPPY_BOX_MATCH, "addNiceApp");
+        } else {
+            final int err = native_addNiceApp(uid);
+            maybeThrow(err, "Unable to add nice app");
+        }
     }
 
     /**
@@ -401,7 +418,13 @@
      */
     public void removeNiceApp(final int uid) {
         throwIfPreT("removeNiceApp is not available on pre-T devices");
-        removeRule(uid, HAPPY_BOX_MATCH, "removeNiceApp");
+
+        if (sEnableJavaBpfMap) {
+            removeRule(uid, HAPPY_BOX_MATCH, "removeNiceApp");
+        } else {
+            final int err = native_removeNiceApp(uid);
+            maybeThrow(err, "Unable to remove nice app");
+        }
     }
 
     /**
@@ -416,16 +439,21 @@
     public void setChildChain(final int childChain, final boolean enable) {
         throwIfPreT("setChildChain is not available on pre-T devices");
 
-        final long match = getMatchByFirewallChain(childChain);
-        try {
-            synchronized (sUidRulesConfigBpfMapLock) {
-                final U32 config = sConfigurationMap.getValue(UID_RULES_CONFIGURATION_KEY);
-                final long newConfig = enable ? (config.val | match) : (config.val & ~match);
-                sConfigurationMap.updateEntry(UID_RULES_CONFIGURATION_KEY, new U32(newConfig));
+        if (sEnableJavaBpfMap) {
+            final long match = getMatchByFirewallChain(childChain);
+            try {
+                synchronized (sUidRulesConfigBpfMapLock) {
+                    final U32 config = sConfigurationMap.getValue(UID_RULES_CONFIGURATION_KEY);
+                    final long newConfig = enable ? (config.val | match) : (config.val & ~match);
+                    sConfigurationMap.updateEntry(UID_RULES_CONFIGURATION_KEY, new U32(newConfig));
+                }
+            } catch (ErrnoException e) {
+                throw new ServiceSpecificException(e.errno,
+                        "Unable to set child chain: " + Os.strerror(e.errno));
             }
-        } catch (ErrnoException e) {
-            throw new ServiceSpecificException(e.errno,
-                    "Unable to set child chain: " + Os.strerror(e.errno));
+        } else {
+            final int err = native_setChildChain(childChain, enable);
+            maybeThrow(err, "Unable to set child chain");
         }
     }
 
@@ -463,36 +491,73 @@
     public void replaceUidChain(final int chain, final int[] uids) {
         throwIfPreT("replaceUidChain is not available on pre-T devices");
 
-        final long match;
-        try {
-            match = getMatchByFirewallChain(chain);
-        } catch (ServiceSpecificException e) {
-            // Throws IllegalArgumentException to keep the behavior of
-            // ConnectivityManager#replaceFirewallChain API
-            throw new IllegalArgumentException("Invalid firewall chain: " + chain);
-        }
-        final Set<Integer> uidSet = Arrays.stream(uids).boxed().collect(Collectors.toSet());
-        final Set<Integer> uidSetToRemoveRule = new HashSet<>();
-        try {
-            synchronized (sUidOwnerMap) {
-                sUidOwnerMap.forEach((uid, config) -> {
-                    // config could be null if there is a concurrent entry deletion.
-                    // http://b/220084230.
-                    if (config != null
-                            && !uidSet.contains((int) uid.val) && (config.rule & match) != 0) {
-                        uidSetToRemoveRule.add((int) uid.val);
-                    }
-                });
-
-                for (final int uid : uidSetToRemoveRule) {
-                    removeRule(uid, match, "replaceUidChain");
-                }
-                for (final int uid : uids) {
-                    addRule(uid, match, "replaceUidChain");
-                }
+        if (sEnableJavaBpfMap) {
+            final long match;
+            try {
+                match = getMatchByFirewallChain(chain);
+            } catch (ServiceSpecificException e) {
+                // Throws IllegalArgumentException to keep the behavior of
+                // ConnectivityManager#replaceFirewallChain API
+                throw new IllegalArgumentException("Invalid firewall chain: " + chain);
             }
-        } catch (ErrnoException | ServiceSpecificException e) {
-            Log.e(TAG, "replaceUidChain failed: " + e);
+            final Set<Integer> uidSet = Arrays.stream(uids).boxed().collect(Collectors.toSet());
+            final Set<Integer> uidSetToRemoveRule = new HashSet<>();
+            try {
+                synchronized (sUidOwnerMap) {
+                    sUidOwnerMap.forEach((uid, config) -> {
+                        // config could be null if there is a concurrent entry deletion.
+                        // http://b/220084230.
+                        if (config != null
+                                && !uidSet.contains((int) uid.val) && (config.rule & match) != 0) {
+                            uidSetToRemoveRule.add((int) uid.val);
+                        }
+                    });
+
+                    for (final int uid : uidSetToRemoveRule) {
+                        removeRule(uid, match, "replaceUidChain");
+                    }
+                    for (final int uid : uids) {
+                        addRule(uid, match, "replaceUidChain");
+                    }
+                }
+            } catch (ErrnoException | ServiceSpecificException e) {
+                Log.e(TAG, "replaceUidChain failed: " + e);
+            }
+        } else {
+            final int err;
+            switch (chain) {
+                case FIREWALL_CHAIN_DOZABLE:
+                    err = native_replaceUidChain("fw_dozable", true /* isAllowList */, uids);
+                    break;
+                case FIREWALL_CHAIN_STANDBY:
+                    err = native_replaceUidChain("fw_standby", false /* isAllowList */, uids);
+                    break;
+                case FIREWALL_CHAIN_POWERSAVE:
+                    err = native_replaceUidChain("fw_powersave", true /* isAllowList */, uids);
+                    break;
+                case FIREWALL_CHAIN_RESTRICTED:
+                    err = native_replaceUidChain("fw_restricted", true /* isAllowList */, uids);
+                    break;
+                case FIREWALL_CHAIN_LOW_POWER_STANDBY:
+                    err = native_replaceUidChain(
+                            "fw_low_power_standby", true /* isAllowList */, uids);
+                    break;
+                case FIREWALL_CHAIN_OEM_DENY_1:
+                    err = native_replaceUidChain("fw_oem_deny_1", false /* isAllowList */, uids);
+                    break;
+                case FIREWALL_CHAIN_OEM_DENY_2:
+                    err = native_replaceUidChain("fw_oem_deny_2", false /* isAllowList */, uids);
+                    break;
+                case FIREWALL_CHAIN_OEM_DENY_3:
+                    err = native_replaceUidChain("fw_oem_deny_3", false /* isAllowList */, uids);
+                    break;
+                default:
+                    throw new IllegalArgumentException("replaceFirewallChain with invalid chain: "
+                            + chain);
+            }
+            if (err != 0) {
+                Log.e(TAG, "replaceUidChain failed: " + Os.strerror(-err));
+            }
         }
     }
 
@@ -508,15 +573,20 @@
     public void setUidRule(final int childChain, final int uid, final int firewallRule) {
         throwIfPreT("setUidRule is not available on pre-T devices");
 
-        final long match = getMatchByFirewallChain(childChain);
-        final boolean isAllowList = isFirewallAllowList(childChain);
-        final boolean add = (firewallRule == FIREWALL_RULE_ALLOW && isAllowList)
-                || (firewallRule == FIREWALL_RULE_DENY && !isAllowList);
+        if (sEnableJavaBpfMap) {
+            final long match = getMatchByFirewallChain(childChain);
+            final boolean isAllowList = isFirewallAllowList(childChain);
+            final boolean add = (firewallRule == FIREWALL_RULE_ALLOW && isAllowList)
+                    || (firewallRule == FIREWALL_RULE_DENY && !isAllowList);
 
-        if (add) {
-            addRule(uid, match, "setUidRule");
+            if (add) {
+                addRule(uid, match, "setUidRule");
+            } else {
+                removeRule(uid, match, "setUidRule");
+            }
         } else {
-            removeRule(uid, match, "setUidRule");
+            final int err = native_setUidRule(childChain, uid, firewallRule);
+            maybeThrow(err, "Unable to set uid rule");
         }
     }
 
@@ -542,24 +612,30 @@
             mNetd.firewallAddUidInterfaceRules(ifName, uids);
             return;
         }
-        // Null ifName is a wildcard to allow apps to receive packets on all interfaces and ifIndex
-        // is set to 0.
-        final int ifIndex;
-        if (ifName == null) {
-            ifIndex = 0;
+
+        if (sEnableJavaBpfMap) {
+            // Null ifName is a wildcard to allow apps to receive packets on all interfaces and
+            // ifIndex is set to 0.
+            final int ifIndex;
+            if (ifName == null) {
+                ifIndex = 0;
+            } else {
+                ifIndex = mDeps.getIfIndex(ifName);
+                if (ifIndex == 0) {
+                    throw new ServiceSpecificException(ENODEV,
+                            "Failed to get index of interface " + ifName);
+                }
+            }
+            for (final int uid : uids) {
+                try {
+                    addRule(uid, IIF_MATCH, ifIndex, "addUidInterfaceRules");
+                } catch (ServiceSpecificException e) {
+                    Log.e(TAG, "addRule failed uid=" + uid + " ifName=" + ifName + ", " + e);
+                }
+            }
         } else {
-            ifIndex = mDeps.getIfIndex(ifName);
-            if (ifIndex == 0) {
-                throw new ServiceSpecificException(ENODEV,
-                        "Failed to get index of interface " + ifName);
-            }
-        }
-        for (final int uid: uids) {
-            try {
-                addRule(uid, IIF_MATCH, ifIndex, "addUidInterfaceRules");
-            } catch (ServiceSpecificException e) {
-                Log.e(TAG, "addRule failed uid=" + uid + " ifName=" + ifName + ", " + e);
-            }
+            final int err = native_addUidInterfaceRules(ifName, uids);
+            maybeThrow(err, "Unable to add uid interface rules");
         }
     }
 
@@ -579,12 +655,18 @@
             mNetd.firewallRemoveUidInterfaceRules(uids);
             return;
         }
-        for (final int uid: uids) {
-            try {
-                removeRule(uid, IIF_MATCH, "removeUidInterfaceRules");
-            } catch (ServiceSpecificException e) {
-                Log.e(TAG, "removeRule failed uid=" + uid + ", " + e);
+
+        if (sEnableJavaBpfMap) {
+            for (final int uid : uids) {
+                try {
+                    removeRule(uid, IIF_MATCH, "removeUidInterfaceRules");
+                } catch (ServiceSpecificException e) {
+                    Log.e(TAG, "removeRule failed uid=" + uid + ", " + e);
+                }
             }
+        } else {
+            final int err = native_removeUidInterfaceRules(uids);
+            maybeThrow(err, "Unable to remove uid interface rules");
         }
     }
 
@@ -598,10 +680,16 @@
      */
     public void updateUidLockdownRule(final int uid, final boolean add) {
         throwIfPreT("updateUidLockdownRule is not available on pre-T devices");
-        if (add) {
-            addRule(uid, LOCKDOWN_VPN_MATCH, "updateUidLockdownRule");
+
+        if (sEnableJavaBpfMap) {
+            if (add) {
+                addRule(uid, LOCKDOWN_VPN_MATCH, "updateUidLockdownRule");
+            } else {
+                removeRule(uid, LOCKDOWN_VPN_MATCH, "updateUidLockdownRule");
+            }
         } else {
-            removeRule(uid, LOCKDOWN_VPN_MATCH, "updateUidLockdownRule");
+            final int err = native_updateUidLockdownRule(uid, add);
+            maybeThrow(err, "Unable to update lockdown rule");
         }
     }
 
@@ -652,24 +740,15 @@
     }
 
     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);
     private native int native_setChildChain(int childChain, boolean enable);
-    @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/ConnectivityService.java b/service/src/com/android/server/ConnectivityService.java
index 7b5026d..960059b 100644
--- a/service/src/com/android/server/ConnectivityService.java
+++ b/service/src/com/android/server/ConnectivityService.java
@@ -3604,10 +3604,9 @@
 
             switch (msg.what) {
                 case NetworkAgent.EVENT_NETWORK_CAPABILITIES_CHANGED: {
-                    nai.declaredCapabilitiesUnsanitized =
-                            new NetworkCapabilities((NetworkCapabilities) arg.second);
-                    final NetworkCapabilities sanitized = sanitizedCapabilitiesFromAgent(
-                            mCarrierPrivilegeAuthenticator, nai);
+                    nai.setDeclaredCapabilities((NetworkCapabilities) arg.second);
+                    final NetworkCapabilities sanitized =
+                            nai.getDeclaredCapabilitiesSanitized(mCarrierPrivilegeAuthenticator);
                     maybeUpdateWifiRoamTimestamp(nai, sanitized);
                     updateCapabilities(nai.getCurrentScore(), nai, sanitized);
                     break;
@@ -7328,11 +7327,11 @@
         // while the network monitor is starting.
         final LinkProperties lp = new LinkProperties(nai.linkProperties);
         // Store a copy of the declared capabilities.
-        nai.declaredCapabilitiesUnsanitized = new NetworkCapabilities(nai.networkCapabilities);
+        nai.setDeclaredCapabilities(nai.networkCapabilities);
         // Make sure the LinkProperties and NetworkCapabilities reflect what the agent info said.
-        final NetworkCapabilities nc =
-                sanitizedCapabilitiesFromAgent(mCarrierPrivilegeAuthenticator, nai);
-        nai.getAndSetNetworkCapabilities(mixInCapabilities(nai, nc));
+        final NetworkCapabilities sanitized =
+                nai.getDeclaredCapabilitiesSanitized(mCarrierPrivilegeAuthenticator);
+        nai.getAndSetNetworkCapabilities(mixInCapabilities(nai, sanitized));
         processLinkPropertiesFromAgent(nai, lp);
         nai.linkProperties = lp;
 
@@ -7795,38 +7794,6 @@
         }
     }
 
-    /**
-     * Sanitize capabilities coming from a network agent.
-     *
-     * Agents have restrictions on what capabilities they can send to Connectivity. For example,
-     * they can't change the owner UID from what they declared before, and complex restrictions
-     * apply to the accessUids field.
-     * They also should not mutate immutable capabilities, although for backward-compatibility
-     * this is not enforced and limited to just a log.
-     *
-     * This method returns a sanitized copy of the passed capabilities to make sure they don't
-     * contain stuff they should not, and should generally be called by code that accesses
-     * {@link NetworkAgentInfo#declaredCapabilitiesUnsanitized}.
-     */
-    // TODO : move this to NetworkAgentInfo
-    private NetworkCapabilities sanitizedCapabilitiesFromAgent(
-            final CarrierPrivilegeAuthenticator carrierPrivilegeAuthenticator,
-            @NonNull final NetworkAgentInfo nai) {
-        final NetworkCapabilities nc = new NetworkCapabilities(nai.declaredCapabilitiesUnsanitized);
-        if (nc.hasConnectivityManagedCapability()) {
-            Log.wtf(TAG, "BUG: " + nai + " has CS-managed capability.");
-        }
-        if (nai.networkCapabilities.getOwnerUid() != nc.getOwnerUid()) {
-            Log.e(TAG, nai.toShortString() + ": ignoring attempt to change owner from "
-                    + nai.networkCapabilities.getOwnerUid() + " to " + nc.getOwnerUid());
-            nc.setOwnerUid(nai.networkCapabilities.getOwnerUid());
-        }
-        NetworkAgentInfo.restrictCapabilitiesFromNetworkAgent(nc, nai.creatorUid,
-                mContext.getPackageManager().hasSystemFeature(PackageManager.FEATURE_AUTOMOTIVE),
-                carrierPrivilegeAuthenticator);
-        return nc;
-    }
-
     /** Modifies |newNc| based on the capabilities of |underlyingNetworks| and |agentCaps|. */
     @VisibleForTesting
     void applyUnderlyingCapabilities(@Nullable Network[] underlyingNetworks,
@@ -7952,7 +7919,7 @@
 
         if (nai.propagateUnderlyingCapabilities()) {
             applyUnderlyingCapabilities(nai.declaredUnderlyingNetworks,
-                    sanitizedCapabilitiesFromAgent(mCarrierPrivilegeAuthenticator, nai),
+                    nai.getDeclaredCapabilitiesSanitized(mCarrierPrivilegeAuthenticator),
                     newNc);
         }
 
diff --git a/service/src/com/android/server/connectivity/NetworkAgentInfo.java b/service/src/com/android/server/connectivity/NetworkAgentInfo.java
index 04031af..c863165 100644
--- a/service/src/com/android/server/connectivity/NetworkAgentInfo.java
+++ b/service/src/com/android/server/connectivity/NetworkAgentInfo.java
@@ -26,6 +26,7 @@
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.content.Context;
+import android.content.pm.PackageManager;
 import android.net.CaptivePortalData;
 import android.net.DscpPolicy;
 import android.net.IDnsResolver;
@@ -184,9 +185,8 @@
     //
     // As the name implies, these capabilities are not sanitized and are not to
     // be trusted. Most callers should simply use the {@link networkCapabilities}
-    // field instead, and callers who need the declared capabilities should generally
-    // pass these to {@link ConnectivityService#sanitizedCapabilitiesFromAgent} before using them.
-    public @Nullable NetworkCapabilities declaredCapabilitiesUnsanitized;
+    // field instead.
+    private @Nullable NetworkCapabilities mDeclaredCapabilitiesUnsanitized;
 
     // Indicates if netd has been told to create this Network. From this point on the appropriate
     // routing rules are setup and routes are added so packets can begin flowing over the Network.
@@ -240,6 +240,53 @@
     // URL, Terms & Conditions URL, and network friendly name.
     public CaptivePortalData networkAgentPortalData;
 
+    /**
+     * Sets the capabilities sent by the agent for later retrieval.
+     *
+     * This method does not sanitize the capabilities ; instead, use
+     * {@link #getDeclaredCapabilitiesSanitized} to retrieve a sanitized
+     * copy of the capabilities as they were passed here.
+     *
+     * This method makes a defensive copy to avoid issues where the passed object is later mutated.
+     *
+     * @param caps the caps sent by the agent
+     */
+    public void setDeclaredCapabilities(@NonNull final NetworkCapabilities caps) {
+        mDeclaredCapabilitiesUnsanitized = new NetworkCapabilities(caps);
+    }
+
+    /**
+     * Get the latest capabilities sent by the network agent, after sanitizing them.
+     *
+     * These are the capabilities as they were sent by the agent (but sanitized to conform to
+     * their restrictions). They are NOT the capabilities currently applying to this agent ;
+     * for that, use {@link #networkCapabilities}.
+     *
+     * Agents have restrictions on what capabilities they can send to Connectivity. For example,
+     * they can't change the owner UID from what they declared before, and complex restrictions
+     * apply to the allowedUids field.
+     * They also should not mutate immutable capabilities, although for backward-compatibility
+     * this is not enforced and limited to just a log.
+     *
+     * @param carrierPrivilegeAuthenticator the authenticator, to check access UIDs.
+     */
+    public NetworkCapabilities getDeclaredCapabilitiesSanitized(
+            final CarrierPrivilegeAuthenticator carrierPrivilegeAuthenticator) {
+        final NetworkCapabilities nc = new NetworkCapabilities(mDeclaredCapabilitiesUnsanitized);
+        if (nc.hasConnectivityManagedCapability()) {
+            Log.wtf(TAG, "BUG: " + this + " has CS-managed capability.");
+        }
+        if (networkCapabilities.getOwnerUid() != nc.getOwnerUid()) {
+            Log.e(TAG, toShortString() + ": ignoring attempt to change owner from "
+                    + networkCapabilities.getOwnerUid() + " to " + nc.getOwnerUid());
+            nc.setOwnerUid(networkCapabilities.getOwnerUid());
+        }
+        restrictCapabilitiesFromNetworkAgent(nc, creatorUid,
+                mContext.getPackageManager().hasSystemFeature(PackageManager.FEATURE_AUTOMOTIVE),
+                carrierPrivilegeAuthenticator);
+        return nc;
+    }
+
     // Networks are lingered when they become unneeded as a result of their NetworkRequests being
     // satisfied by a higher-scoring network. so as to allow communication to wrap up before the
     // network is taken down.  This usually only happens to the default network. Lingering ends with
diff --git a/tests/unit/java/com/android/server/ConnectivityServiceTest.java b/tests/unit/java/com/android/server/ConnectivityServiceTest.java
index 1a624e3..06f7300 100644
--- a/tests/unit/java/com/android/server/ConnectivityServiceTest.java
+++ b/tests/unit/java/com/android/server/ConnectivityServiceTest.java
@@ -241,6 +241,7 @@
 import android.net.ConnectivityThread;
 import android.net.DataStallReportParcelable;
 import android.net.EthernetManager;
+import android.net.EthernetNetworkSpecifier;
 import android.net.IConnectivityDiagnosticsCallback;
 import android.net.IDnsResolver;
 import android.net.INetd;
@@ -8531,8 +8532,7 @@
 
         // Enable always-on VPN lockdown. The main user loses network access because no VPN is up.
         final ArrayList<String> allowList = new ArrayList<>();
-        mVpnManagerService.setAlwaysOnVpnPackage(PRIMARY_USER, ALWAYS_ON_PACKAGE,
-                true /* lockdown */, allowList);
+        mMockVpn.setAlwaysOnPackage(ALWAYS_ON_PACKAGE, true /* lockdown */, allowList);
         waitForIdle();
         assertNull(mCm.getActiveNetworkForUid(uid));
         // This is arguably overspecified: a UID that is not running doesn't have an active network.
@@ -8564,8 +8564,7 @@
         assertNull(mCm.getActiveNetworkForUid(uid));
         assertNotNull(mCm.getActiveNetworkForUid(restrictedUid));
 
-        mVpnManagerService.setAlwaysOnVpnPackage(PRIMARY_USER, null, false /* lockdown */,
-                allowList);
+        mMockVpn.setAlwaysOnPackage(null, false /* lockdown */, allowList);
         waitForIdle();
     }
 
@@ -9023,10 +9022,8 @@
                 new Handler(ConnectivityThread.getInstanceLooper()));
 
         final int uid = Process.myUid();
-        final int userId = UserHandle.getUserId(uid);
         final ArrayList<String> allowList = new ArrayList<>();
-        mVpnManagerService.setAlwaysOnVpnPackage(userId, ALWAYS_ON_PACKAGE, true /* lockdown */,
-                allowList);
+        mMockVpn.setAlwaysOnPackage(ALWAYS_ON_PACKAGE, true /* lockdown */, allowList);
         waitForIdle();
 
         final Set<Integer> excludedUids = new ArraySet<Integer>();
@@ -9056,7 +9053,7 @@
         assertNetworkInfo(TYPE_WIFI, DetailedState.BLOCKED);
 
         // Disable lockdown, expect to see the network unblocked.
-        mVpnManagerService.setAlwaysOnVpnPackage(userId, null, false /* lockdown */, allowList);
+        mMockVpn.setAlwaysOnPackage(null, false /* lockdown */, allowList);
         callback.expectBlockedStatusCallback(false, mWiFiNetworkAgent);
         defaultCallback.expectBlockedStatusCallback(false, mWiFiNetworkAgent);
         vpnUidCallback.assertNoCallback();
@@ -9071,8 +9068,7 @@
 
         // Add our UID to the allowlist and re-enable lockdown, expect network is not blocked.
         allowList.add(TEST_PACKAGE_NAME);
-        mVpnManagerService.setAlwaysOnVpnPackage(userId, ALWAYS_ON_PACKAGE, true /* lockdown */,
-                allowList);
+        mMockVpn.setAlwaysOnPackage(ALWAYS_ON_PACKAGE, true /* lockdown */, allowList);
         callback.assertNoCallback();
         defaultCallback.assertNoCallback();
         vpnUidCallback.assertNoCallback();
@@ -9110,12 +9106,11 @@
 
         // Disable lockdown, remove our UID from the allowlist, and re-enable lockdown.
         // Everything should now be blocked.
-        mVpnManagerService.setAlwaysOnVpnPackage(userId, null, false /* lockdown */, allowList);
+        mMockVpn.setAlwaysOnPackage(null, false /* lockdown */, allowList);
         waitForIdle();
         expectNetworkRejectNonSecureVpn(inOrder, false, uidRangeParcelsAlsoExcludingUs);
         allowList.clear();
-        mVpnManagerService.setAlwaysOnVpnPackage(userId, ALWAYS_ON_PACKAGE, true /* lockdown */,
-                allowList);
+        mMockVpn.setAlwaysOnPackage(ALWAYS_ON_PACKAGE, true /* lockdown */, allowList);
         waitForIdle();
         expectNetworkRejectNonSecureVpn(inOrder, true, uidRangeParcels);
         defaultCallback.expectBlockedStatusCallback(true, mWiFiNetworkAgent);
@@ -9130,7 +9125,7 @@
         assertNetworkInfo(TYPE_WIFI, DetailedState.BLOCKED);
 
         // Disable lockdown. Everything is unblocked.
-        mVpnManagerService.setAlwaysOnVpnPackage(userId, null, false /* lockdown */, allowList);
+        mMockVpn.setAlwaysOnPackage(null, false /* lockdown */, allowList);
         defaultCallback.expectBlockedStatusCallback(false, mWiFiNetworkAgent);
         assertBlockedCallbackInAnyOrder(callback, false, mWiFiNetworkAgent, mCellNetworkAgent);
         vpnUidCallback.assertNoCallback();
@@ -9144,8 +9139,7 @@
 
         // Enable and disable an always-on VPN package without lockdown. Expect no changes.
         reset(mMockNetd);
-        mVpnManagerService.setAlwaysOnVpnPackage(userId, ALWAYS_ON_PACKAGE, false /* lockdown */,
-                allowList);
+        mMockVpn.setAlwaysOnPackage(ALWAYS_ON_PACKAGE, false /* lockdown */, allowList);
         inOrder.verify(mMockNetd, never()).networkRejectNonSecureVpn(anyBoolean(), any());
         callback.assertNoCallback();
         defaultCallback.assertNoCallback();
@@ -9158,7 +9152,7 @@
         assertNetworkInfo(TYPE_MOBILE, DetailedState.DISCONNECTED);
         assertNetworkInfo(TYPE_WIFI, DetailedState.CONNECTED);
 
-        mVpnManagerService.setAlwaysOnVpnPackage(userId, null, false /* lockdown */, allowList);
+        mMockVpn.setAlwaysOnPackage(null, false /* lockdown */, allowList);
         inOrder.verify(mMockNetd, never()).networkRejectNonSecureVpn(anyBoolean(), any());
         callback.assertNoCallback();
         defaultCallback.assertNoCallback();
@@ -9172,8 +9166,7 @@
         assertNetworkInfo(TYPE_WIFI, DetailedState.CONNECTED);
 
         // Enable lockdown and connect a VPN. The VPN is not blocked.
-        mVpnManagerService.setAlwaysOnVpnPackage(userId, ALWAYS_ON_PACKAGE, true /* lockdown */,
-                allowList);
+        mMockVpn.setAlwaysOnPackage(ALWAYS_ON_PACKAGE, true /* lockdown */, allowList);
         defaultCallback.expectBlockedStatusCallback(true, mWiFiNetworkAgent);
         assertBlockedCallbackInAnyOrder(callback, true, mWiFiNetworkAgent, mCellNetworkAgent);
         vpnUidCallback.assertNoCallback();
@@ -15707,6 +15700,45 @@
         mCm.unregisterNetworkCallback(cb);
     }
 
+    @Test
+    public void testSanitizedCapabilitiesFromAgentDoesNotMutateArgument()
+            throws Exception {
+        // This NetworkCapabilities builds an usual object to maximize the chance that this requires
+        // sanitization, so we have a high chance to detect any changes to the original.
+        final NetworkCapabilities unsanitized = new NetworkCapabilities.Builder()
+                .withoutDefaultCapabilities()
+                .addTransportType(TRANSPORT_WIFI)
+                .addCapability(NET_CAPABILITY_INTERNET)
+                .setOwnerUid(12345)
+                .setAdministratorUids(new int[] {12345, 23456, 34567})
+                .setLinkUpstreamBandwidthKbps(20)
+                .setLinkDownstreamBandwidthKbps(10)
+                .setNetworkSpecifier(new EthernetNetworkSpecifier("foobar"))
+                .setTransportInfo(new WifiInfo.Builder().setBssid("AA:AA:AA:AA:AA:AA").build())
+                .setSignalStrength(-75)
+                .setSsid("SSID1")
+                .setRequestorUid(98765)
+                .setRequestorPackageName("TestPackage")
+                .setSubscriptionIds(Collections.singleton(Process.myUid()))
+                .setUids(UidRange.toIntRanges(uidRangesForUids(
+                        UserHandle.getUid(PRIMARY_USER, 10100),
+                        UserHandle.getUid(SECONDARY_USER, 10101),
+                        UserHandle.getUid(TERTIARY_USER, 10043))))
+                .setAllowedUids(Set.of(45678, 56789, 65432))
+                .setUnderlyingNetworks(List.of(new Network(99999)))
+                .build();
+        final NetworkCapabilities copyOfUnsanitized = new NetworkCapabilities(unsanitized);
+        final NetworkInfo info = new NetworkInfo(TYPE_MOBILE, TelephonyManager.NETWORK_TYPE_LTE,
+                ConnectivityManager.getNetworkTypeName(TYPE_MOBILE),
+                TelephonyManager.getNetworkTypeName(TelephonyManager.NETWORK_TYPE_LTE));
+        final NetworkAgentInfo agent = fakeNai(unsanitized, info);
+        agent.setDeclaredCapabilities(unsanitized);
+        final NetworkCapabilities sanitized = agent.getDeclaredCapabilitiesSanitized(
+                null /* carrierPrivilegeAuthenticator */);
+        assertEquals(copyOfUnsanitized, unsanitized);
+        assertNotEquals(sanitized, unsanitized);
+    }
+
     /**
      * Validate request counts are counted accurately on setProfileNetworkPreference on set/replace.
      */