Prepare some ConnectivityService logic for fallback to Cellular change

Reduce the duplication of some logic so when falling back to Cellular
when WiFi fails to validate is enabled, there's less chance for bugs
and failures:
1. De-duplicate several Network vs NetworkRequest matching functions
2. Remove the very tricky nascent logic by adding a simple "lingering" bit.

Bug:20896761
Change-Id: I21da9e827eec9cfd6835fcaa650192b9186ed053

3 files changed, 144 insertions, 163 deletions

diff --git a/services/core/java/com/android/server/ConnectivityService.java b/services/core/java/com/android/server/ConnectivityService.java
index 7f124dc..6878caf 100644
--- a/services/core/java/com/android/server/ConnectivityService.java
+++ b/services/core/java/com/android/server/ConnectivityService.java
@@ -220,23 +220,13 @@ public class ConnectivityService extends IConnectivityManager.Stub
     private static final int ENABLED  = 1;
     private static final int DISABLED = 0;
 
-    // Arguments to rematchNetworkAndRequests()
-    private enum NascentState {
-        // Indicates a network was just validated for the first time.  If the network is found to
-        // be unwanted (i.e. not satisfy any NetworkRequests) it is torn down.
-        JUST_VALIDATED,
-        // Indicates a network was not validated for the first time immediately prior to this call.
-        NOT_JUST_VALIDATED
-    };
     private enum ReapUnvalidatedNetworks {
-        // Tear down unvalidated networks that have no chance (i.e. even if validated) of becoming
-        // the highest scoring network satisfying a NetworkRequest.  This should be passed when it's
-        // known that there may be unvalidated networks that could potentially be reaped, and when
+        // Tear down networks that have no chance (e.g. even if validated) of becoming
+        // the highest scoring network satisfying a NetworkRequest.  This should be passed when
         // all networks have been rematched against all NetworkRequests.
         REAP,
-        // Don't reap unvalidated networks.  This should be passed when it's known that there are
-        // no unvalidated networks that could potentially be reaped, and when some networks have
-        // not yet been rematched against all NetworkRequests.
+        // Don't reap networks.  This should be passed when some networks have not yet been
+        // rematched against all NetworkRequests.
         DONT_REAP
     };
 
@@ -873,7 +863,7 @@ public class ConnectivityService extends IConnectivityManager.Stub
         Network network = null;
         String subscriberId = null;
 
-        NetworkAgentInfo nai = mNetworkForRequestId.get(mDefaultRequest.requestId);
+        NetworkAgentInfo nai = getDefaultNetwork();
 
         final Network[] networks = getVpnUnderlyingNetworks(uid);
         if (networks != null) {
@@ -1778,7 +1768,7 @@ public class ConnectivityService extends IConnectivityManager.Stub
         pw.println();
         pw.println();
 
-        NetworkAgentInfo defaultNai = mNetworkForRequestId.get(mDefaultRequest.requestId);
+        final NetworkAgentInfo defaultNai = getDefaultNetwork();
         pw.print("Active default network: ");
         if (defaultNai == null) {
             pw.println("none");
@@ -1903,8 +1893,7 @@ public class ConnectivityService extends IConnectivityManager.Stub
                                 networkCapabilities.hasCapability(NET_CAPABILITY_VALIDATED)) {
                             Slog.wtf(TAG, "BUG: " + nai + " has stateful capability.");
                         }
-                        updateCapabilities(nai, networkCapabilities,
-                                NascentState.NOT_JUST_VALIDATED);
+                        updateCapabilities(nai, networkCapabilities);
                     }
                     break;
                 }
@@ -1994,11 +1983,9 @@ public class ConnectivityService extends IConnectivityManager.Stub
                         if (DBG) log(nai.name() + " validation " + (valid ? " passed" : "failed"));
                         if (valid != nai.lastValidated) {
                             final int oldScore = nai.getCurrentScore();
-                            final NascentState nascent = (valid && !nai.everValidated) ?
-                                    NascentState.JUST_VALIDATED : NascentState.NOT_JUST_VALIDATED;
                             nai.lastValidated = valid;
                             nai.everValidated |= valid;
-                            updateCapabilities(nai, nai.networkCapabilities, nascent);
+                            updateCapabilities(nai, nai.networkCapabilities);
                             // If score has changed, rebroadcast to NetworkFactories. b/17726566
                             if (oldScore != nai.getCurrentScore()) sendUpdatedScoreToFactories(nai);
                         }
