Snap for 11166233 from 3255d59c03ace580cdfd13a4046a026c23b16cad to mainline-healthfitness-release

Change-Id: Ic0c343456184af8efeb7749cb9cd46f9c5655e61
diff --git a/Android.bp b/Android.bp
index 61b4a1e..8ee5d6f 100644
--- a/Android.bp
+++ b/Android.bp
@@ -89,7 +89,7 @@
     sdk_version: module_34_version,
     libs: [
         "framework-configinfrastructure",
-        "framework-connectivity",
+        "framework-connectivity.stubs.module_lib",
         "framework-connectivity-t",
         "framework-statsd",
         "framework-wifi",
@@ -311,6 +311,7 @@
         "androidx.annotation_annotation",
         "modules-utils-build_system",
         "modules-utils-preconditions",
+        "modules-utils-shell-command-handler",
         "modules-utils-statemachine",
         "netd_aidl_interface-lateststable-java",
         "networkstack-client",
diff --git a/src/android/net/dhcp6/Dhcp6Client.java b/src/android/net/dhcp6/Dhcp6Client.java
index 1511669..a107b9c 100644
--- a/src/android/net/dhcp6/Dhcp6Client.java
+++ b/src/android/net/dhcp6/Dhcp6Client.java
@@ -638,6 +638,33 @@
         }
     }
 
+    // Create an IPv6 address from the interface mac address with IFA_F_MANAGETEMPADDR
+    // flag, kernel will create another privacy IPv6 address on behalf of user space.
+    // We don't need to remember IPv6 addresses that need to extend the lifetime every
+    // time it enters BoundState.
+    private boolean addInterfaceAddress(@NonNull final Inet6Address address,
+            @NonNull final IaPrefixOption ipo) {
+        final int flags = IFA_F_NOPREFIXROUTE | IFA_F_MANAGETEMPADDR | IFA_F_NODAD;
+        final long now = SystemClock.elapsedRealtime();
+        final long deprecationTime = now + ipo.preferred;
+        final long expirationTime = now + ipo.valid;
+        final LinkAddress la = new LinkAddress(address, RFC7421_PREFIX_LENGTH, flags,
+                RT_SCOPE_UNIVERSE /* scope */, deprecationTime, expirationTime);
+        if (!la.isGlobalPreferred()) {
+            Log.e(TAG, la + " is not a global preferred IPv6 address");
+            return false;
+        }
+        if (!NetlinkUtils.sendRtmNewAddressRequest(mIface.index, address,
+                (short) RFC7421_PREFIX_LENGTH,
+                flags, (byte) RT_SCOPE_UNIVERSE /* scope */,
+                ipo.preferred, ipo.valid)) {
+            Log.e(TAG, "Failed to set IPv6 address " + address.getHostAddress()
+                    + "%" + mIface.index);
+            return false;
+        }
+        return true;
+    }
+
     /**
      * Client has already obtained the lease(e.g. IA_PD option) from server and stays in Bound
      * state until T1 expires, and then transition to Renew state to extend the lease duration.
@@ -650,35 +677,21 @@
 
             // TODO: roll back to SOLICIT state after a delay if something wrong happens
             // instead of returning directly.
-            // The server may assign a prefix with length less than 64. To support automatic address
-            // generation (with IFA_F_MANAGETEMPADDR), we always set the address prefix length to
-            // 64, even if the delegated prefix length is less than 64. However, the unreachable
-            // route should still use the assigned prefix length.
-            final IpPrefix routePrefix = mReply.ipos.get(0).getIpPrefix();
-            final IpPrefix addressPrefix = new IpPrefix(routePrefix.getAddress(),
-                    RFC7421_PREFIX_LENGTH);
-            // Create EUI-64, so we don't need to remember IPv6 addresses that need to extend the
-            // lifetime every time it enters BoundState.
-            final Inet6Address address = createInet6AddressFromEui64(addressPrefix,
-                    macAddressToEui64(mIface.macAddr));
-            final int flags = IFA_F_NOPREFIXROUTE | IFA_F_MANAGETEMPADDR | IFA_F_NODAD;
-            final long now = SystemClock.elapsedRealtime();
-            final IaPrefixOption ipo = mReply.ipos.get(0);
-            final long deprecationTime = now + ipo.preferred;
-            final long expirationTime = now + ipo.valid;
-            final LinkAddress la = new LinkAddress(address, RFC7421_PREFIX_LENGTH, flags,
-                    RT_SCOPE_UNIVERSE /* scope */, deprecationTime, expirationTime);
-            if (!la.isGlobalPreferred()) {
-                Log.e(TAG, la + " is not a global IPv6 address, ignoring");
-                return;
-            }
-            if (!NetlinkUtils.sendRtmNewAddressRequest(mIface.index, address,
-                    (short) RFC7421_PREFIX_LENGTH,
-                    flags, (byte) RT_SCOPE_UNIVERSE /* scope */,
-                    ipo.preferred, ipo.valid)) {
-                Log.e(TAG, "Failed to set IPv6 address " + address.getHostAddress()
-                        + "%" + mIface.index);
-                return;
+            for (IaPrefixOption ipo : mReply.getValidIaPrefixes()) {
+                // TODO: The prefix with preferred/valid lifetime of 0 is valid, but client
+                // should stop using the prefix immediately. Actually kernel doesn't accept
+                // the address with valid lifetime of 0 and returns EINVAL when it sees that.
+                // We should send RTM_DELADDR netlink message to kernel to delete these addresses
+                // from the interface if any.
+                // Configure IPv6 addresses based on the delegated prefix(es) on the interface.
+                // We've checked that delegated prefix is valid upon receiving the response from
+                // DHCPv6 server, and the server may assign a prefix with length less than 64. So
+                // for SLAAC use case we always set the prefix length to 64 even if the delegated
+                // prefix length is less than 64.
+                final IpPrefix prefix = ipo.getIpPrefix();
+                final Inet6Address address = createInet6AddressFromEui64(prefix,
+                        macAddressToEui64(mIface.macAddr));
+                if (!addInterfaceAddress(address, ipo)) continue;
             }
             notifyPrefixDelegation(DHCP6_PD_SUCCESS, mReply);
         }
diff --git a/src/android/net/dhcp6/Dhcp6Packet.java b/src/android/net/dhcp6/Dhcp6Packet.java
index 250aa73..7a977e5 100644
--- a/src/android/net/dhcp6/Dhcp6Packet.java
+++ b/src/android/net/dhcp6/Dhcp6Packet.java
@@ -221,7 +221,7 @@
         }
 
         /**
-         * Check whether or not the delegated prefix in DHCPv6 packet is valid.
+         * Check whether or not the IA_PD option in DHCPv6 message is valid.
          *
          * TODO: ensure that the prefix has a reasonable lifetime, and the timers aren't too short.
          */
diff --git a/src/com/android/networkstack/metrics/stats.proto b/src/com/android/networkstack/metrics/stats.proto
index c09f082..bd5e62b 100644
--- a/src/com/android/networkstack/metrics/stats.proto
+++ b/src/com/android/networkstack/metrics/stats.proto
@@ -188,3 +188,73 @@
     // NUD neighbor type, default gateway, DNS server or both.
     optional .android.stats.connectivity.NudNeighborType neighbor_type = 3;
 }
