[automerger skipped] [SP28.1] Add onSetWarningAndLimit to test utilities am: 952e18da40 am: c8b2f91ce8 am: dede9c3455 -s ours

am skip reason: skip tag Change-Id I4cc242e5fba895bb0756733bafc5be1f39b10ac0 with SHA-1 b4d6a7152f is already in history

Original change: https://android-review.googlesource.com/c/platform/frameworks/libs/net/+/1469182

Change-Id: I0919439a464ad3bc7f2927138b7482b8f90ed706
diff --git a/staticlibs/Android.bp b/staticlibs/Android.bp
index 03296e7..a341b0f 100644
--- a/staticlibs/Android.bp
+++ b/staticlibs/Android.bp
@@ -121,7 +121,10 @@
     name: "net-utils-framework-common-srcs",
     srcs: ["framework/**/*.java"],
     path: "framework",
-    visibility: ["//frameworks/base"],
+    visibility: [
+        "//frameworks/base",
+        "//frameworks/base/packages/Connectivity/framework",
+    ],
 }
 
 java_library {
diff --git a/staticlibs/device/android/net/NetworkFactory.java b/staticlibs/device/android/net/NetworkFactory.java
index bcc6089..68db57b 100644
--- a/staticlibs/device/android/net/NetworkFactory.java
+++ b/staticlibs/device/android/net/NetworkFactory.java
@@ -17,7 +17,9 @@
 package android.net;
 
 import android.annotation.NonNull;
+import android.annotation.Nullable;
 import android.content.Context;
+import android.net.NetworkProvider.NetworkOfferCallback;
 import android.os.Handler;
 import android.os.Looper;
 import android.os.Message;
@@ -29,6 +31,7 @@
 import java.io.PrintWriter;
 import java.util.LinkedHashMap;
 import java.util.Map;
+import java.util.concurrent.Executor;
 
 /**
  * A NetworkFactory is an entity that creates NetworkAgent objects.
@@ -48,6 +51,14 @@
 public class NetworkFactory extends Handler {
     private static final boolean DBG = true;
     private static final boolean VDBG = false;
+
+    // A score that will win against everything, so that score filtering will let all requests
+    // through
+    // TODO : remove this and replace with an API to listen to all requests.
+    @NonNull
+    private static final NetworkScore INVINCIBLE_SCORE =
+            new NetworkScore.Builder().setLegacyInt(1000).build();
+
     /**
      * Pass a network request to the bearer.  If the bearer believes it can
      * satisfy the request it should connect to the network and create a
@@ -59,19 +70,9 @@
      * disregard any that it will never be able to service, for example
      * those requiring a different bearer.
      * msg.obj = NetworkRequest
-     * msg.arg1 = score - the score of the network currently satisfying this
-     *            request.  If this bearer knows in advance it cannot
-     *            exceed this score it should not try to connect, holding the request
-     *            for the future.
-     *            Note that subsequent events may give a different (lower
-     *            or higher) score for this request, transmitted to each
-     *            NetworkFactory through additional CMD_REQUEST_NETWORK msgs
-     *            with the same NetworkRequest but an updated score.
-     *            Also, network conditions may change for this bearer
-     *            allowing for a better score in the future.
-     * msg.arg2 = the ID of the NetworkProvider currently responsible for the
-     *            NetworkAgent handling this request, or NetworkProvider.ID_NONE if none.
      */
+    // TODO : this and CANCEL_REQUEST are only used by telephony tests. Replace it in the tests
+    // and remove them and the associated code.
     public static final int CMD_REQUEST_NETWORK = 1;
 
     /**
@@ -82,7 +83,7 @@
 
     /**
      * Internally used to set our best-guess score.
-     * msg.arg1 = new score
+     * msg.obj = new score
      */
     private static final int CMD_SET_SCORE = 3;
 
@@ -93,28 +94,72 @@
      */
     private static final int CMD_SET_FILTER = 4;
 
+    /**
+     * Internally used to send the network offer associated with this factory.
+     * No arguments, will read from members
+     */
+    private static final int CMD_OFFER_NETWORK = 5;
+
+    /**
+     * Internally used to send the request to listen to all requests.
+     * No arguments, will read from members
+     */
+    private static final int CMD_LISTEN_TO_ALL_REQUESTS = 6;
+
     private final Context mContext;
     private final String LOG_TAG;
 
     private final Map<NetworkRequest, NetworkRequestInfo> mNetworkRequests =
             new LinkedHashMap<>();
 
-    private int mScore;
-    private NetworkCapabilities mCapabilityFilter;
+    @NonNull private NetworkScore mScore = new NetworkScore.Builder().setLegacyInt(0).build();
+    @NonNull private NetworkCapabilities mCapabilityFilter;
 
     private int mRefCount = 0;
     private NetworkProvider mProvider = null;
 
+    private final NetworkOfferCallback mRequestCallback = new NetworkOfferCallback() {
+        @Override
+        public void onNetworkNeeded(@NonNull final NetworkRequest request) {
+            handleAddRequest(request);
+        }
+
+        @Override
+        public void onNetworkUnneeded(@NonNull final NetworkRequest request) {
+            handleRemoveRequest(request);
+        }
+    };
+    @NonNull private final Executor mExecutor = command -> post(command);
+
+    // Ideally the filter argument would be non-null, but null has historically meant to see
+    // no requests and telephony passes null.
     public NetworkFactory(Looper looper, Context context, String logTag,
-            NetworkCapabilities filter) {
+            @Nullable final NetworkCapabilities filter) {
         super(looper);
         LOG_TAG = logTag;
         mContext = context;
-        mCapabilityFilter = filter;
+        if (null != filter) {
+            mCapabilityFilter = filter;
+        } else {
+            mCapabilityFilter = new NetworkCapabilities.Builder().clearAll().build();
+        }
     }
 
     /* Registers this NetworkFactory with the system. May only be called once per factory. */
     public void register() {
+        register(false);
+    }
+
+    /**
+     * Registers this NetworkFactory with the system ignoring the score filter. This will let
+     * the factory always see all network requests matching its capabilities filter.
+     * May only be called once per factory.
+     */
+    public void registerIgnoringScore() {
+        register(true);
+    }
+
+    private void register(final boolean listenToAllRequests) {
         if (mProvider != null) {
             throw new IllegalStateException("A NetworkFactory must only be registered once");
         }
@@ -124,7 +169,7 @@
             @Override
             public void onNetworkRequested(@NonNull NetworkRequest request, int score,
                     int servingProviderId) {
-                handleAddRequest(request, score, servingProviderId);
+                handleAddRequest(request);
             }
 
             @Override
@@ -134,7 +179,19 @@
         };
 
         ((ConnectivityManager) mContext.getSystemService(
-            Context.CONNECTIVITY_SERVICE)).registerNetworkProvider(mProvider);
+                Context.CONNECTIVITY_SERVICE)).registerNetworkProvider(mProvider);
+
+        // The mScore and mCapabilityFilter members can only be accessed on the handler thread.
+        // TODO : offer a separate API to listen to all requests instead
+        if (listenToAllRequests) {
+            sendMessage(obtainMessage(CMD_LISTEN_TO_ALL_REQUESTS));
+        } else {
+            sendMessage(obtainMessage(CMD_OFFER_NETWORK));
+        }
+    }
+
+    private void handleOfferNetwork(@NonNull final NetworkScore score) {
+        mProvider.registerNetworkOffer(score, mCapabilityFilter, mExecutor, mRequestCallback);
     }
 
     /** Unregisters this NetworkFactory. After this call, the object can no longer be used. */