@@ -2029,8 +2016,7 @@ public class ConnectivityService extends IConnectivityManager.Stub
                     if (nai != null && (visible != nai.lastCaptivePortalDetected)) {
                         nai.lastCaptivePortalDetected = visible;
                         nai.everCaptivePortalDetected |= visible;
-                        updateCapabilities(nai, nai.networkCapabilities,
-                                NascentState.NOT_JUST_VALIDATED);
+                        updateCapabilities(nai, nai.networkCapabilities);
                     }
                     if (!visible) {
                         setProvNotificationVisibleIntent(false, netId, null, 0, null, null, false);
@@ -2049,12 +2035,19 @@ public class ConnectivityService extends IConnectivityManager.Stub
         }
     }
 
+    private void linger(NetworkAgentInfo nai) {
+        nai.lingering = true;
+        nai.networkMonitor.sendMessage(NetworkMonitor.CMD_NETWORK_LINGER);
+        notifyNetworkCallbacks(nai, ConnectivityManager.CALLBACK_LOSING);
+    }
+
     // Cancel any lingering so the linger timeout doesn't teardown a network.
     // This should be called when a network begins satisfying a NetworkRequest.
     // Note: depending on what state the NetworkMonitor is in (e.g.,
     // if it's awaiting captive portal login, or if validation failed), this
     // may trigger a re-evaluation of the network.
     private void unlinger(NetworkAgentInfo nai) {
+        nai.lingering = false;
         if (VDBG) log("Canceling linger of " + nai.name());
         // If network has never been validated, it cannot have been lingered, so don't bother
         // needlessly triggering a re-evaluation.
@@ -2130,33 +2123,13 @@ public class ConnectivityService extends IConnectivityManager.Stub
                 // available until we've told netd to delete it below.
                 mNetworkForNetId.remove(nai.network.netId);
             }
-            // Since we've lost the network, go through all the requests that
-            // it was satisfying and see if any other factory can satisfy them.
-            // TODO: This logic may be better replaced with a call to rematchAllNetworksAndRequests
-            final ArrayList<NetworkAgentInfo> toActivate = new ArrayList<NetworkAgentInfo>();
+            // Remove all previously satisfied requests.
             for (int i = 0; i < nai.networkRequests.size(); i++) {
                 NetworkRequest request = nai.networkRequests.valueAt(i);
                 NetworkAgentInfo currentNetwork = mNetworkForRequestId.get(request.requestId);
                 if (currentNetwork != null && currentNetwork.network.netId == nai.network.netId) {
-                    if (DBG) {
-                        log("Checking for replacement network to handle request " + request );
-                    }
                     mNetworkForRequestId.remove(request.requestId);
                     sendUpdatedScoreToFactories(request, 0);
-                    NetworkAgentInfo alternative = null;
-                    for (NetworkAgentInfo existing : mNetworkAgentInfos.values()) {
-                        if (existing.satisfies(request) &&
-                                (alternative == null ||
-                                 alternative.getCurrentScore() < existing.getCurrentScore())) {
-                            alternative = existing;
-                        }
-                    }
-                    if (alternative != null) {
-                        if (DBG) log(" found replacement in " + alternative.name());
-                        if (!toActivate.contains(alternative)) {
-                            toActivate.add(alternative);
-                        }
-                    }
                 }
             }
             if (nai.networkRequests.get(mDefaultRequest.requestId) != null) {
@@ -2165,11 +2138,7 @@ public class ConnectivityService extends IConnectivityManager.Stub
                 requestNetworkTransitionWakelock(nai.name());
             }
             mLegacyTypeTracker.remove(nai, wasDefault);
