Have registerNetworkAgent return a Network.

This is more useful.

Bug: 135998869
Test: FrameworkNetTests FrameworksWifiTests
Change-Id: I74710b396fa3cec6fbae08d3dcab5bc90dc2e669
diff --git a/core/java/android/net/ConnectivityManager.java b/core/java/android/net/ConnectivityManager.java
index 4bf3e90..dda9ff4cb 100644
--- a/core/java/android/net/ConnectivityManager.java
+++ b/core/java/android/net/ConnectivityManager.java
@@ -3169,10 +3169,10 @@
     /**
      * @hide
      * Register a NetworkAgent with ConnectivityService.
-     * @return NetID corresponding to NetworkAgent.
+     * @return Network corresponding to NetworkAgent.
      */
     @RequiresPermission(android.Manifest.permission.NETWORK_FACTORY)
-    public int registerNetworkAgent(Messenger messenger, NetworkInfo ni, LinkProperties lp,
+    public Network registerNetworkAgent(Messenger messenger, NetworkInfo ni, LinkProperties lp,
             NetworkCapabilities nc, int score, NetworkAgentConfig config) {
         return registerNetworkAgent(messenger, ni, lp, nc, score, config, NetworkProvider.ID_NONE);
     }
@@ -3180,10 +3180,10 @@
     /**
      * @hide
      * Register a NetworkAgent with ConnectivityService.
-     * @return NetID corresponding to NetworkAgent.
+     * @return Network corresponding to NetworkAgent.
      */
     @RequiresPermission(android.Manifest.permission.NETWORK_FACTORY)
-    public int registerNetworkAgent(Messenger messenger, NetworkInfo ni, LinkProperties lp,
+    public Network registerNetworkAgent(Messenger messenger, NetworkInfo ni, LinkProperties lp,
             NetworkCapabilities nc, int score, NetworkAgentConfig config, int providerId) {
         try {
             return mService.registerNetworkAgent(messenger, ni, lp, nc, score, config, providerId);
diff --git a/core/java/android/net/IConnectivityManager.aidl b/core/java/android/net/IConnectivityManager.aidl
index 7691beb..186196bd 100644
--- a/core/java/android/net/IConnectivityManager.aidl
+++ b/core/java/android/net/IConnectivityManager.aidl
@@ -152,7 +152,7 @@
 
     void declareNetworkRequestUnfulfillable(in NetworkRequest request);
 
-    int registerNetworkAgent(in Messenger messenger, in NetworkInfo ni, in LinkProperties lp,
+    Network registerNetworkAgent(in Messenger messenger, in NetworkInfo ni, in LinkProperties lp,
             in NetworkCapabilities nc, int score, in NetworkAgentConfig config,
             in int factorySerialNumber);
 
diff --git a/core/java/android/net/NetworkAgent.java b/core/java/android/net/NetworkAgent.java
index fc72eec..d286204 100644
--- a/core/java/android/net/NetworkAgent.java
+++ b/core/java/android/net/NetworkAgent.java
@@ -44,9 +44,10 @@
  * @hide
  */
 public abstract class NetworkAgent {
-    // Guaranteed to be valid (not NETID_UNSET), otherwise registerNetworkAgent() would have thrown
-    // an exception.
-    public final int netId;
+    // Guaranteed to be non-null, otherwise registerNetworkAgent() would have thrown
+    // an exception. Be careful in tests when mocking though.
+    @NonNull
+    public final Network network;
 
     private final Handler mHandler;
     private volatile AsyncChannel mAsyncChannel;
@@ -246,7 +247,7 @@
         if (VDBG) log("Registering NetworkAgent");
         ConnectivityManager cm = (ConnectivityManager)mContext.getSystemService(
                 Context.CONNECTIVITY_SERVICE);
-        netId = cm.registerNetworkAgent(new Messenger(mHandler), new NetworkInfo(ni),
+        network = cm.registerNetworkAgent(new Messenger(mHandler), new NetworkInfo(ni),
                 new LinkProperties(lp), new NetworkCapabilities(nc), score, config,
                 providerId);
     }
diff --git a/services/core/java/com/android/server/ConnectivityService.java b/services/core/java/com/android/server/ConnectivityService.java
index f4962e5..5435cba 100644
--- a/services/core/java/com/android/server/ConnectivityService.java
+++ b/services/core/java/com/android/server/ConnectivityService.java
@@ -5505,7 +5505,10 @@
     // changes that would conflict throughout the automerger graph. Having this method temporarily
     // helps with the process of going through with all these dependent changes across the entire
     // tree.
-    public int registerNetworkAgent(Messenger messenger, NetworkInfo networkInfo,
+    /**
+     * Register a new agent. {@see #registerNetworkAgent} below.
+     */
+    public Network registerNetworkAgent(Messenger messenger, NetworkInfo networkInfo,
             LinkProperties linkProperties, NetworkCapabilities networkCapabilities,
             int currentScore, NetworkAgentConfig networkAgentConfig) {
         return registerNetworkAgent(messenger, networkInfo, linkProperties, networkCapabilities,
@@ -5526,8 +5529,9 @@
      *         {@link NetworkAgentInfo#getCurrentScore}.
      * @param networkAgentConfig metadata about the network. This is never updated.
      * @param providerId the ID of the provider owning this NetworkAgent.
+     * @return the network created for this agent.
      */
-    public int registerNetworkAgent(Messenger messenger, NetworkInfo networkInfo,
+    public Network registerNetworkAgent(Messenger messenger, NetworkInfo networkInfo,
             LinkProperties linkProperties, NetworkCapabilities networkCapabilities,
             int currentScore, NetworkAgentConfig networkAgentConfig, int providerId) {
         enforceNetworkFactoryPermission();
@@ -5560,7 +5564,7 @@
         // If the network disconnects or sends any other event before that, messages are deferred by
         // NetworkAgent until nai.asyncChannel.connect(), which will be called when finalizing the
         // registration.
-        return nai.network.netId;
+        return nai.network;
     }
 
     private void handleRegisterNetworkAgent(NetworkAgentInfo nai, INetworkMonitor networkMonitor) {
diff --git a/services/core/java/com/android/server/TestNetworkService.java b/services/core/java/com/android/server/TestNetworkService.java
index c27b0da..ed3bab9 100644
--- a/services/core/java/com/android/server/TestNetworkService.java
+++ b/services/core/java/com/android/server/TestNetworkService.java
@@ -218,7 +218,7 @@
             // Has to be in TestNetworkAgent to ensure all teardown codepaths properly clean up
             // resources, even for binder death or unwanted calls.
             synchronized (mTestNetworkTracker) {
-                mTestNetworkTracker.remove(netId);
+                mTestNetworkTracker.remove(network.netId);
             }
         }
     }
@@ -337,7 +337,7 @@
                                             callingUid,
                                             binder);
 
-                            mTestNetworkTracker.put(agent.netId, agent);
+                            mTestNetworkTracker.put(agent.network.netId, agent);
                         }
                     } catch (SocketException e) {
                         throw new UncheckedIOException(e);
diff --git a/services/core/java/com/android/server/connectivity/Vpn.java b/services/core/java/com/android/server/connectivity/Vpn.java
index 476f3f8..1d2a905 100644
--- a/services/core/java/com/android/server/connectivity/Vpn.java
+++ b/services/core/java/com/android/server/connectivity/Vpn.java
@@ -848,7 +848,7 @@
     }
 
     public int getNetId() {
-        return mNetworkAgent != null ? mNetworkAgent.netId : NETID_UNSET;
+        return mNetworkAgent != null ? mNetworkAgent.network.netId : NETID_UNSET;
     }
 
     private LinkProperties makeLinkProperties() {
diff --git a/tests/net/integration/util/com/android/server/NetworkAgentWrapper.java b/tests/net/integration/util/com/android/server/NetworkAgentWrapper.java
index 0ab5c97..11d5b25 100644
--- a/tests/net/integration/util/com/android/server/NetworkAgentWrapper.java
+++ b/tests/net/integration/util/com/android/server/NetworkAgentWrapper.java
@@ -222,7 +222,7 @@
 
     @Override
     public Network getNetwork() {
-        return new Network(mNetworkAgent.netId);
+        return mNetworkAgent.network;
     }
 
     public void expectPreventReconnectReceived(long timeoutMs) {
diff --git a/tests/net/java/com/android/server/ConnectivityServiceTest.java b/tests/net/java/com/android/server/ConnectivityServiceTest.java
index b2d363e..1901a1d 100644
--- a/tests/net/java/com/android/server/ConnectivityServiceTest.java
+++ b/tests/net/java/com/android/server/ConnectivityServiceTest.java
@@ -575,7 +575,7 @@
                 }
             };
 
-            assertEquals(na.netId, nmNetworkCaptor.getValue().netId);
+            assertEquals(na.network.netId, nmNetworkCaptor.getValue().netId);
             mNmCallbacks = nmCbCaptor.getValue();
 
             mNmCallbacks.onNetworkMonitorCreated(mNetworkMonitor);