+
+/**
+ * Logs Ip client RA(Router Advertisement) info
+ * Logged from:
+ * packages/modules/NetworkStack/src/android/net/ip/IpClient.java
+ */
+message IpClientRaInfoReported {
+    // The maximum number of distinct RAs (Router Advertisements).
+    optional int32 max_number_of_distinct_ras = 1;
+
+    // The number of zero lifetime RAs (Router Advertisements).
+    optional int32 number_of_zero_lifetime_ras = 2;
+
+    // The number of parsing error for RAs (Router Advertisements).
+    optional int32 number_of_parsing_error_ras = 3;
+
+    // The lowest router lifetime in seconds.
+    optional int32 lowest_router_lifetime_seconds = 4;
+
+    // The lowest valid lifetime of PIO (Prefix Information Option) in seconds.
+    optional int32 lowest_pio_valid_lifetime_seconds = 5;
+
+    // The lowest route lifetime of RIO (Route Information Option) in seconds.
+    optional int32 lowest_rio_route_lifetime_seconds = 6;
+
+    // The lowest lifetime of RDNSS (Recursive DNS Server Option) in seconds.
+    optional int32 lowest_rdnss_lifetime_seconds = 7;
+}
+
+/**
+ * Logs value of the APF counter.
+ */
+message ApfCounter {
+    // The name of APF counter.
+    optional .android.stats.connectivity.CounterName counter_name = 1;
+
+    // The value of APF counter.
+    optional int32 counter_value = 2;
+}
+
+
+message ApfCounterList {
+    repeated ApfCounter apf_counter = 1;
+}
+
+/**
+ * Logs APF session information event.
+ * Logged from:
+ * packages/modules/NetworkStack/src/android/net/apf/ApfFilter.java or
+ * packages/modules/NetworkStack/src/android/net/apf/LegacyApfFilter.java
+ */
+message ApfSessionInfoReported {
+    // The version of APF, where version = -1 equals APF disable.
+    optional int32 version = 1;
+
+    // The memory size of APF module.
+    optional int32 memory_size = 2;
+
+    // The values of all APF counters.
+    optional ApfCounterList apf_counter_list = 3;
+
+    // The duration of ip client in milliseconds.
+    optional int32 ip_client_session_duration_ms = 4;
+
+    // Number of times APF program updated.
+    optional int32 num_of_times_apf_program_updated = 5;
+
+    // Record the maximum of program size.
+    optional int32 max_program_size = 6;
+}
diff --git a/src/com/android/networkstack/netlink/TcpSocketTracker.java b/src/com/android/networkstack/netlink/TcpSocketTracker.java
index 4b66c0a..658fe8a 100644
--- a/src/com/android/networkstack/netlink/TcpSocketTracker.java
+++ b/src/com/android/networkstack/netlink/TcpSocketTracker.java
@@ -27,12 +27,14 @@
 import static android.system.OsConstants.SOL_SOCKET;
 import static android.system.OsConstants.SO_SNDTIMEO;
 
+import static com.android.net.module.util.FeatureVersions.FEATURE_IS_UID_NETWORKING_BLOCKED;
 import static com.android.net.module.util.NetworkStackConstants.DNS_OVER_TLS_PORT;
 import static com.android.net.module.util.netlink.NetlinkConstants.NLMSG_DONE;
 import static com.android.net.module.util.netlink.NetlinkConstants.SOCKDIAG_MSG_HEADER_SIZE;
 import static com.android.net.module.util.netlink.NetlinkConstants.SOCK_DIAG_BY_FAMILY;
 import static com.android.net.module.util.netlink.NetlinkUtils.DEFAULT_RECV_BUFSIZE;
 import static com.android.net.module.util.netlink.NetlinkUtils.IO_TIMEOUT_MS;
+import static com.android.networkstack.util.NetworkStackUtils.IGNORE_TCP_INFO_FOR_BLOCKED_UIDS;
 import static com.android.networkstack.util.NetworkStackUtils.SKIP_TCP_POLL_IN_LIGHT_DOZE;
 
 import android.annotation.TargetApi;
@@ -40,10 +42,12 @@
 import android.content.Context;
 import android.content.Intent;
 import android.content.IntentFilter;
+import android.net.ConnectivityManager;
 import android.net.INetd;
 import android.net.LinkProperties;
 import android.net.MarkMaskParcel;
 import android.net.Network;
+import android.net.NetworkCapabilities;
 import android.os.AsyncTask;
 import android.os.Build;
 import android.os.IBinder;
@@ -129,6 +133,8 @@
     private int mMinPacketsThreshold = DEFAULT_DATA_STALL_MIN_PACKETS_THRESHOLD;
     private int mTcpPacketsFailRateThreshold = DEFAULT_TCP_PACKETS_FAIL_PERCENTAGE;
 
+    // TODO: Remove doze mode solution since uid networking blocked traffic is filtered out by
+    //  the info provided by bpf maps.
     private final Object mDozeModeLock = new Object();
     @GuardedBy("mDozeModeLock")
     private boolean mInDozeMode = false;
@@ -139,8 +145,13 @@
     private boolean mInOpportunisticMode;
     @NonNull
     private LinkProperties mLinkProperties;
+    @NonNull
+    private NetworkCapabilities mNetworkCapabilities;
 
+    private final boolean mShouldDisableInDeepDoze;
     private final boolean mShouldDisableInLightDoze;
+    private final boolean mShouldIgnoreTcpInfoForBlockedUids;
+    private final ConnectivityManager mCm;
 
     @VisibleForTesting
     protected final DeviceConfig.OnPropertiesChangedListener mConfigListener =
@@ -158,8 +169,9 @@
                 }
             };
 
-    private static boolean isDeviceIdleModeChangedAction(Intent intent) {
-        return ACTION_DEVICE_IDLE_MODE_CHANGED.equals(intent.getAction());
+    private boolean isDeviceIdleModeChangedAction(Intent intent) {
+        return mShouldDisableInDeepDoze
+                && ACTION_DEVICE_IDLE_MODE_CHANGED.equals(intent.getAction());
     }
 
     @TargetApi(Build.VERSION_CODES.TIRAMISU)
@@ -180,7 +192,8 @@
                 // For tcp polling mechanism, there is no difference between deep doze mode and
                 // light doze mode. The deep doze mode and light doze mode block networking
                 // for uids in the same way, use single variable to control.
-                final boolean deviceIdle = powerManager.isDeviceIdleMode()
+                final boolean deviceIdle = (mShouldDisableInDeepDoze
+                        && powerManager.isDeviceIdleMode())
                         || (mShouldDisableInLightDoze && powerManager.isDeviceLightIdleMode());
                 setDozeMode(deviceIdle);
             }
@@ -191,7 +204,15 @@
         mDependencies = dps;
         mNetwork = network;
         mNetd = mDependencies.getNetd();
-        mShouldDisableInLightDoze = mDependencies.shouldDisableInLightDoze();
+        mShouldIgnoreTcpInfoForBlockedUids = mDependencies.shouldIgnoreTcpInfoForBlockedUids();
+
+        // Previous workarounds can be disabled if the device supports ignore blocked uids feature.
+        // To prevent inconsistencies and issues like broadcast receiver leaks, the feature flags
+        // are fixed after being read.
+        // TODO: Remove these workarounds when pre-T devices are no longer supported.
+        mShouldDisableInLightDoze = mDependencies.shouldDisableInLightDoze(
+                mShouldIgnoreTcpInfoForBlockedUids);
+        mShouldDisableInDeepDoze = !mShouldIgnoreTcpInfoForBlockedUids;
 
         // If the parcel is null, nothing should be matched which is achieved by the combination of
         // {@code NetlinkUtils#NULL_MASK} and {@code NetlinkUtils#UNKNOWN_MARK}.