@@ -156,7 +213,7 @@
     public void handleMessage(Message msg) {
         switch (msg.what) {
             case CMD_REQUEST_NETWORK: {
-                handleAddRequest((NetworkRequest) msg.obj, msg.arg1, msg.arg2);
+                handleAddRequest((NetworkRequest) msg.obj);
                 break;
             }
             case CMD_CANCEL_REQUEST: {
@@ -164,32 +221,36 @@
                 break;
             }
             case CMD_SET_SCORE: {
-                handleSetScore(msg.arg1);
+                handleSetScore((NetworkScore) msg.obj);
                 break;
             }
             case CMD_SET_FILTER: {
                 handleSetFilter((NetworkCapabilities) msg.obj);
                 break;
             }
+            case CMD_OFFER_NETWORK: {
+                handleOfferNetwork(mScore);
+                break;
+            }
+            case CMD_LISTEN_TO_ALL_REQUESTS: {
+                handleOfferNetwork(INVINCIBLE_SCORE);
+                break;
+            }
         }
     }
 
     private static class NetworkRequestInfo {
-        public final NetworkRequest request;
-        public int score;
+        @NonNull public final NetworkRequest request;
         public boolean requested; // do we have a request outstanding, limited by score
-        public int providerId;
 
-        NetworkRequestInfo(NetworkRequest request, int score, int providerId) {
+        NetworkRequestInfo(@NonNull final NetworkRequest request) {
             this.request = request;
-            this.score = score;
             this.requested = false;
-            this.providerId = providerId;
         }
 
         @Override
         public String toString() {
-            return "{" + request + ", score=" + score + ", requested=" + requested + "}";
+            return "{" + request + ", requested=" + requested + "}";
         }
     }
 
@@ -198,35 +259,27 @@
      * @see #CMD_REQUEST_NETWORK
      *
      * @param request the request to handle.
-     * @param score the score of the NetworkAgent currently satisfying this request.
-     * @param servingProviderId the ID of the NetworkProvider that created the NetworkAgent
-     *        currently satisfying this request.
      */
-    @VisibleForTesting
-    protected void handleAddRequest(NetworkRequest request, int score, int servingProviderId) {
+    private void handleAddRequest(@NonNull final NetworkRequest request) {
         NetworkRequestInfo n = mNetworkRequests.get(request);
         if (n == null) {
-            if (DBG) {
-                log("got request " + request + " with score " + score
-                        + " and providerId " + servingProviderId);
-            }
-            n = new NetworkRequestInfo(request, score, servingProviderId);
+            if (DBG) log("got request " + request);
+            n = new NetworkRequestInfo(request);
             mNetworkRequests.put(n.request, n);
         } else {
             if (VDBG) {
-                log("new score " + score + " for existing request " + request
-                        + " and providerId " + servingProviderId);
+                log("handle existing request " + request);
             }
-            n.score = score;
-            n.providerId = servingProviderId;
         }
         if (VDBG) log("  my score=" + mScore + ", my filter=" + mCapabilityFilter);
 
-        evalRequest(n);
+        if (acceptRequest(request)) {
+            n.requested = true;
+            needNetworkFor(request);
+        }
     }
 
-    @VisibleForTesting
-    protected void handleRemoveRequest(NetworkRequest request) {
+    private void handleRemoveRequest(NetworkRequest request) {
         NetworkRequestInfo n = mNetworkRequests.get(request);
         if (n != null) {
             mNetworkRequests.remove(request);
@@ -234,20 +287,21 @@
         }
     }
 
-    private void handleSetScore(int score) {
+    private void handleSetScore(@NonNull final NetworkScore score) {
+        if (mScore.equals(score)) return;
         mScore = score;
-        evalRequests();
+        reevaluateAllRequests();
     }
 
-    private void handleSetFilter(NetworkCapabilities netCap) {
+    private void handleSetFilter(@NonNull final NetworkCapabilities netCap) {
+        if (netCap.equals(mCapabilityFilter)) return;
         mCapabilityFilter = netCap;
-        evalRequests();
+        reevaluateAllRequests();
     }
 
-    /** @deprecated None of the implementors use the score, remove this method */
-    @Deprecated
-    public boolean acceptRequest(NetworkRequest request, int score) {
-        return acceptRequest(request);
+    protected final void reevaluateAllRequests() {
+        if (mProvider == null) return;
+        mProvider.registerNetworkOffer(mScore, mCapabilityFilter, mExecutor, mRequestCallback);
     }
 
     /**
@@ -269,75 +323,9 @@
      *
      * @return {@code true} to accept the request.
      */
-    public boolean acceptRequest(NetworkRequest request) {
+    public boolean acceptRequest(@NonNull final NetworkRequest request) {
         return true;
     }
-
-    private void evalRequest(NetworkRequestInfo n) {
-        if (VDBG) {
-            log("evalRequest");
-            log(" n.requests = " + n.requested);
-            log(" n.score = " + n.score);
-            log(" mScore = " + mScore);
-            log(" request.providerId = " + n.providerId);
-            log(" mProvider.id = " + mProvider.getProviderId());
-        }
-        if (shouldNeedNetworkFor(n)) {
-            if (VDBG) log("  needNetworkFor");
-            needNetworkFor(n.request, n.score);
-            n.requested = true;
-        } else if (shouldReleaseNetworkFor(n)) {
-            if (VDBG) log("  releaseNetworkFor");
-            releaseNetworkFor(n.request);
-            n.requested = false;
-        } else {
-            if (VDBG) log("  done");
-        }
-    }
-
-    private boolean shouldNeedNetworkFor(NetworkRequestInfo n) {
-        // If this request is already tracked, it doesn't qualify for need
-        return !n.requested
-            // If the score of this request is higher or equal to that of this factory and some
-            // other factory is responsible for it, then this factory should not track the request
-            // because it has no hope of satisfying it.
-            && (n.score < mScore || n.providerId == mProvider.getProviderId())
-            // If this factory can't satisfy the capability needs of this request, then it
-            // should not be tracked.
-            && n.request.canBeSatisfiedBy(mCapabilityFilter)
-            // Finally if the concrete implementation of the factory rejects the request, then
-            // don't track it.
-            && acceptRequest(n.request, n.score);
-    }
-
-    private boolean shouldReleaseNetworkFor(NetworkRequestInfo n) {
-        // Don't release a request that's not tracked.
-        return n.requested
-            // The request should be released if it can't be satisfied by this factory. That
-            // means either of the following conditions are met :
-            // - Its score is too high to be satisfied by this factory and it's not already
-            //   assigned to the factory
-            // - This factory can't satisfy the capability needs of the request
-            // - The concrete implementation of the factory rejects the request
-            && ((n.score > mScore && n.providerId != mProvider.getProviderId())
-                    || !n.request.canBeSatisfiedBy(mCapabilityFilter)
-                    || !acceptRequest(n.request, n.score));
-    }
-
-    private void evalRequests() {
-        for (NetworkRequestInfo n : mNetworkRequests.values()) {
-            evalRequest(n);
-        }
-    }
-
-    /**
-     * Post a command, on this NetworkFactory Handler, to re-evaluate all
-     * outstanding requests. Can be called from a factory implementation.
-     */
-    protected void reevaluateAllRequests() {
-        post(this::evalRequests);
-    }
-
     /**
      * Can be called by a factory to release a request as unfulfillable: the request will be
      * removed, and the caller will get a
@@ -363,23 +351,36 @@
     protected void startNetwork() { }
     protected void stopNetwork() { }
 
-    /** @deprecated none of the implementors use the score : migrate them */
-    @Deprecated
-    protected void needNetworkFor(NetworkRequest networkRequest, int score) {
-        needNetworkFor(networkRequest);
-    }
-
     // override to do fancier stuff
-    protected void needNetworkFor(NetworkRequest networkRequest) {
+    protected void needNetworkFor(@NonNull final NetworkRequest networkRequest) {
         if (++mRefCount == 1) startNetwork();
     }
 
-    protected void releaseNetworkFor(NetworkRequest networkRequest) {
+    protected void releaseNetworkFor(@NonNull final NetworkRequest networkRequest) {
         if (--mRefCount == 0) stopNetwork();
     }
 
-    public void setScoreFilter(int score) {
-        sendMessage(obtainMessage(CMD_SET_SCORE, score, 0));
+    /**
+     * @deprecated this method was never part of the API (system or public) and is only added
+     *   for migration of existing clients. It will be removed before S ships.
+     */
+    @Deprecated
+    public void setScoreFilter(final int score) {
+        setScoreFilter(new NetworkScore.Builder().setLegacyInt(score).build());
+    }
+
+    /**
+     * Set a score filter for this factory.
+     *
+     * This should include the transports the factory knows its networks will have, and
+     * an optimistic view of the attributes it may have. This does not commit the factory
+     * to being able to bring up such a network ; it only lets it avoid hearing about
+     * requests that it has no chance of fulfilling.
+     *
+     * @param score the filter
+     */
+    public void setScoreFilter(@NonNull final NetworkScore score) {
+        sendMessage(obtainMessage(CMD_SET_SCORE, score));
     }
 
     public void setCapabilityFilter(NetworkCapabilities netCap) {
diff --git a/staticlibs/framework/com/android/net/module/util/CollectionUtils.java b/staticlibs/framework/com/android/net/module/util/CollectionUtils.java
index 2223443..44d0a6e 100644
--- a/staticlibs/framework/com/android/net/module/util/CollectionUtils.java
+++ b/staticlibs/framework/com/android/net/module/util/CollectionUtils.java
@@ -26,6 +26,7 @@
 
 /**
  * Utilities for {@link Collection} and arrays.
+ * @hide
  */
 public final class CollectionUtils {
     private CollectionUtils() {}
diff --git a/staticlibs/framework/com/android/net/module/util/ConnectivityUtils.java b/staticlibs/framework/com/android/net/module/util/ConnectivityUtils.java
index 382912b..c135e46 100644
--- a/staticlibs/framework/com/android/net/module/util/ConnectivityUtils.java
+++ b/staticlibs/framework/com/android/net/module/util/ConnectivityUtils.java
@@ -24,6 +24,7 @@
 
 /**
  * Various utilities used in connectivity code.
+ * @hide
  */
 public final class ConnectivityUtils {
     private ConnectivityUtils() {}
diff --git a/staticlibs/framework/com/android/net/module/util/NetworkCapabilitiesUtils.java b/staticlibs/framework/com/android/net/module/util/NetworkCapabilitiesUtils.java
index 3de78c6..e51c5f2 100644
--- a/staticlibs/framework/com/android/net/module/util/NetworkCapabilitiesUtils.java
+++ b/staticlibs/framework/com/android/net/module/util/NetworkCapabilitiesUtils.java
@@ -16,6 +16,25 @@
 
 package com.android.net.module.util;
 
+import static android.net.NetworkCapabilities.NET_CAPABILITY_BIP;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_CBS;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_DUN;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_EIMS;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_ENTERPRISE;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_FOTA;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_IA;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_IMS;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_INTERNET;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_MCX;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_MMS;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_OEM_PAID;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_OEM_PRIVATE;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_RCS;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_SUPL;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_VEHICLE_INTERNAL;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_VSIM;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_WIFI_P2P;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_XCAP;
 import static android.net.NetworkCapabilities.TRANSPORT_BLUETOOTH;
 import static android.net.NetworkCapabilities.TRANSPORT_CELLULAR;
 import static android.net.NetworkCapabilities.TRANSPORT_ETHERNET;
@@ -24,10 +43,13 @@
 import static android.net.NetworkCapabilities.TRANSPORT_WIFI_AWARE;
 
 import android.annotation.NonNull;
+import android.net.NetworkCapabilities;
 
+import com.android.internal.annotations.VisibleForTesting;
 
 /**
  * Utilities to examine {@link android.net.NetworkCapabilities}.
+ * @hide
  */
 public final class NetworkCapabilitiesUtils {
     // Transports considered to classify networks in UI, in order of which transport should be
@@ -55,6 +77,46 @@
     };
 
     /**
+     * Capabilities that suggest that a network is restricted.
+     * See {@code NetworkCapabilities#maybeMarkCapabilitiesRestricted},
+      * and {@code FORCE_RESTRICTED_CAPABILITIES}.
+     */
+    @VisibleForTesting
+    static final long RESTRICTED_CAPABILITIES =
+            (1 << NET_CAPABILITY_BIP)
+            | (1 << NET_CAPABILITY_CBS)
+            | (1 << NET_CAPABILITY_DUN)
+            | (1 << NET_CAPABILITY_EIMS)
+            | (1 << NET_CAPABILITY_ENTERPRISE)
+            | (1 << NET_CAPABILITY_FOTA)
+            | (1 << NET_CAPABILITY_IA)
+            | (1 << NET_CAPABILITY_IMS)
+            | (1 << NET_CAPABILITY_MCX)
+            | (1 << NET_CAPABILITY_RCS)
+            | (1 << NET_CAPABILITY_VEHICLE_INTERNAL)
+            | (1 << NET_CAPABILITY_VSIM)
+            | (1 << NET_CAPABILITY_XCAP);
+
+    /**
+     * Capabilities that force network to be restricted.
+     * See {@code NetworkCapabilities#maybeMarkCapabilitiesRestricted}.
+     */
+    private static final long FORCE_RESTRICTED_CAPABILITIES =
+            (1 << NET_CAPABILITY_OEM_PAID)
+            | (1 << NET_CAPABILITY_OEM_PRIVATE);
+
+    /**
+     * Capabilities that suggest that a network is unrestricted.
+     * See {@code NetworkCapabilities#maybeMarkCapabilitiesRestricted}.
+     */
+    @VisibleForTesting
+    static final long UNRESTRICTED_CAPABILITIES =
+            (1 << NET_CAPABILITY_INTERNET)
+            | (1 << NET_CAPABILITY_MMS)
+            | (1 << NET_CAPABILITY_SUPL)
+            | (1 << NET_CAPABILITY_WIFI_P2P);
+
+    /**
      * Get a transport that can be used to classify a network when displaying its info to users.
      *
      * While networks can have multiple transports, users generally think of them as "wifi",
@@ -79,6 +141,33 @@
         return transports[0];
     }
 
+
+    /**
+     * Infers that all the capabilities it provides are typically provided by restricted networks
+     * or not.
+     *
+     * @param nc the {@link NetworkCapabilities} to infer the restricted capabilities.
+     *
+     * @return {@code true} if the network should be restricted.
+     */
+    public static boolean inferRestrictedCapability(NetworkCapabilities nc) {
+        final long capabilities = packBits(nc.getCapabilities());
+        // Check if we have any capability that forces the network to be restricted.
+        final boolean forceRestrictedCapability =
+                (capabilities & FORCE_RESTRICTED_CAPABILITIES) != 0;
+
+        // Verify there aren't any unrestricted capabilities.  If there are we say
+        // the whole thing is unrestricted unless it is forced to be restricted.
+        final boolean hasUnrestrictedCapabilities =
+                (capabilities & UNRESTRICTED_CAPABILITIES) != 0;
+
+        // Must have at least some restricted capabilities.
+        final boolean hasRestrictedCapabilities = (capabilities & RESTRICTED_CAPABILITIES) != 0;
+
+        return forceRestrictedCapability
+                || (hasRestrictedCapabilities && !hasUnrestrictedCapabilities);
+    }
+
     /**
      * Unpacks long value into an array of bits.
      */
diff --git a/staticlibs/framework/com/android/net/module/util/NetworkIdentityUtils.java b/staticlibs/framework/com/android/net/module/util/NetworkIdentityUtils.java
index 94e6017..b641753 100644
--- a/staticlibs/framework/com/android/net/module/util/NetworkIdentityUtils.java
+++ b/staticlibs/framework/com/android/net/module/util/NetworkIdentityUtils.java
@@ -21,6 +21,7 @@
 
 /**
  * Utilities to examine {@link android.net.NetworkIdentity}.
+ * @hide
  */
 public class NetworkIdentityUtils {
     /**
diff --git a/staticlibs/framework/com/android/net/module/util/NetworkStackConstants.java b/staticlibs/framework/com/android/net/module/util/NetworkStackConstants.java
index 499297c..b7062e7 100644
--- a/staticlibs/framework/com/android/net/module/util/NetworkStackConstants.java
+++ b/staticlibs/framework/com/android/net/module/util/NetworkStackConstants.java
@@ -186,6 +186,27 @@
      */
     public static final int VENDOR_SPECIFIC_IE_ID = 0xdd;
 
+
+    /**
+     * TrafficStats constants.
+     */
+    // These tags are used by the network stack to do traffic for its own purposes. Traffic
+    // tagged with these will be counted toward the network stack and must stay inside the
+    // range defined by
+    // {@link android.net.TrafficStats#TAG_NETWORK_STACK_RANGE_START} and
+    // {@link android.net.TrafficStats#TAG_NETWORK_STACK_RANGE_END}.
+    public static final int TAG_SYSTEM_DHCP = 0xFFFFFE01;
+    public static final int TAG_SYSTEM_NEIGHBOR = 0xFFFFFE02;
+    public static final int TAG_SYSTEM_DHCP_SERVER = 0xFFFFFE03;
+
+    // These tags are used by the network stack to do traffic on behalf of apps. Traffic
+    // tagged with these will be counted toward the app on behalf of which the network
+    // stack is doing this traffic. These values must stay inside the range defined by
+    // {@link android.net.TrafficStats#TAG_NETWORK_STACK_IMPERSONATION_RANGE_START} and
+    // {@link android.net.TrafficStats#TAG_NETWORK_STACK_IMPERSONATION_RANGE_END}.
+    public static final int TAG_SYSTEM_PROBE = 0xFFFFFF81;
+    public static final int TAG_SYSTEM_DNS = 0xFFFFFF82;
+
     // TODO: Move to Inet4AddressUtils
     // See aosp/1455936: NetworkStackConstants can't depend on it as it causes jarjar-related issues
     // for users of both the net-utils-device-common and net-utils-framework-common libraries.
diff --git a/staticlibs/tests/unit/src/com/android/net/module/util/NetworkCapabilitiesUtilsTest.kt b/staticlibs/tests/unit/src/com/android/net/module/util/NetworkCapabilitiesUtilsTest.kt
index df2f459..5f15c6a 100644
--- a/staticlibs/tests/unit/src/com/android/net/module/util/NetworkCapabilitiesUtilsTest.kt
+++ b/staticlibs/tests/unit/src/com/android/net/module/util/NetworkCapabilitiesUtilsTest.kt
@@ -16,6 +16,12 @@
 
 package com.android.net.module.util
 
+import android.net.NetworkCapabilities
+import android.net.NetworkCapabilities.NET_CAPABILITY_CBS
+import android.net.NetworkCapabilities.NET_CAPABILITY_EIMS
+import android.net.NetworkCapabilities.NET_CAPABILITY_INTERNET
+import android.net.NetworkCapabilities.NET_CAPABILITY_INTERNET
+import android.net.NetworkCapabilities.NET_CAPABILITY_OEM_PAID
 import android.net.NetworkCapabilities.TRANSPORT_BLUETOOTH
 import android.net.NetworkCapabilities.TRANSPORT_CELLULAR
 import android.net.NetworkCapabilities.TRANSPORT_ETHERNET
@@ -25,6 +31,8 @@
 import android.net.NetworkCapabilities.TRANSPORT_WIFI_AWARE
 import androidx.test.filters.SmallTest
 import androidx.test.runner.AndroidJUnit4
+import com.android.net.module.util.NetworkCapabilitiesUtils.RESTRICTED_CAPABILITIES
+import com.android.net.module.util.NetworkCapabilitiesUtils.UNRESTRICTED_CAPABILITIES
 import com.android.net.module.util.NetworkCapabilitiesUtils.getDisplayTransport
 import com.android.net.module.util.NetworkCapabilitiesUtils.packBits
 import com.android.net.module.util.NetworkCapabilitiesUtils.unpackBits
@@ -33,6 +41,7 @@
 import java.lang.IllegalArgumentException
 import kotlin.test.assertEquals
 import kotlin.test.assertFailsWith
+import kotlin.test.assertFalse
 import kotlin.test.assertTrue
 
 @RunWith(AndroidJUnit4::class)
@@ -78,4 +87,41 @@
         assertEquals(packedBits, packBits(bits))
         assertTrue(bits contentEquals unpackBits(packedBits))
     }
+
+    @Test
+    fun testInferRestrictedCapability() {
+        val nc = NetworkCapabilities()
+        // Default capabilities don't have restricted capability.
+        assertFalse(NetworkCapabilitiesUtils.inferRestrictedCapability(nc))
+        // If there is a force restricted capability, then the network capabilities is restricted.
+        nc.addCapability(NET_CAPABILITY_OEM_PAID)
+        nc.addCapability(NET_CAPABILITY_INTERNET)
+        assertTrue(NetworkCapabilitiesUtils.inferRestrictedCapability(nc))
+        // Except for the force restricted capability, if there is any unrestricted capability in
+        // capabilities, then the network capabilities is not restricted.
+        nc.removeCapability(NET_CAPABILITY_OEM_PAID)
+        nc.addCapability(NET_CAPABILITY_CBS)
+        assertFalse(NetworkCapabilitiesUtils.inferRestrictedCapability(nc))
+        // Except for the force restricted capability, the network capabilities will only be treated
+        // as restricted when there is no any unrestricted capability.
+        nc.removeCapability(NET_CAPABILITY_INTERNET)
+        assertTrue(NetworkCapabilitiesUtils.inferRestrictedCapability(nc))
+    }
+
+    @Test
+    fun testRestrictedUnrestrictedCapabilities() {
+        // verify EIMS is restricted
+        assertEquals((1 shl NET_CAPABILITY_EIMS).toLong() and RESTRICTED_CAPABILITIES,
+                (1 shl NET_CAPABILITY_EIMS).toLong())
+
+        // verify CBS is also restricted
+        assertEquals((1 shl NET_CAPABILITY_CBS).toLong() and RESTRICTED_CAPABILITIES,
+                (1 shl NET_CAPABILITY_CBS).toLong())
+
+        // verify default is not restricted
+        assertEquals((1 shl NET_CAPABILITY_INTERNET).toLong() and RESTRICTED_CAPABILITIES, 0)
+
+        // just to see
+        assertEquals(RESTRICTED_CAPABILITIES and UNRESTRICTED_CAPABILITIES, 0)
+    }
 }