-            for (NetworkAgentInfo networkToActivate : toActivate) {
-                unlinger(networkToActivate);
-                rematchNetworkAndRequests(networkToActivate, NascentState.NOT_JUST_VALIDATED,
-                        ReapUnvalidatedNetworks.DONT_REAP);
-            }
+            rematchAllNetworksAndRequests(null, 0);
             if (nai.created) {
                 // Tell netd to clean up the configuration for this network
                 // (routing rules, DNS, etc).
@@ -2223,49 +2192,9 @@ public class ConnectivityService extends IConnectivityManager.Stub
 
     private void handleRegisterNetworkRequest(NetworkRequestInfo nri) {
         mNetworkRequests.put(nri.request, nri);
-
-        // TODO: This logic may be better replaced with a call to rematchNetworkAndRequests
-
-        // Check for the best currently alive network that satisfies this request
-        NetworkAgentInfo bestNetwork = null;
-        for (NetworkAgentInfo network : mNetworkAgentInfos.values()) {
-            if (DBG) log("handleRegisterNetworkRequest checking " + network.name());
-            if (network.satisfies(nri.request)) {
-                if (DBG) log("apparently satisfied.  currentScore=" + network.getCurrentScore());
-                if (!nri.isRequest) {
-                    // Not setting bestNetwork here as a listening NetworkRequest may be
-                    // satisfied by multiple Networks.  Instead the request is added to
-                    // each satisfying Network and notified about each.
-                    if (!network.addRequest(nri.request)) {
-                        Slog.wtf(TAG, "BUG: " + network.name() + " already has " + nri.request);
-                    }
-                    notifyNetworkCallback(network, nri);
-                } else if (bestNetwork == null ||
-                        bestNetwork.getCurrentScore() < network.getCurrentScore()) {
-                    bestNetwork = network;
-                }
-            }
-        }
-        if (bestNetwork != null) {
-            if (DBG) log("using " + bestNetwork.name());
-            unlinger(bestNetwork);
-            if (!bestNetwork.addRequest(nri.request)) {
-                Slog.wtf(TAG, "BUG: " + bestNetwork.name() + " already has " + nri.request);
-            }
-            mNetworkForRequestId.put(nri.request.requestId, bestNetwork);
-            notifyNetworkCallback(bestNetwork, nri);
-            if (nri.request.legacyType != TYPE_NONE) {
-                mLegacyTypeTracker.add(nri.request.legacyType, bestNetwork);
-            }
-        }
-
-        if (nri.isRequest) {
-            if (DBG) log("sending new NetworkRequest to factories");
-            final int score = bestNetwork == null ? 0 : bestNetwork.getCurrentScore();
-            for (NetworkFactoryInfo nfi : mNetworkFactoryInfos.values()) {
-                nfi.asyncChannel.sendMessage(android.net.NetworkFactory.CMD_REQUEST_NETWORK, score,
-                        0, nri.request);
-            }
+        rematchAllNetworksAndRequests(null, 0);
+        if (nri.isRequest && mNetworkForRequestId.get(nri.request.requestId) == null) {
+            sendUpdatedScoreToFactories(nri.request, 0);
         }
     }
 
@@ -2282,7 +2211,7 @@ public class ConnectivityService extends IConnectivityManager.Stub
     // For unvalidated Networks this is whether it is satsifying any NetworkRequests or
     // were it to become validated, would it have a chance of satisfying any NetworkRequests.
     private boolean unneeded(NetworkAgentInfo nai) {
-        if (!nai.created || nai.isVPN()) return false;
+        if (!nai.created || nai.isVPN() || nai.lingering) return false;
         boolean unneeded = true;
         if (nai.everValidated) {
             for (int i = 0; i < nai.networkRequests.size() && unneeded; i++) {
@@ -2424,7 +2353,7 @@ public class ConnectivityService extends IConnectivityManager.Stub
         if (accept != nai.networkMisc.acceptUnvalidated) {
             int oldScore = nai.getCurrentScore();
             nai.networkMisc.acceptUnvalidated = accept;
-            rematchAllNetworksAndRequests(nai, oldScore, NascentState.NOT_JUST_VALIDATED);
+            rematchAllNetworksAndRequests(nai, oldScore);
             sendUpdatedScoreToFactories(nai);
         }
 
@@ -4033,7 +3962,7 @@ public class ConnectivityService extends IConnectivityManager.Stub
             } catch (Exception e) {
                 loge("Exception in setDnsServersForNetwork: " + e);
             }
-            NetworkAgentInfo defaultNai = mNetworkForRequestId.get(mDefaultRequest.requestId);
+            final NetworkAgentInfo defaultNai = getDefaultNetwork();
             if (defaultNai != null && defaultNai.network.netId == netId) {
                 setDefaultDnsSystemProperties(dnses);
             }
@@ -4070,11 +3999,9 @@ public class ConnectivityService extends IConnectivityManager.Stub
      *
      * @param networkAgent the network having its capabilities updated.
      * @param networkCapabilities the new network capabilities.
-     * @param nascent indicates whether {@code networkAgent} was validated
-     *         (i.e. had everValidated set for the first time) immediately prior to this call.
      */
     private void updateCapabilities(NetworkAgentInfo networkAgent,
-            NetworkCapabilities networkCapabilities, NascentState nascent) {
+            NetworkCapabilities networkCapabilities) {
         // Don't modify caller's NetworkCapabilities.
         networkCapabilities = new NetworkCapabilities(networkCapabilities);
         if (networkAgent.lastValidated) {
@@ -4091,7 +4018,7 @@ public class ConnectivityService extends IConnectivityManager.Stub
             synchronized (networkAgent) {
                 networkAgent.networkCapabilities = networkCapabilities;
             }
-            rematchAllNetworksAndRequests(networkAgent, networkAgent.getCurrentScore(), nascent);
+            rematchAllNetworksAndRequests(networkAgent, networkAgent.getCurrentScore());
             notifyNetworkCallbacks(networkAgent, ConnectivityManager.CALLBACK_CAP_CHANGED);
         }
     }
@@ -4234,7 +4161,7 @@ public class ConnectivityService extends IConnectivityManager.Stub
     //   one or more NetworkRequests, or if it is a VPN.
     //
     // - Tears down newNetwork if it just became validated
-    //   (i.e. nascent==JUST_VALIDATED) but turns out to be unneeded.
+    //   but turns out to be unneeded.
     //
     // - If reapUnvalidatedNetworks==REAP, tears down unvalidated
     //   networks that have no chance (i.e. even if validated)
@@ -4247,17 +4174,12 @@ public class ConnectivityService extends IConnectivityManager.Stub
     // as it performs better by a factor of the number of Networks.
     //
     // @param newNetwork is the network to be matched against NetworkRequests.
-    // @param nascent indicates if newNetwork just became validated, in which case it should be
-    //               torn down if unneeded.
     // @param reapUnvalidatedNetworks indicates if an additional pass over all networks should be
     //               performed to tear down unvalidated networks that have no chance (i.e. even if
     //               validated) of becoming the highest scoring network.
-    private void rematchNetworkAndRequests(NetworkAgentInfo newNetwork, NascentState nascent,
+    private void rematchNetworkAndRequests(NetworkAgentInfo newNetwork,
             ReapUnvalidatedNetworks reapUnvalidatedNetworks) {
         if (!newNetwork.created) return;
-        if (nascent == NascentState.JUST_VALIDATED && !newNetwork.everValidated) {
-            loge("ERROR: nascent network not validated.");
-        }
         boolean keep = newNetwork.isVPN();
         boolean isNewDefault = false;
         NetworkAgentInfo oldDefaultNetwork = null;
@@ -4270,7 +4192,7 @@ public class ConnectivityService extends IConnectivityManager.Stub
         for (NetworkRequestInfo nri : mNetworkRequests.values()) {
             NetworkAgentInfo currentNetwork = mNetworkForRequestId.get(nri.request.requestId);
             if (newNetwork == currentNetwork) {
-                if (DBG) {
+                if (VDBG) {
                     log("Network " + newNetwork.name() + " was already satisfying" +
                             " request " + nri.request.requestId + ". No change.");
                 }
@@ -4328,8 +4250,7 @@ public class ConnectivityService extends IConnectivityManager.Stub
         // Linger any networks that are no longer needed.
         for (NetworkAgentInfo nai : affectedNetworks) {
             if (nai.everValidated && unneeded(nai)) {
-                nai.networkMonitor.sendMessage(NetworkMonitor.CMD_NETWORK_LINGER);
-                notifyNetworkCallbacks(nai, ConnectivityManager.CALLBACK_LOSING);
+                linger(nai);
             } else {
                 unlinger(nai);
             }
@@ -4415,19 +4336,10 @@ public class ConnectivityService extends IConnectivityManager.Stub
             if (newNetwork.isVPN()) {
                 mLegacyTypeTracker.add(TYPE_VPN, newNetwork);
             }
-        } else if (nascent == NascentState.JUST_VALIDATED) {
-            // Only tear down newly validated networks here.  Leave unvalidated to either become
-            // validated (and get evaluated against peers, one losing here), or get reaped (see
-            // reapUnvalidatedNetworks) if they have no chance of becoming the highest scoring
-            // network.  Networks that have been up for a while and are validated should be torn
-            // down via the lingering process so communication on that network is given time to
-            // wrap up.
-            if (DBG) log("Validated network turns out to be unwanted.  Tear it down.");
-            teardownUnneededNetwork(newNetwork);
         }
         if (reapUnvalidatedNetworks == ReapUnvalidatedNetworks.REAP) {
             for (NetworkAgentInfo nai : mNetworkAgentInfos.values()) {
-                if (!nai.everValidated && unneeded(nai)) {
+                if (unneeded(nai)) {
                     if (DBG) log("Reaping " + nai.name());
                     teardownUnneededNetwork(nai);
                 }
@@ -4446,10 +4358,8 @@ public class ConnectivityService extends IConnectivityManager.Stub
      *         this argument, otherwise pass {@code changed.getCurrentScore()} or 0 if
      *         {@code changed} is {@code null}. This is because NetworkCapabilities influence a
      *         network's score.
-     * @param nascent indicates if {@code changed} has just been validated.
      */
-    private void rematchAllNetworksAndRequests(NetworkAgentInfo changed, int oldScore,
-            NascentState nascent) {
+    private void rematchAllNetworksAndRequests(NetworkAgentInfo changed, int oldScore) {
         // TODO: This may get slow.  The "changed" parameter is provided for future optimization
         // to avoid the slowness.  It is not simply enough to process just "changed", for
         // example in the case where "changed"'s score decreases and another network should begin
@@ -4458,17 +4368,21 @@ public class ConnectivityService extends IConnectivityManager.Stub
         // Optimization: Only reprocess "changed" if its score improved.  This is safe because it
         // can only add more NetworkRequests satisfied by "changed", and this is exactly what
         // rematchNetworkAndRequests() handles.
-        if (changed != null &&
-                (oldScore < changed.getCurrentScore() || nascent == NascentState.JUST_VALIDATED)) {
-            rematchNetworkAndRequests(changed, nascent, ReapUnvalidatedNetworks.REAP);
+        if (changed != null && oldScore < changed.getCurrentScore()) {
+            rematchNetworkAndRequests(changed, ReapUnvalidatedNetworks.REAP);
         } else {
-            for (Iterator i = mNetworkAgentInfos.values().iterator(); i.hasNext(); ) {
-                rematchNetworkAndRequests((NetworkAgentInfo)i.next(),
-                        NascentState.NOT_JUST_VALIDATED,
+            final NetworkAgentInfo[] nais = mNetworkAgentInfos.values().toArray(
+                    new NetworkAgentInfo[mNetworkAgentInfos.size()]);
+            // Rematch higher scoring networks first to prevent requests first matching a lower
+            // scoring network and then a higher scoring network, which could produce multiple
+            // callbacks and inadvertently unlinger networks.
+            Arrays.sort(nais);
+            for (NetworkAgentInfo nai : nais) {
+                rematchNetworkAndRequests(nai,
                         // Only reap the last time through the loop.  Reaping before all rematching
                         // is complete could incorrectly teardown a network that hasn't yet been
                         // rematched.
-                        i.hasNext() ? ReapUnvalidatedNetworks.DONT_REAP
+                        (nai != nais[nais.length-1]) ? ReapUnvalidatedNetworks.DONT_REAP
                                 : ReapUnvalidatedNetworks.REAP);
             }
         }
@@ -4555,8 +4469,7 @@ public class ConnectivityService extends IConnectivityManager.Stub
             }
 
             // Consider network even though it is not yet validated.
-            rematchNetworkAndRequests(networkAgent, NascentState.NOT_JUST_VALIDATED,
-                    ReapUnvalidatedNetworks.REAP);
+            rematchNetworkAndRequests(networkAgent, ReapUnvalidatedNetworks.REAP);
 
             // This has to happen after matching the requests, because callbacks are just requests.
             notifyNetworkCallbacks(networkAgent, ConnectivityManager.CALLBACK_PRECHECK);
@@ -4587,7 +4500,7 @@ public class ConnectivityService extends IConnectivityManager.Stub
         final int oldScore = nai.getCurrentScore();
         nai.setCurrentScore(score);
 
-        rematchAllNetworksAndRequests(nai, oldScore, NascentState.NOT_JUST_VALIDATED);
+        rematchAllNetworksAndRequests(nai, oldScore);
 
         sendUpdatedScoreToFactories(nai);
     }
@@ -4637,7 +4550,7 @@ public class ConnectivityService extends IConnectivityManager.Stub
             }
             NetworkAgentInfo newDefaultAgent = null;
             if (nai.networkRequests.get(mDefaultRequest.requestId) != null) {
-                newDefaultAgent = mNetworkForRequestId.get(mDefaultRequest.requestId);
+                newDefaultAgent = getDefaultNetwork();
                 if (newDefaultAgent != null) {
                     intent.putExtra(ConnectivityManager.EXTRA_OTHER_NETWORK_INFO,
                             newDefaultAgent.networkInfo);
diff --git a/services/core/java/com/android/server/connectivity/NetworkAgentInfo.java b/services/core/java/com/android/server/connectivity/NetworkAgentInfo.java
index 51c6628..0a0c096 100644
--- a/services/core/java/com/android/server/connectivity/NetworkAgentInfo.java
+++ b/services/core/java/com/android/server/connectivity/NetworkAgentInfo.java
@@ -31,14 +31,15 @@ import com.android.internal.util.AsyncChannel;
 import com.android.server.connectivity.NetworkMonitor;
 
 import java.util.ArrayList;
+import java.util.Comparator;
 
 /**
  * A bag class used by ConnectivityService for holding a collection of most recent
  * information published by a particular NetworkAgent as well as the
  * AsyncChannel/messenger for reaching that NetworkAgent and lists of NetworkRequests
- * interested in using it.
+ * interested in using it.  Default sort order is descending by score.
  */
-public class NetworkAgentInfo {
+public class NetworkAgentInfo implements Comparable<NetworkAgentInfo> {
     public NetworkInfo networkInfo;
     // This Network object should always be used if possible, so as to encourage reuse of the
     // enclosed socket factory and connection pool.  Avoid creating other Network objects.
@@ -72,6 +73,13 @@ public class NetworkAgentInfo {
     // Whether a captive portal was found during the last network validation attempt.
     public boolean lastCaptivePortalDetected;
 
+    // Indicates whether the network is lingering.  Networks are lingered when they become unneeded
+    // as a result of their NetworkRequests being satisfied by a different 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 either the linger timeout expiring and the network
+    // being taken down, or the network satisfying a request again.
+    public boolean lingering;
+
     // This represents the last score received from the NetworkAgent.
     private int currentScore;
     // Penalty applied to scores of Networks that have not been validated.
@@ -175,7 +183,7 @@ public class NetworkAgentInfo {
                 linkProperties + "}  nc{" +
                 networkCapabilities + "}  Score{" + getCurrentScore() + "}  " +
                 "everValidated{" + everValidated + "}  lastValidated{" + lastValidated + "}  " +
-                "created{" + created + "}  " +
+                "created{" + created + "} lingering{" + lingering + "} " +
                 "explicitlySelected{" + networkMisc.explicitlySelected + "} " +
                 "acceptUnvalidated{" + networkMisc.acceptUnvalidated + "} " +
                 "everCaptivePortalDetected{" + everCaptivePortalDetected + "} " +
@@ -188,4 +196,10 @@ public class NetworkAgentInfo {
                 networkInfo.getSubtypeName() + ") - " +
                 (network == null ? "null" : network.toString()) + "]";
     }
+
+    // Enables sorting in descending order of score.
+    @Override
+    public int compareTo(NetworkAgentInfo other) {
+        return other.getCurrentScore() - getCurrentScore();
+    }
 }
diff --git a/services/tests/servicestests/src/com/android/server/ConnectivityServiceTest.java b/services/tests/servicestests/src/com/android/server/ConnectivityServiceTest.java
index fb8a5bb..c1311ed 100644
--- a/services/tests/servicestests/src/com/android/server/ConnectivityServiceTest.java
+++ b/services/tests/servicestests/src/com/android/server/ConnectivityServiceTest.java
@@ -20,8 +20,24 @@ import static android.net.ConnectivityManager.CONNECTIVITY_ACTION;
 import static android.net.ConnectivityManager.TYPE_MOBILE;
 import static android.net.ConnectivityManager.TYPE_WIFI;
 import static android.net.ConnectivityManager.getNetworkTypeName;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_CAPTIVE_PORTAL;
+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_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_MMS;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_NOT_METERED;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_NOT_RESTRICTED;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_NOT_VPN;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_RCS;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_SUPL;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_TRUSTED;
 import static android.net.NetworkCapabilities.NET_CAPABILITY_VALIDATED;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_WIFI_P2P;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_XCAP;
 import static android.net.NetworkCapabilities.TRANSPORT_CELLULAR;
 import static android.net.NetworkCapabilities.TRANSPORT_WIFI;
 import static org.mockito.Matchers.anyInt;
@@ -176,6 +192,11 @@ public class ConnectivityServiceTest extends AndroidTestCase {
             mNetworkAgent.sendNetworkScore(mScore);
         }
 
+        public void addCapability(int capability) {
+            mNetworkCapabilities.addCapability(capability);
+            mNetworkAgent.sendNetworkCapabilities(mNetworkCapabilities);
+        }
+
         /**
          * Transition this NetworkAgent to CONNECTED state.
          * @param validated Indicate if network should pretend to be validated.
@@ -736,10 +757,9 @@ public class ConnectivityServiceTest extends AndroidTestCase {
         assertEquals(CallbackState.NONE, wifiNetworkCallback.getLastCallback());
     }
 
-    @LargeTest
-    public void testNetworkFactoryRequests() throws Exception {
+    private void tryNetworkFactoryRequests(int capability) throws Exception {
         NetworkCapabilities filter = new NetworkCapabilities();
-        filter.addCapability(NET_CAPABILITY_INTERNET);
+        filter.addCapability(capability);
         final HandlerThread handlerThread = new HandlerThread("testNetworkFactoryRequests");
         handlerThread.start();
         final MockNetworkFactory testFactory = new MockNetworkFactory(handlerThread.getLooper(),
@@ -747,57 +767,91 @@ public class ConnectivityServiceTest extends AndroidTestCase {
         testFactory.setScoreFilter(40);
         ConditionVariable cv = testFactory.getNetworkStartedCV();
         testFactory.register();
+        int expectedRequestCount = 1;
+        NetworkCallback networkCallback = null;
+        // For non-INTERNET capabilities we cannot rely on the default request being present, so
+        // add one.
+        if (capability != NET_CAPABILITY_INTERNET) {
+            testFactory.waitForNetworkRequests(1);
+            assertFalse(testFactory.getMyStartRequested());
+            NetworkRequest request = new NetworkRequest.Builder().addCapability(capability).build();
+            networkCallback = new NetworkCallback();
+            mCm.requestNetwork(request, networkCallback);
+            expectedRequestCount++;
+        }
         waitFor(cv);
-        assertEquals(1, testFactory.getMyRequestCount());
-        assertEquals(true, testFactory.getMyStartRequested());
+        assertEquals(expectedRequestCount, testFactory.getMyRequestCount());
+        assertTrue(testFactory.getMyStartRequested());
 
-        // now bring in a higher scored network
+        // Now bring in a higher scored network.
         MockNetworkAgent testAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
-        cv = waitForConnectivityBroadcasts(1);
-        ConditionVariable cvRelease = testFactory.getNetworkStoppedCV();
-        testAgent.connect(true);
+        // Rather than create a validated network which complicates things by registering it's
+        // own NetworkRequest during startup, just bump up the score to cancel out the
+        // unvalidated penalty.
+        testAgent.adjustScore(40);
+        cv = testFactory.getNetworkStoppedCV();
+        testAgent.connect(false);
+        testAgent.addCapability(capability);
         waitFor(cv);
-        // part of the bringup makes another network request and then releases it
-        // wait for the release
-        waitFor(cvRelease);
-        assertEquals(false, testFactory.getMyStartRequested());
-        testFactory.waitForNetworkRequests(1);
+        assertEquals(expectedRequestCount, testFactory.getMyRequestCount());
+        assertFalse(testFactory.getMyStartRequested());
 
-        // bring in a bunch of requests..
+        // Bring in a bunch of requests.
         ConnectivityManager.NetworkCallback[] networkCallbacks =
                 new ConnectivityManager.NetworkCallback[10];
         for (int i = 0; i< networkCallbacks.length; i++) {
             networkCallbacks[i] = new ConnectivityManager.NetworkCallback();
             NetworkRequest.Builder builder = new NetworkRequest.Builder();
-            builder.addCapability(NetworkCapabilities.NET_CAPABILITY_INTERNET);
+            builder.addCapability(capability);
             mCm.requestNetwork(builder.build(), networkCallbacks[i]);
         }
-        testFactory.waitForNetworkRequests(11);
-        assertEquals(false, testFactory.getMyStartRequested());
+        testFactory.waitForNetworkRequests(10 + expectedRequestCount);
+        assertFalse(testFactory.getMyStartRequested());
 
-        // remove the requests
+        // Remove the requests.
         for (int i = 0; i < networkCallbacks.length; i++) {
             mCm.unregisterNetworkCallback(networkCallbacks[i]);
         }
-        testFactory.waitForNetworkRequests(1);
-        assertEquals(false, testFactory.getMyStartRequested());
+        testFactory.waitForNetworkRequests(expectedRequestCount);
+        assertFalse(testFactory.getMyStartRequested());
 
-        // drop the higher scored network
-        cv = waitForConnectivityBroadcasts(1);
+        // Drop the higher scored network.
+        cv = testFactory.getNetworkStartedCV();
         testAgent.disconnect();
         waitFor(cv);
-        assertEquals(1, testFactory.getMyRequestCount());
-        assertEquals(true, testFactory.getMyStartRequested());
+        assertEquals(expectedRequestCount, testFactory.getMyRequestCount());
+        assertTrue(testFactory.getMyStartRequested());
 
         testFactory.unregister();
+        if (networkCallback != null) mCm.unregisterNetworkCallback(networkCallback);
         handlerThread.quit();
     }
 
     @LargeTest
+    public void testNetworkFactoryRequests() throws Exception {
+        tryNetworkFactoryRequests(NET_CAPABILITY_MMS);
+        tryNetworkFactoryRequests(NET_CAPABILITY_SUPL);
+        tryNetworkFactoryRequests(NET_CAPABILITY_DUN);
+        tryNetworkFactoryRequests(NET_CAPABILITY_FOTA);
+        tryNetworkFactoryRequests(NET_CAPABILITY_IMS);
+        tryNetworkFactoryRequests(NET_CAPABILITY_CBS);
+        tryNetworkFactoryRequests(NET_CAPABILITY_WIFI_P2P);
+        tryNetworkFactoryRequests(NET_CAPABILITY_IA);
+        tryNetworkFactoryRequests(NET_CAPABILITY_RCS);
+        tryNetworkFactoryRequests(NET_CAPABILITY_XCAP);
+        tryNetworkFactoryRequests(NET_CAPABILITY_EIMS);
+        tryNetworkFactoryRequests(NET_CAPABILITY_NOT_METERED);
+        tryNetworkFactoryRequests(NET_CAPABILITY_INTERNET);
+        tryNetworkFactoryRequests(NET_CAPABILITY_TRUSTED);
+        tryNetworkFactoryRequests(NET_CAPABILITY_NOT_VPN);
+        // Skipping VALIDATED and CAPTIVE_PORTAL as they're disallowed.
+    }
+
+    @LargeTest
     public void testNoMutableNetworkRequests() throws Exception {
         PendingIntent pendingIntent = PendingIntent.getBroadcast(mContext, 0, new Intent("a"), 0);
         NetworkRequest.Builder builder = new NetworkRequest.Builder();
-        builder.addCapability(NetworkCapabilities.NET_CAPABILITY_VALIDATED);
+        builder.addCapability(NET_CAPABILITY_VALIDATED);
         try {
             mCm.requestNetwork(builder.build(), new NetworkCallback());
             fail();
@@ -807,7 +861,7 @@ public class ConnectivityServiceTest extends AndroidTestCase {
             fail();
         } catch (IllegalArgumentException expected) {}
         builder = new NetworkRequest.Builder();
-        builder.addCapability(NetworkCapabilities.NET_CAPABILITY_CAPTIVE_PORTAL);
+        builder.addCapability(NET_CAPABILITY_CAPTIVE_PORTAL);
         try {
             mCm.requestNetwork(builder.build(), new NetworkCallback());
             fail();