@@ -205,7 +226,9 @@
                     family, InetDiagMessage.buildInetDiagReqForAliveTcpSockets(family));
         }
         mDependencies.addDeviceConfigChangedListener(mConfigListener);
-        mDependencies.addDeviceIdleReceiver(mDeviceIdleReceiver, mShouldDisableInLightDoze);
+        mDependencies.addDeviceIdleReceiver(mDeviceIdleReceiver, mShouldDisableInDeepDoze,
+                mShouldDisableInLightDoze);
+        mCm = mDependencies.getContext().getSystemService(ConnectivityManager.class);
     }
 
     @Nullable
@@ -252,6 +275,7 @@
 
             // Append TcpStats based on previous and current socket info.
             final TcpStat stat = new TcpStat();
+            final ArrayList<Integer> skippedBlockedUids = new ArrayList<>();
             mLatestReportedUids.clear();
             for (final SocketInfo newInfo : newSocketInfoList) {
                 final TcpStat diff = calculateLatestPacketsStat(newInfo,
@@ -271,11 +295,28 @@
                     continue;
                 }
 
+                if (mShouldIgnoreTcpInfoForBlockedUids) {
+                    // For backward-compatibility, NET_CAPABILITY_TEMPORARILY_NOT_METERED
+                    // is not referenced when deciding meteredness in NetworkPolicyManagerService.
+                    // Thus, whether to block metered networking should only be judged with
+                    // NET_CAPABILITY_NOT_METERED.
+                    final boolean metered = !mNetworkCapabilities.hasCapability(
+                            NetworkCapabilities.NET_CAPABILITY_NOT_METERED);
+                    final boolean uidBlocked = mCm.isUidNetworkingBlocked(newInfo.uid, metered);
+                    if (uidBlocked) {
+                        skippedBlockedUids.add(newInfo.uid);
+                        continue;
+                    }
+                }
+
                 if (diff != null) {
                     mLatestReportedUids.add(newInfo.uid);
                     stat.accumulate(diff);
                 }
             }
+            if (!skippedBlockedUids.isEmpty()) {
+                logd("Skip blocked uids: " + skippedBlockedUids);
+            }
 
             // Calculate mLatestReceiveCount, mSentSinceLastRecv and mLatestPacketFailPercentage.
             mSentSinceLastRecv = (stat.receivedCount == 0)
@@ -521,7 +562,8 @@
     /** Stops monitoring and releases resources. */
     public void quit() {
         mDependencies.removeDeviceConfigChangedListener(mConfigListener);
-        mDependencies.removeBroadcastReceiver(mDeviceIdleReceiver);
+        mDependencies.removeBroadcastReceiver(mDeviceIdleReceiver,
+                mShouldDisableInDeepDoze, mShouldDisableInLightDoze);
     }
 
     /**
@@ -616,6 +658,10 @@
         mLinkProperties = lp;
     }
 
+    public void setNetworkCapabilities(@NonNull NetworkCapabilities caps) {
+        mNetworkCapabilities = caps;
+    }
+
     /**
      * Dependencies class for testing.
      */
@@ -702,8 +748,14 @@
         /** Add receiver for detecting doze mode change to control TCP detection. */
         @TargetApi(Build.VERSION_CODES.TIRAMISU)
         public void addDeviceIdleReceiver(@NonNull final BroadcastReceiver receiver,
-                boolean shouldDisableInLightDoze) {
-            final IntentFilter intentFilter = new IntentFilter(ACTION_DEVICE_IDLE_MODE_CHANGED);
+                boolean shouldDisableInDeepDoze, boolean shouldDisableInLightDoze) {
+            // No need to register receiver if no related feature is enabled.
+            if (!shouldDisableInDeepDoze && !shouldDisableInLightDoze) return;
+
+            final IntentFilter intentFilter = new IntentFilter();
+            if (shouldDisableInDeepDoze) {
+                intentFilter.addAction(ACTION_DEVICE_IDLE_MODE_CHANGED);
+            }
             if (shouldDisableInLightDoze) {
                 intentFilter.addAction(ACTION_DEVICE_LIGHT_IDLE_MODE_CHANGED);
             }
@@ -711,7 +763,9 @@
         }
 
         /** Remove broadcast receiver. */
-        public void removeBroadcastReceiver(@NonNull final BroadcastReceiver receiver) {
+        public void removeBroadcastReceiver(@NonNull final BroadcastReceiver receiver,
+                boolean shouldDisableInDeepDoze, boolean shouldDisableInLightDoze) {
+            if (!shouldDisableInDeepDoze && !shouldDisableInLightDoze) return;
             mContext.unregisterReceiver(receiver);
         }
 
@@ -721,10 +775,27 @@
          * to deal with flag values changing at runtime.
          */
         @TargetApi(Build.VERSION_CODES.TIRAMISU)
-        public boolean shouldDisableInLightDoze() {
+        public boolean shouldDisableInLightDoze(boolean ignoreBlockedUidsSupported) {
             // Light doze mode status checking API is only available at T or later releases.
-            return SdkLevel.isAtLeastT() && DeviceConfigUtils.isNetworkStackFeatureNotChickenedOut(
+            if (!SdkLevel.isAtLeastT()) return false;
+
+            // Disable light doze mode design is replaced by ignoring blocked uids design.
+            if (ignoreBlockedUidsSupported) return false;
+
+            return DeviceConfigUtils.isNetworkStackFeatureNotChickenedOut(
                     mContext, SKIP_TCP_POLL_IN_LIGHT_DOZE);
         }
+
+        /**
+         * Get whether the ignore Tcp info for blocked uids is supported. This method should
+         * only be called once in the constructor, to ensure that the code does not need
+         * to deal with flag values changing at runtime.
+         */
+        public boolean shouldIgnoreTcpInfoForBlockedUids() {
+            return SdkLevel.isAtLeastT() && DeviceConfigUtils.isFeatureSupported(
+                    mContext, FEATURE_IS_UID_NETWORKING_BLOCKED)
+                    && DeviceConfigUtils.isNetworkStackFeatureNotChickenedOut(mContext,
+                    IGNORE_TCP_INFO_FOR_BLOCKED_UIDS);
+        }
     }
 }
diff --git a/src/com/android/networkstack/util/NetworkStackUtils.java b/src/com/android/networkstack/util/NetworkStackUtils.java
index 84a4491..2dbf86e 100755
--- a/src/com/android/networkstack/util/NetworkStackUtils.java
+++ b/src/com/android/networkstack/util/NetworkStackUtils.java
@@ -282,6 +282,13 @@
      */
     public static final String REEVALUATE_WHEN_RESUME = "reevaluate_when_resume";
 
+    /**
+     * Kill switch flag to disable the feature of ignoring Tcp socket info for uids which
+     * networking are blocked.
+     */
+    public static final String IGNORE_TCP_INFO_FOR_BLOCKED_UIDS =
+            "ignore_tcp_info_for_blocked_uids";
+
     static {
         System.loadLibrary("networkstackutilsjni");
     }
diff --git a/src/com/android/server/NetworkStackService.java b/src/com/android/server/NetworkStackService.java
index 368a6d4..40aee28 100644
--- a/src/com/android/server/NetworkStackService.java
+++ b/src/com/android/server/NetworkStackService.java
@@ -21,11 +21,15 @@
 import static android.net.dhcp.IDhcpServer.STATUS_UNKNOWN_ERROR;
 
 import static com.android.net.module.util.DeviceConfigUtils.getResBooleanConfig;
+import static com.android.net.module.util.FeatureVersions.FEATURE_IS_UID_NETWORKING_BLOCKED;
+import static com.android.networkstack.util.NetworkStackUtils.IGNORE_TCP_INFO_FOR_BLOCKED_UIDS;
+import static com.android.networkstack.util.NetworkStackUtils.SKIP_TCP_POLL_IN_LIGHT_DOZE;
 import static com.android.server.util.PermissionUtil.checkDumpPermission;
 
 import android.app.Service;
 import android.content.Context;
 import android.content.Intent;
+import android.net.ConnectivityManager;
 import android.net.IIpMemoryStore;
 import android.net.IIpMemoryStoreCallbacks;
 import android.net.INetd;
@@ -49,6 +53,7 @@
 import android.os.HandlerThread;
 import android.os.IBinder;
 import android.os.Looper;
+import android.os.ParcelFileDescriptor;
 import android.os.RemoteException;
 import android.text.TextUtils;
 import android.util.ArraySet;
@@ -59,6 +64,8 @@
 
 import com.android.internal.annotations.GuardedBy;
 import com.android.internal.util.IndentingPrintWriter;
+import com.android.modules.utils.BasicShellCommandHandler;
+import com.android.net.module.util.DeviceConfigUtils;
 import com.android.net.module.util.SharedLog;
 import com.android.networkstack.NetworkStackNotifier;
 import com.android.networkstack.R;
@@ -435,6 +442,20 @@
                 return;
             }
 
