Get NetworkCapabilities from NetworkCallback

Register network callback for all networks and record
NetworkCapabilities for every networks. Once onDnsEvent is
triggered, use the netId it passes in to find the corresponding
NetworkCapabilities instead of using netId to create a Network
object(hidden API) then get the NetworkCapabilities by
ConnectivityManager#getNetworkCapabilities.

Bug: 182963397
Test: m
Test: atest IpConnectivityMetricsTest
Test: atest NetdEventListenerServiceTest
Change-Id: I91d68ca33253831b78def1ddeb074ba944a5d6ad

services/core/java/com/android/server/connectivity/NetdEventListenerService.java

@@ -18,12 +18,14 @@ package com.android.server.connectivity;
 
 import static android.util.TimeUtils.NANOS_PER_MS;
 
+import android.annotation.Nullable;
 import android.content.Context;
 import android.net.ConnectivityManager;
 import android.net.INetdEventCallback;
 import android.net.MacAddress;
 import android.net.Network;
 import android.net.NetworkCapabilities;
+import android.net.NetworkRequest;
 import android.net.metrics.ConnectStats;
 import android.net.metrics.DnsEvent;
 import android.net.metrics.INetdEventListener;
@@ -98,6 +100,7 @@ public class NetdEventListenerService extends INetdEventListener.Stub {
     private final TokenBucket mConnectTb =
             new TokenBucket(CONNECT_LATENCY_FILL_RATE, CONNECT_LATENCY_BURST_LIMIT);
 
+    final TransportForNetIdNetworkCallback mCallback = new TransportForNetIdNetworkCallback();
 
     /**
      * There are only 3 possible callbacks.
@@ -158,6 +161,9 @@ public class NetdEventListenerService extends INetdEventListener.Stub {
     public NetdEventListenerService(ConnectivityManager cm) {
         // We are started when boot is complete, so ConnectivityService should already be running.
         mCm = cm;
+        // Clear all capabilities to listen all networks.
+        mCm.registerNetworkCallback(new NetworkRequest.Builder().clearCapabilities().build(),
+                mCallback);
     }
 
     private static long projectSnapshotTime(long timeMs) {
@@ -389,18 +395,13 @@ public class NetdEventListenerService extends INetdEventListener.Stub {
     }
 
     private long getTransports(int netId) {
-        // TODO: directly query ConnectivityService instead of going through Binder interface.
+        final NetworkCapabilities nc = mCallback.getNetworkCapabilities(netId);
-        NetworkCapabilities nc = mCm.getNetworkCapabilities(new Network(netId));
         if (nc == null) {
             return 0;
         }
         return BitUtils.packBits(nc.getTransportTypes());
     }
 
-    private static void maybeLog(String s, Object... args) {
-        if (DBG) Log.d(TAG, String.format(s, args));
-    }
-
     /** Helper class for buffering summaries of NetworkMetrics at regular time intervals */
     static class NetworkMetricsSnapshot {
 
@@ -428,4 +429,29 @@ public class NetdEventListenerService extends INetdEventListener.Stub {
             return String.format("%tT.%tL: %s", timeMs, timeMs, j.toString());
         }
     }
+
+    private class TransportForNetIdNetworkCallback extends ConnectivityManager.NetworkCallback {
+        private final SparseArray<NetworkCapabilities> mCapabilities = new SparseArray<>();
+
+        @Override
+        public void onCapabilitiesChanged(Network network, NetworkCapabilities nc) {
+            synchronized (mCapabilities) {
+                mCapabilities.put(network.getNetId(), nc);
+            }
+        }
+
+        @Override
+        public void onLost(Network network) {
+            synchronized (mCapabilities) {
+                mCapabilities.remove(network.getNetId());
+            }
+        }
+
+        @Nullable
+        public NetworkCapabilities getNetworkCapabilities(int netId) {
+            synchronized (mCapabilities) {
+                return mCapabilities.get(netId);
+            }
+        }
+    }
 }

tests/net/java/com/android/server/connectivity/IpConnectivityMetricsTest.java