+            pw.println("Device Configs:");
+            pw.increaseIndent();
+            pw.println("SKIP_TCP_POLL_IN_LIGHT_DOZE="
+                    + DeviceConfigUtils.isNetworkStackFeatureNotChickenedOut(
+                            mContext, SKIP_TCP_POLL_IN_LIGHT_DOZE));
+            pw.println("FEATURE_IS_UID_NETWORKING_BLOCKED=" + DeviceConfigUtils.isFeatureSupported(
+                            mContext, FEATURE_IS_UID_NETWORKING_BLOCKED));
+            pw.println("IGNORE_TCP_INFO_FOR_BLOCKED_UIDS="
+                    + DeviceConfigUtils.isNetworkStackFeatureNotChickenedOut(mContext,
+                            IGNORE_TCP_INFO_FOR_BLOCKED_UIDS));
+            pw.decreaseIndent();
+            pw.println();
+
+
             pw.println("NetworkStack logs:");
             mLog.dump(fd, pw, args);
 
@@ -482,6 +503,69 @@
                     R.bool.config_no_sim_card_uses_neighbor_mcc, false));
         }
 
+        @Override
+        public int handleShellCommand(@NonNull ParcelFileDescriptor in,
+                @NonNull ParcelFileDescriptor out, @NonNull ParcelFileDescriptor err,
+                @NonNull String[] args) {
+            return new ShellCmd().exec(this, in.getFileDescriptor(), out.getFileDescriptor(),
+                    err.getFileDescriptor(), args);
+        }
+
+        private class ShellCmd extends BasicShellCommandHandler {
+            @Override
+            public int onCommand(String cmd) {
+                if (cmd == null) {
+                    return handleDefaultCommands(cmd);
+                }
+                final PrintWriter pw = getOutPrintWriter();
+                try {
+                    switch (cmd) {
+                        case "is-uid-networking-blocked":
+                            if (!DeviceConfigUtils.isFeatureSupported(mContext,
+                                    FEATURE_IS_UID_NETWORKING_BLOCKED)) {
+                                pw.println("API is unsupported");
+                                return -1;
+                            }
+
+                            // Usage : cmd network_stack is-uid-networking-blocked <uid> <metered>
+                            // If no argument, get and display the usage help.
+                            if (getRemainingArgsCount() != 2) {
+                                onHelp();
+                                return -1;
+                            }
+                            final int uid;
+                            final boolean metered;
+                            // If any fail, throws and output to the stdout.
+                            // Let the caller handle it.
+                            uid = Integer.parseInt(getNextArg());
+                            metered = Boolean.parseBoolean(getNextArg());
+                            final ConnectivityManager cm =
+                                    mContext.getSystemService(ConnectivityManager.class);
+                            pw.println(cm.isUidNetworkingBlocked(
+                                    uid, metered /* isNetworkMetered */));
+                            return 0;
+                        default:
+                            return handleDefaultCommands(cmd);
+                    }
+                } catch (Exception e) {
+                    pw.println(e);
+                }
+                return -1;
+            }
+
+            @Override
+            public void onHelp() {
+                PrintWriter pw = getOutPrintWriter();
+                pw.println("NetworkStack service commands:");
+                pw.println("  help");
+                pw.println("    Print this help text.");
+                pw.println("  is-uid-networking-blocked <uid> <metered>");
+                pw.println("    Get whether the networking is blocked for given uid and metered.");
+                pw.println("    <uid>: The target uid.");
+                pw.println("    <metered>: [true|false], Whether the target network is metered.");
+            }
+        }
+
         /**
          * Dump version information of the module and detected system version.
          */
diff --git a/src/com/android/server/connectivity/NetworkMonitor.java b/src/com/android/server/connectivity/NetworkMonitor.java
index 5432a13..8c10138 100755
--- a/src/com/android/server/connectivity/NetworkMonitor.java
+++ b/src/com/android/server/connectivity/NetworkMonitor.java
@@ -942,6 +942,7 @@
                 // Initialization.
                 tst.setOpportunisticMode(false);
                 tst.setLinkProperties(mLinkProperties);
+                tst.setNetworkCapabilities(mNetworkCapabilities);
             }
             Log.d(TAG, "Starting on network " + mNetwork
                     + " with capport HTTPS URL " + Arrays.toString(mCaptivePortalHttpsUrls)
@@ -1155,6 +1156,10 @@
             // VPN.
             sendMessage(CMD_FORCE_REEVALUATION, NO_UID, 0 /* forceAccept */);
         }
+        final TcpSocketTracker tst = getTcpSocketTracker();
+        if (tst != null) {
+            tst.setNetworkCapabilities(newCap);
+        }
 
         mNetworkCapabilities = newCap;
         suppressNotificationIfNetworkRestricted();
diff --git a/tests/integration/common/android/net/ip/IpClientIntegrationTestCommon.java b/tests/integration/common/android/net/ip/IpClientIntegrationTestCommon.java
index 1dd5e40..d239379 100644
--- a/tests/integration/common/android/net/ip/IpClientIntegrationTestCommon.java
+++ b/tests/integration/common/android/net/ip/IpClientIntegrationTestCommon.java
@@ -5034,6 +5034,44 @@
         ));
     }
 
+    @Test
+    public void testDhcp6Pd_multiplePrefixesWithInvalidPrefix() throws Exception {
+        final IpPrefix valid = new IpPrefix("2001:db8:1::/64");
+        final IpPrefix invalid = new IpPrefix("2001:db8:2::/64"); // preferred lft > valid lft
+        final IaPrefixOption validIpo = buildIaPrefixOption(valid, 4500 /* preferred */,
+                7200 /* valid */);
+        final IaPrefixOption invalidIpo = buildIaPrefixOption(invalid, 4500 /* preferred */,
+                3000 /* valid */);
+
+        prepareDhcp6PdTest();
+        handleDhcp6Packets(Arrays.asList(invalidIpo, validIpo), 3600 /* t1 */, 4500 /* t2 */,
+                true /* shouldReplyRapidCommit */);
+        final ArgumentCaptor<LinkProperties> captor = ArgumentCaptor.forClass(LinkProperties.class);
+        verify(mCb, timeout(TEST_TIMEOUT_MS)).onProvisioningSuccess(captor.capture());
+        final LinkProperties lp = captor.getValue();
+        assertTrue(hasIpv6AddressPrefixedWith(lp, valid));
+        assertFalse(hasIpv6AddressPrefixedWith(lp, invalid));
+    }
+
+    @Test
+    public void testDhcp6Pd_multiplePrefixesWithPrefixValidLifetimeOfZero() throws Exception {
+        final IpPrefix valid = new IpPrefix("2001:db8:1::/64");
+        final IpPrefix invalid = new IpPrefix("2001:db8:2::/64"); // preferred/valid lft 0
+        final IaPrefixOption validIpo = buildIaPrefixOption(valid, 4500 /* preferred */,
+                7200 /* valid */);
+        final IaPrefixOption invalidIpo = buildIaPrefixOption(invalid, 0 /* preferred */,
+                0 /* valid */);
+
+        prepareDhcp6PdTest();
+        handleDhcp6Packets(Arrays.asList(invalidIpo, validIpo), 3600 /* t1 */, 4500 /* t2 */,
+                true /* shouldReplyRapidCommit */);
+        final ArgumentCaptor<LinkProperties> captor = ArgumentCaptor.forClass(LinkProperties.class);
+        verify(mCb, timeout(TEST_TIMEOUT_MS)).onProvisioningSuccess(captor.capture());
+        final LinkProperties lp = captor.getValue();
+        assertTrue(hasIpv6AddressPrefixedWith(lp, valid));
+        assertFalse(hasIpv6AddressPrefixedWith(lp, invalid));
+    }
+
     private void prepareDhcp6PdRenewTest() throws Exception {
         final IpPrefix prefix = new IpPrefix("2001:db8:1::/64");
         prepareDhcp6PdTest();
diff --git a/tests/unit/src/com/android/networkstack/netlink/TcpSocketTrackerTest.java b/tests/unit/src/com/android/networkstack/netlink/TcpSocketTrackerTest.java
index ab9816c..3857b04 100644
--- a/tests/unit/src/com/android/networkstack/netlink/TcpSocketTrackerTest.java
+++ b/tests/unit/src/com/android/networkstack/netlink/TcpSocketTrackerTest.java
@@ -16,8 +16,12 @@
 
 package com.android.networkstack.netlink;
 
+import static android.net.NetworkCapabilities.NET_CAPABILITY_INTERNET;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_NOT_METERED;
+import static android.net.NetworkCapabilities.TRANSPORT_CELLULAR;
 import static android.net.util.DataStallUtils.CONFIG_TCP_PACKETS_FAIL_PERCENTAGE;
 import static android.net.util.DataStallUtils.DEFAULT_TCP_PACKETS_FAIL_PERCENTAGE;
+import static android.os.PowerManager.ACTION_DEVICE_IDLE_MODE_CHANGED;
 import static android.os.PowerManager.ACTION_DEVICE_LIGHT_IDLE_MODE_CHANGED;
 import static android.provider.DeviceConfig.NAMESPACE_CONNECTIVITY;
 import static android.system.OsConstants.AF_INET;
@@ -33,18 +37,22 @@
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import android.annotation.IntDef;
 import android.content.BroadcastReceiver;
 import android.content.Context;
 import android.content.Intent;
+import android.net.ConnectivityManager;
 import android.net.INetd;
 import android.net.InetAddresses;
 import android.net.LinkProperties;
 import android.net.MarkMaskParcel;
 import android.net.Network;
+import android.net.NetworkCapabilities;
 import android.os.Build;
 import android.os.PowerManager;
 import android.util.Log;
@@ -57,6 +65,7 @@
 import com.android.net.module.util.netlink.NetlinkUtils;
 import com.android.net.module.util.netlink.StructNlMsgHdr;
 import com.android.testutils.DevSdkIgnoreRule;
+import com.android.testutils.DevSdkIgnoreRule.IgnoreAfter;
 import com.android.testutils.DevSdkIgnoreRule.IgnoreUpTo;
 
 import libcore.util.HexEncoding;
@@ -71,6 +80,8 @@
 import org.mockito.MockitoAnnotations;
 
 import java.io.FileDescriptor;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
 import java.net.InetAddress;
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
@@ -118,10 +129,22 @@
     private static final int TEST_NETID2_FWMARK = 0x1A85;
     private static final int NETID_MASK = 0xffff;
     private static final int TEST_UID1 = 1234;
+    private static final int TEST_UID2 = TEST_UID1 + 1;
     private static final short TEST_DST_PORT = 29113;
     private static final long TEST_COOKIE1 = 43387759684916L;
     private static final long TEST_COOKIE2 = TEST_COOKIE1 + 1;
     private static final InetAddress TEST_DNS1 = InetAddresses.parseNumericAddress("8.8.8.8");
+
+    private static final NetworkCapabilities CELL_METERED_CAPABILITIES =
+            new NetworkCapabilities()
+                    .addTransportType(TRANSPORT_CELLULAR)
+                    .addCapability(NET_CAPABILITY_INTERNET);
+
+    private static final NetworkCapabilities CELL_NOT_METERED_CAPABILITIES =
+            new NetworkCapabilities()
+                    .addTransportType(TRANSPORT_CELLULAR)
+                    .addCapability(NET_CAPABILITY_INTERNET)
+                    .addCapability(NET_CAPABILITY_NOT_METERED);
     @Mock private TcpSocketTracker.Dependencies mDependencies;
     @Mock private INetd mNetd;
     private final Network mNetwork = new Network(TEST_NETID1);
@@ -129,6 +152,7 @@
     private TerribleFailureHandler mOldWtfHandler;
     @Mock private Context mContext;
     @Mock private PowerManager mPowerManager;
+    @Mock private ConnectivityManager mCm;
 
     @Rule
     public final DevSdkIgnoreRule mIgnoreRule = new DevSdkIgnoreRule();
@@ -146,11 +170,13 @@
                 eq(NAMESPACE_CONNECTIVITY),
                 eq(CONFIG_TCP_PACKETS_FAIL_PERCENTAGE),
                 anyInt())).thenReturn(DEFAULT_TCP_PACKETS_FAIL_PERCENTAGE);
-        when(mDependencies.shouldDisableInLightDoze()).thenReturn(true);
+        when(mDependencies.shouldDisableInLightDoze(anyBoolean())).thenReturn(true);
 
         when(mNetd.getFwmarkForNetwork(eq(TEST_NETID1)))
                 .thenReturn(makeMarkMaskParcel(NETID_MASK, TEST_NETID1_FWMARK));