@@ -22,7 +22,9 @@ import static android.net.metrics.INetdEventListener.EVENT_GETHOSTBYNAME;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.fail;
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import android.content.Context;
@@ -56,6 +58,7 @@ import com.android.server.connectivity.metrics.nano.IpConnectivityLogClass;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
@@ -80,6 +83,12 @@ public class IpConnectivityMetricsTest {
 
     IpConnectivityMetrics mService;
     NetdEventListenerService mNetdListener;
+    final NetworkCapabilities mNcWifi = new NetworkCapabilities.Builder()
+            .addTransportType(NetworkCapabilities.TRANSPORT_WIFI)
+            .build();
+    final NetworkCapabilities mNcCell = new NetworkCapabilities.Builder()
+            .addTransportType(NetworkCapabilities.TRANSPORT_CELLULAR)
+            .build();
 
     @Before
     public void setUp() {
@@ -263,14 +272,6 @@ public class IpConnectivityMetricsTest {
         // TODO: instead of comparing textpb to textpb, parse textpb and compare proto to proto.
         IpConnectivityLog logger = new IpConnectivityLog(mService.impl);
 
-        NetworkCapabilities ncWifi = new NetworkCapabilities();
-        NetworkCapabilities ncCell = new NetworkCapabilities();
-        ncWifi.addTransportType(NetworkCapabilities.TRANSPORT_WIFI);
-        ncCell.addTransportType(NetworkCapabilities.TRANSPORT_CELLULAR);
-
-        when(mCm.getNetworkCapabilities(new Network(100))).thenReturn(ncWifi);
-        when(mCm.getNetworkCapabilities(new Network(101))).thenReturn(ncCell);
-
         ApfStats apfStats = new ApfStats.Builder()
                 .setDurationMs(45000)
                 .setReceivedRas(10)
@@ -584,11 +585,21 @@ public class IpConnectivityMetricsTest {
         return buffer.toString();
     }
 
-    void connectEvent(int netid, int error, int latencyMs, String ipAddr) throws Exception {
+    private void setCapabilities(int netId) {
-        mNetdListener.onConnectEvent(netid, error, latencyMs, ipAddr, 80, 1);
+        final ArgumentCaptor<ConnectivityManager.NetworkCallback> networkCallback =
+                ArgumentCaptor.forClass(ConnectivityManager.NetworkCallback.class);
+        verify(mCm).registerNetworkCallback(any(), networkCallback.capture());
+        networkCallback.getValue().onCapabilitiesChanged(new Network(netId),
+                netId == 100 ? mNcWifi : mNcCell);
+    }
+
+    void connectEvent(int netId, int error, int latencyMs, String ipAddr) throws Exception {
+        setCapabilities(netId);
+        mNetdListener.onConnectEvent(netId, error, latencyMs, ipAddr, 80, 1);
     }
 
     void dnsEvent(int netId, int type, int result, int latency) throws Exception {
+        setCapabilities(netId);
         mNetdListener.onDnsEvent(netId, type, result, latency, "", null, 0, 0);
     }
 

tests/net/java/com/android/server/connectivity/NetdEventListenerServiceTest.java

@@ -23,8 +23,9 @@ import static com.android.testutils.MiscAsserts.assertStringContains;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.fail;
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.when;
+import static org.mockito.Mockito.verify;
 
 import android.content.Context;
 import android.net.ConnectivityManager;
@@ -42,6 +43,7 @@ import com.android.server.connectivity.metrics.nano.IpConnectivityLogClass.IpCon
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
 
 import java.io.FileOutputStream;
 import java.io.PrintWriter;
@@ -61,18 +63,16 @@ public class NetdEventListenerServiceTest {
 
     NetdEventListenerService mService;
     ConnectivityManager mCm;
+    final NetworkCapabilities mNcWifi = new NetworkCapabilities.Builder()
+            .addTransportType(NetworkCapabilities.TRANSPORT_WIFI)
+            .build();
+    final NetworkCapabilities mNcCell = new NetworkCapabilities.Builder()
+            .addTransportType(NetworkCapabilities.TRANSPORT_CELLULAR)
+            .build();
 
     @Before
     public void setUp() {
-        NetworkCapabilities ncWifi = new NetworkCapabilities();
-        NetworkCapabilities ncCell = new NetworkCapabilities();
-        ncWifi.addTransportType(NetworkCapabilities.TRANSPORT_WIFI);
-        ncCell.addTransportType(NetworkCapabilities.TRANSPORT_CELLULAR);
-
         mCm = mock(ConnectivityManager.class);
-        when(mCm.getNetworkCapabilities(new Network(100))).thenReturn(ncWifi);
-        when(mCm.getNetworkCapabilities(new Network(101))).thenReturn(ncCell);
-
         mService = new NetdEventListenerService(mCm);
     }
 
@@ -470,7 +470,16 @@ public class NetdEventListenerServiceTest {
         assertEquals(want, got);
     }
 
+    private void setCapabilities(int netId) {
+        final ArgumentCaptor<ConnectivityManager.NetworkCallback> networkCallback =
+                ArgumentCaptor.forClass(ConnectivityManager.NetworkCallback.class);
+        verify(mCm).registerNetworkCallback(any(), networkCallback.capture());
+        networkCallback.getValue().onCapabilitiesChanged(new Network(netId),
+                netId == 100 ? mNcWifi : mNcCell);
+    }
+
     Thread connectEventAction(int netId, int error, int latencyMs, String ipAddr) {
+        setCapabilities(netId);
         return new Thread(() -> {
             try {
                 mService.onConnectEvent(netId, error, latencyMs, ipAddr, 80, 1);
@@ -481,6 +490,7 @@ public class NetdEventListenerServiceTest {
     }
 
     void dnsEvent(int netId, int type, int result, int latency) throws Exception {
+        setCapabilities(netId);
         mService.onDnsEvent(netId, type, result, latency, "", null, 0, 0);
     }