+        doReturn(mContext).when(mDependencies).getContext();
         doReturn(mPowerManager).when(mContext).getSystemService(PowerManager.class);
+        doReturn(mCm).when(mContext).getSystemService(ConnectivityManager.class);
     }
 
     @After
@@ -264,7 +290,7 @@
         testLp.addDnsServer(TEST_DNS1);
         tst.setLinkProperties(testLp);
         doReturn(getByteBufferFromHexString(composeSockDiagTcpHex(9, 10)
-                + composeSockDiagTcpHex(9, 10, DNS_OVER_TLS_PORT, TEST_COOKIE2)
+                + composeSockDiagTcpHex(9, 10, DNS_OVER_TLS_PORT, TEST_COOKIE2, TEST_UID1)
                 + NLMSG_DONE_HEX))
                 .when(mDependencies).recvMessage(any());
         assertTrue(tst.pollSocketsInfo());
@@ -282,7 +308,7 @@
         testLp.addValidatedPrivateDnsServer(TEST_DNS1);
         tst.setLinkProperties(testLp);
         doReturn(getByteBufferFromHexString(composeSockDiagTcpHex(10, 12)
-                + composeSockDiagTcpHex(11, 12, DNS_OVER_TLS_PORT, TEST_COOKIE2)
+                + composeSockDiagTcpHex(11, 12, DNS_OVER_TLS_PORT, TEST_COOKIE2, TEST_UID1)
                 + NLMSG_DONE_HEX))
                 .when(mDependencies).recvMessage(any());
         assertTrue(tst.pollSocketsInfo());
@@ -296,7 +322,7 @@
         // polling cycle.
         tst.setOpportunisticMode(false);
         doReturn(getByteBufferFromHexString(composeSockDiagTcpHex(11, 14)
-                + composeSockDiagTcpHex(13, 14, DNS_OVER_TLS_PORT, TEST_COOKIE2)
+                + composeSockDiagTcpHex(13, 14, DNS_OVER_TLS_PORT, TEST_COOKIE2, TEST_UID1)
                 + NLMSG_DONE_HEX))
                 .when(mDependencies).recvMessage(any());
         assertTrue(tst.pollSocketsInfo());
@@ -306,6 +332,101 @@
         assertFalse(tst.isDataStallSuspected());
     }
 
+    @IgnoreAfter(Build.VERSION_CODES.S_V2)
+    @Test
+    public void testPollSocketsInfo_ignoreBlockedUid_featureDisabled_beforeT() throws Exception {
+        doTestPollSocketsInfo_ignoreBlockedUid_featureDisabled();
+    }
+
+    @IgnoreUpTo(Build.VERSION_CODES.S_V2)
+    @Test
+    public void testPollSocketsInfo_ignoreBlockedUid_featureDisabled_TOrAbove() throws Exception {
+        doTestPollSocketsInfo_ignoreBlockedUid_featureDisabled();
+        verify(mCm, never()).isUidNetworkingBlocked(anyInt(), anyBoolean());
+    }
+
+    private void doTestPollSocketsInfo_ignoreBlockedUid_featureDisabled() throws Exception {
+        doReturn(false).when(mDependencies).shouldIgnoreTcpInfoForBlockedUids();
+        final TcpSocketTracker tst = new TcpSocketTracker(mDependencies, mNetwork);
+        // Simulate 1 message with data stall happened.
+        doReturn(getByteBufferFromHexString(
+                composeSockDiagTcpHex(4, 10) + NLMSG_DONE_HEX))
+                .when(mDependencies).recvMessage(any());
+        assertTrue(tst.pollSocketsInfo());
+        // 4 retran / 10 sent = 40 percent.
+        assertEquals(40, tst.getLatestPacketFailPercentage());
+        assertEquals(10, tst.getSentSinceLastRecv());
+        assertFalse(tst.isDataStallSuspected());
+
+        // With the feature disabled, append another message with blocked uid, verify the
+        // traffic of networking-blocked uid is not filtered.
+        doReturn(getByteBufferFromHexString(composeSockDiagTcpHex(9, 10)
+                + composeSockDiagTcpHex(5, 10, DNS_OVER_TLS_PORT, TEST_COOKIE2, TEST_UID2)
+                + NLMSG_DONE_HEX))
+                .when(mDependencies).recvMessage(any());
+        assertTrue(tst.pollSocketsInfo());
+        // 5 + 5 retrans / 10 sent = 100 percent.
+        assertEquals(100, tst.getLatestPacketFailPercentage());
+        assertEquals(20, tst.getSentSinceLastRecv());
+        assertTrue(tst.isDataStallSuspected());
+    }
+
+    // The feature is not enabled on pre-T device, because it needs bpf support.
+    @IgnoreUpTo(Build.VERSION_CODES.S_V2)
+    @Test
+    public void testPollSocketsInfo_ignoreBlockedUid_featureEnabled() throws Exception {
+        doReturn(true).when(mDependencies).shouldIgnoreTcpInfoForBlockedUids();
+        final TcpSocketTracker tst = new TcpSocketTracker(mDependencies, mNetwork);
+        tst.setNetworkCapabilities(CELL_NOT_METERED_CAPABILITIES);
+        doReturn(true).when(mCm).isUidNetworkingBlocked(TEST_UID2, false /* metered */);
+        // With the feature enabled, append another message with blocked uid, verify the
+        // traffic of networking-blocked uid is filtered out.
+        doReturn(getByteBufferFromHexString(composeSockDiagTcpHex(4, 10)
+                + composeSockDiagTcpHex(6, 12, DNS_OVER_TLS_PORT, TEST_COOKIE2, TEST_UID2)
+                + NLMSG_DONE_HEX))
+                .when(mDependencies).recvMessage(any());
+        assertTrue(tst.pollSocketsInfo());
+        assertEquals(40, tst.getLatestPacketFailPercentage());
+        assertEquals(10, tst.getSentSinceLastRecv());
+        assertFalse(tst.isDataStallSuspected());
+
+        // Unblock traffic of the uid, verify the traffic of the uid is not filtered.
+        doReturn(false).when(mCm).isUidNetworkingBlocked(TEST_UID2, false /* metered */);
+        doReturn(getByteBufferFromHexString(composeSockDiagTcpHex(4, 10)
+                + composeSockDiagTcpHex(8, 14, DNS_OVER_TLS_PORT, TEST_COOKIE2, TEST_UID2)
+                + NLMSG_DONE_HEX))
+                .when(mDependencies).recvMessage(any());
+        assertTrue(tst.pollSocketsInfo());
+        // Lost 2 / 2 sent = 100 percent.
+        assertEquals(100, tst.getLatestPacketFailPercentage());
+        assertEquals(12, tst.getSentSinceLastRecv());
+        assertTrue(tst.isDataStallSuspected());
+    }
+
+    // The feature is not enabled on pre-T device, because it needs bpf support.
+    @IgnoreUpTo(Build.VERSION_CODES.S_V2)
+    @Test
+    public void testPollSocketsInfo_ignoreBlockedUid_featureEnabled_dataSaver() throws Exception {
+        doReturn(true).when(mDependencies).shouldIgnoreTcpInfoForBlockedUids();
+        final TcpSocketTracker tst = new TcpSocketTracker(mDependencies, mNetwork);
+
+        tst.setNetworkCapabilities(CELL_NOT_METERED_CAPABILITIES);
+        final ByteBuffer mockMessage = getByteBufferFromHexString(composeSockDiagTcpHex(4, 10)
+                + NLMSG_DONE_HEX);
+        doReturn(mockMessage).when(mDependencies).recvMessage(any());
+        assertTrue(tst.pollSocketsInfo());
+        verify(mCm).isUidNetworkingBlocked(TEST_UID1, false /* metered */);
+
+        // Verify the metered parameter will be correctly passed to ConnectivityManager.
+        tst.setNetworkCapabilities(CELL_METERED_CAPABILITIES);
+        mockMessage.rewind(); // Reset read position to 0 since the same buffer is used.
+        assertTrue(tst.pollSocketsInfo());
+        verify(mCm).isUidNetworkingBlocked(TEST_UID1, true /* metered */);
+
+        // Correctness of the logic which handling different blocked status is
+        // verified in other tests, see {@code testPollSocketsInfo_ignoreBlockedUid_featureEnabled}.
+    }
+
     @Test
     public void testTcpInfoParsingWithMultipleMsgs() throws Exception {
         final TcpSocketTracker tst = new TcpSocketTracker(mDependencies, mNetwork);
@@ -470,10 +591,11 @@
     }
 
     private static String composeSockDiagTcpHex(int retrans, int sent) {
-        return composeSockDiagTcpHex(retrans, sent, TEST_DST_PORT, TEST_COOKIE1);
+        return composeSockDiagTcpHex(retrans, sent, TEST_DST_PORT, TEST_COOKIE1, TEST_UID1);
     }
 
-    private static String composeSockDiagTcpHex(int retrans, int sent, short dstPort, long cookie) {
+    private static String composeSockDiagTcpHex(int retrans, int sent, short dstPort,
+            long cookie, int uid) {
         return // struct nlmsghdr.
                 "14010000"          // length = 276
                 + "1400"            // type = SOCK_DIAG_BY_FAMILY
@@ -487,17 +609,17 @@
                 + "00"              // retrans
                 // inet_diag_sockid: ports and addresses are always in big endian,
                 // see StructInetDiagSockId.
-                + "DEA5"                                               // idiag_sport = 56997
-                + getHexStringFromShort(dstPort, ByteOrder.BIG_ENDIAN) // idiag_dport
-                + "0a006402000000000000000000000000"                   // idiag_src = 10.0.100.2
-                + "08080808000000000000000000000000"                   // idiag_dst = 8.8.8.8
-                + "00000000"                                           // idiag_if
-                + getHexStringFromLong(cookie)                         // idiag_cookie
-                + "00000000"                                           // idiag_expires
-                + "00000000"                                           // idiag_rqueue
-                + "00000000"                                           // idiag_wqueue
-                + getHexStringFromInt(TEST_UID1)                       // idiag_uid
-                + "00000000"                                           // idiag_inode
+                + "DEA5"                                                // idiag_sport = 56997
+                + getHexStringFromShort(dstPort, ByteOrder.BIG_ENDIAN)  // idiag_dport
+                + "0a006402000000000000000000000000"                    // idiag_src = 10.0.100.2
+                + "08080808000000000000000000000000"                    // idiag_dst = 8.8.8.8
+                + "00000000"                                            // idiag_if
+                + getHexStringFromLong(cookie)                          // idiag_cookie
+                + "00000000"                                            // idiag_expires
+                + "00000000"                                            // idiag_rqueue
+                + "00000000"                                            // idiag_wqueue
+                + getHexStringFromInt(uid)                              // idiag_uid
+                + "00000000"                                            // idiag_inode
                 // rtattr
                 + "0500"            // len = 5
                 + "0800"            // type = 8
@@ -553,50 +675,76 @@
                 + "0000000000000000"; // deliverRate = 0
     }
 
+    private static final int DEEP_DOZE = 0;
+    private static final int LIGHT_DOZE = 1;
+
+    @Retention(RetentionPolicy.SOURCE)
+    @IntDef(value = {
+            DEEP_DOZE,
+            LIGHT_DOZE
+    })
+    private @interface DozeModeType {}
+
     @Test
-    public void testTcpInfoParsingWithDozeMode() throws Exception {
-        final TcpSocketTracker tst = new TcpSocketTracker(mDependencies, mNetwork);
-        final ArgumentCaptor<BroadcastReceiver> receiverCaptor =
-                ArgumentCaptor.forClass(BroadcastReceiver.class);
+    public void testTcpInfoParsingWithDozeMode_enabled() throws Exception {
+        doReturn(false).when(mDependencies).shouldIgnoreTcpInfoForBlockedUids();
+        doReturn(false).when(mDependencies).shouldDisableInLightDoze(anyBoolean());
+        doTestTcpInfoDisableParsingWithDozeMode(DEEP_DOZE, true /* featureEnabled */);
+    }
 
-        verify(mDependencies).addDeviceIdleReceiver(receiverCaptor.capture(), anyBoolean());
-        setupNormalTestTcpInfo();
-        assertTrue(tst.pollSocketsInfo());
-
-        // Lower the threshold.
-        when(mDependencies.getDeviceConfigPropertyInt(any(), eq(CONFIG_TCP_PACKETS_FAIL_PERCENTAGE),
-                anyInt())).thenReturn(40);
-
-        // Trigger a config update.
-        tst.mConfigListener.onPropertiesChanged(null /* properties */);
-        assertEquals(10, tst.getSentSinceLastRecv());
-        assertEquals(50, tst.getLatestPacketFailPercentage());
-        assertTrue(tst.isDataStallSuspected());
-
-        // Enable doze mode, verify counters are not updated.
-        doReturn(true).when(mPowerManager).isDeviceIdleMode();
-        final BroadcastReceiver receiver = receiverCaptor.getValue();
-        receiver.onReceive(mContext, new Intent(PowerManager.ACTION_DEVICE_IDLE_MODE_CHANGED));
-        assertFalse(tst.pollSocketsInfo());
-        assertEquals(10, tst.getSentSinceLastRecv());
-        assertEquals(50, tst.getLatestPacketFailPercentage());
-        assertFalse(tst.isDataStallSuspected());
+    // Ignore blocked uids is supported on T. Thus, for pre-T device this feature is always
+    // needed since there is no replacement.
+    @IgnoreUpTo(Build.VERSION_CODES.S_V2)
+    @Test
+    public void testTcpInfoParsingWithDozeMode_disabled() throws Exception {
+        doReturn(true).when(mDependencies).shouldIgnoreTcpInfoForBlockedUids();
+        doReturn(false).when(mDependencies).shouldDisableInLightDoze(anyBoolean());
+        doTestTcpInfoDisableParsingWithDozeMode(DEEP_DOZE, false /* featureEnabled */);
     }
 
     @Test @IgnoreUpTo(Build.VERSION_CODES.S_V2)
     public void testTcpInfoDisableParsingWithLightDozeMode_enabled() throws Exception {
+        doReturn(true).when(mDependencies).shouldDisableInLightDoze(anyBoolean());
+        doTestTcpInfoDisableParsingWithDozeMode(LIGHT_DOZE, true /* featureEnabled */);
+    }
+
+    @Test @IgnoreUpTo(Build.VERSION_CODES.S_V2)
+    public void testTcpInfoDisableParsingWithLightDozeMode_disabled() throws Exception {
+        doReturn(false).when(mDependencies).shouldDisableInLightDoze(anyBoolean());
+        doTestTcpInfoDisableParsingWithDozeMode(LIGHT_DOZE, false /* featureEnabled */);
+    }
+
+    private void doTestTcpInfoDisableParsingWithDozeMode(@DozeModeType int dozeModeType,
+            boolean featureEnabled) throws Exception {
         final TcpSocketTracker tst = new TcpSocketTracker(mDependencies, mNetwork);
+        tst.setNetworkCapabilities(CELL_NOT_METERED_CAPABILITIES);
         final ArgumentCaptor<BroadcastReceiver> receiverCaptor =
                 ArgumentCaptor.forClass(BroadcastReceiver.class);
 
-        // Enable light doze mode with 1 netlink message.
-        verify(mDependencies).addDeviceIdleReceiver(receiverCaptor.capture(), anyBoolean());
+        // Enable doze mode with 1 netlink message.
+        verify(mDependencies).addDeviceIdleReceiver(receiverCaptor.capture(),
+                anyBoolean(), anyBoolean());
         final BroadcastReceiver receiver = receiverCaptor.getValue();
-        doReturn(true).when(mPowerManager).isDeviceLightIdleMode();
-        receiver.onReceive(mContext, new Intent(ACTION_DEVICE_LIGHT_IDLE_MODE_CHANGED));
+        if (dozeModeType == DEEP_DOZE) {
+            doReturn(true).when(mPowerManager).isDeviceIdleMode();
+            receiver.onReceive(mContext, new Intent(ACTION_DEVICE_IDLE_MODE_CHANGED));
+        } else {
+            doReturn(true).when(mPowerManager).isDeviceLightIdleMode();
+            receiver.onReceive(mContext, new Intent(ACTION_DEVICE_LIGHT_IDLE_MODE_CHANGED));
+        }
         doReturn(getByteBufferFromHexString(composeSockDiagTcpHex(9, 10)
                 + NLMSG_DONE_HEX)).when(mDependencies).recvMessage(any());
 
+        if (!featureEnabled) {
+            // Verify TcpInfo is still processed.
+            assertTrue(tst.pollSocketsInfo());
+            assertEquals(10, tst.getSentSinceLastRecv());
+            // Lost 4 + default 5 retrans / 10 sent.
+            assertEquals(90, tst.getLatestPacketFailPercentage());
+            assertTrue(tst.isDataStallSuspected());
+            return;
+        }
+
         // Verify counters are not updated.
         assertFalse(tst.pollSocketsInfo());
         assertEquals(0, tst.getSentSinceLastRecv());
@@ -604,9 +752,14 @@
         assertEquals(-1, tst.getLatestPacketFailPercentage());
         assertFalse(tst.isDataStallSuspected());
 
-        // Disable light doze mode, verify polling are processed and counters are updated.
-        doReturn(false).when(mPowerManager).isDeviceLightIdleMode();
-        receiver.onReceive(mContext, new Intent(ACTION_DEVICE_LIGHT_IDLE_MODE_CHANGED));
+        // Disable deep/light doze mode, verify polling are processed and counters are updated.
+        if (dozeModeType == DEEP_DOZE) {
+            doReturn(false).when(mPowerManager).isDeviceIdleMode();
+            receiver.onReceive(mContext, new Intent(ACTION_DEVICE_IDLE_MODE_CHANGED));
+        } else {
+            doReturn(false).when(mPowerManager).isDeviceLightIdleMode();
+            receiver.onReceive(mContext, new Intent(ACTION_DEVICE_LIGHT_IDLE_MODE_CHANGED));
+        }
         assertTrue(tst.pollSocketsInfo());
         assertEquals(10, tst.getSentSinceLastRecv());
         // Lost 4 + default 5 retrans / 10 sent.
@@ -614,28 +767,6 @@
         assertTrue(tst.isDataStallSuspected());
     }
 
-    @Test @IgnoreUpTo(Build.VERSION_CODES.S_V2)
-    public void testTcpInfoDisableParsingWithLightDozeMode_disabled() throws Exception {
-        when(mDependencies.shouldDisableInLightDoze()).thenReturn(false);
-        final TcpSocketTracker tst = new TcpSocketTracker(mDependencies, mNetwork);
-        final ArgumentCaptor<BroadcastReceiver> receiverCaptor =
-                ArgumentCaptor.forClass(BroadcastReceiver.class);
-
-        // Enable light doze mode with 1 netlink message.
-        verify(mDependencies).addDeviceIdleReceiver(receiverCaptor.capture(), anyBoolean());
-        final BroadcastReceiver receiver = receiverCaptor.getValue();
-        doReturn(true).when(mPowerManager).isDeviceLightIdleMode();
-        receiver.onReceive(mContext, new Intent(ACTION_DEVICE_LIGHT_IDLE_MODE_CHANGED));
-        doReturn(getByteBufferFromHexString(composeSockDiagTcpHex(9, 10)
-                + NLMSG_DONE_HEX)).when(mDependencies).recvMessage(any());
-
-        // Verify TcpInfo is still processed.
-        assertTrue(tst.pollSocketsInfo());
-        assertEquals(10, tst.getSentSinceLastRecv());
-        assertEquals(90, tst.getLatestPacketFailPercentage());
-        assertTrue(tst.isDataStallSuspected());
-    }
-
     private void setupNormalTestTcpInfo() throws Exception {
         final ByteBuffer tcpBufferV6 = getByteBuffer(TEST_RESPONSE_BYTES);
         final ByteBuffer tcpBufferV4 = getByteBuffer(TEST_RESPONSE_BYTES);
diff --git a/tests/unit/src/com/android/server/connectivity/NetworkMonitorTest.java b/tests/unit/src/com/android/server/connectivity/NetworkMonitorTest.java
index b29ea16..5be2573 100644
--- a/tests/unit/src/com/android/server/connectivity/NetworkMonitorTest.java
+++ b/tests/unit/src/com/android/server/connectivity/NetworkMonitorTest.java
@@ -2345,6 +2345,19 @@
     }
 
     @Test
+    public void testTcpSocketTracker_setCapabilities() throws Exception {
+        setDataStallEvaluationType(DATA_STALL_EVALUATION_TYPE_TCP);
+        final InOrder inOrder = inOrder(mTst);
+        final WrappedNetworkMonitor wnm = prepareValidatedStateNetworkMonitor(
+                CELL_METERED_CAPABILITIES);
+        inOrder.verify(mTst).setNetworkCapabilities(eq(CELL_METERED_CAPABILITIES));
+
+        // Suspend the network. Verify the capabilities would be passed to TcpSocketTracker.
+        setNetworkCapabilities(wnm, CELL_SUSPENDED_METERED_CAPABILITIES);
+        inOrder.verify(mTst).setNetworkCapabilities(eq(CELL_SUSPENDED_METERED_CAPABILITIES));
+    }
+
+    @Test
     public void testDataStall_setOpportunisticMode() {
         setDataStallEvaluationType(DATA_STALL_EVALUATION_TYPE_TCP);
         WrappedNetworkMonitor wnm = makeCellNotMeteredNetworkMonitor();