Make VPN more testable and update NC during network change

This commit:
- Updates the NetworkCapabilities with new underlying network during
  IKE Session re-establishment
- Create wrapper classes for IkeSession
- Allow injecting executor and NetworkAgent with
  the Dependencies class

Bug: 192077544
Test: atest VpnTest(new tests)
Change-Id: Ic98e3906c2b02fa8b9f4a8e1366b1dd8a363bb47

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

@@ -498,6 +498,29 @@ public class Vpn {
                 return IKEV2_VPN_RETRY_DELAYS_SEC[retryCount];
             }
         }
+
+        /** Get single threaded executor for IKEv2 VPN */
+        public ScheduledThreadPoolExecutor newScheduledThreadPoolExecutor() {
+            return new ScheduledThreadPoolExecutor(1);
+        }
+
+        /** Get a NetworkAgent instance */
+        public NetworkAgent newNetworkAgent(
+                @NonNull Context context,
+                @NonNull Looper looper,
+                @NonNull String logTag,
+                @NonNull NetworkCapabilities nc,
+                @NonNull LinkProperties lp,
+                @NonNull NetworkScore score,
+                @NonNull NetworkAgentConfig config,
+                @Nullable NetworkProvider provider) {
+            return new NetworkAgent(context, looper, logTag, nc, lp, score, config, provider) {
+                @Override
+                public void onNetworkUnwanted() {
+                    // We are user controlled, not driven by NetworkRequest.
+                }
+            };
+        }
     }
 
     public Vpn(Looper looper, Context context, INetworkManagementService netService, INetd netd,
@@ -1474,15 +1497,10 @@ public class Vpn {
                 ? Arrays.asList(mConfig.underlyingNetworks) : null);
 
         mNetworkCapabilities = capsBuilder.build();
-        mNetworkAgent = new NetworkAgent(mContext, mLooper, NETWORKTYPE /* logtag */,
+        mNetworkAgent = mDeps.newNetworkAgent(mContext, mLooper, NETWORKTYPE /* logtag */,
                 mNetworkCapabilities, lp,
                 new NetworkScore.Builder().setLegacyInt(VPN_DEFAULT_SCORE).build(),
-                networkAgentConfig, mNetworkProvider) {
-            @Override
-            public void onNetworkUnwanted() {
-                // We are user controlled, not driven by NetworkRequest.
-            }
-        };
+                networkAgentConfig, mNetworkProvider);
         final long token = Binder.clearCallingIdentity();
         try {
             mNetworkAgent.register();
@@ -2692,8 +2710,7 @@ public class Vpn {
          * of the mutable Ikev2VpnRunner fields. The Ikev2VpnRunner is built mostly lock-free by
          * virtue of everything being serialized on this executor.
          */
-        @NonNull
-        private final ScheduledThreadPoolExecutor mExecutor = new ScheduledThreadPoolExecutor(1);
+        @NonNull private final ScheduledThreadPoolExecutor mExecutor;
 
         @Nullable private ScheduledFuture<?> mScheduledHandleNetworkLostTimeout;
         @Nullable private ScheduledFuture<?> mScheduledHandleRetryIkeSessionTimeout;
@@ -2714,7 +2731,7 @@ public class Vpn {
         @Nullable private LinkProperties mUnderlyingLinkProperties;
         private final String mSessionKey;
 
-        @Nullable private IkeSession mSession;
+        @Nullable private IkeSessionWrapper mSession;
         @Nullable private IkeSessionConnectionInfo mIkeConnectionInfo;
 
         // mMobikeEnabled can only be updated after IKE AUTH is finished.
@@ -2728,9 +2745,11 @@ public class Vpn {
          */
         private int mRetryCount = 0;
 
-        IkeV2VpnRunner(@NonNull Ikev2VpnProfile profile) {
+        IkeV2VpnRunner(
+                @NonNull Ikev2VpnProfile profile, @NonNull ScheduledThreadPoolExecutor executor) {
             super(TAG);
             mProfile = profile;
+            mExecutor = executor;
             mIpSecManager = (IpSecManager) mContext.getSystemService(Context.IPSEC_SERVICE);
             mNetworkCallback = new VpnIkev2Utils.Ikev2VpnNetworkCallback(TAG, this, mExecutor);
             mSessionKey = UUID.randomUUID().toString();
@@ -2743,7 +2762,7 @@ public class Vpn {
 
             // To avoid hitting RejectedExecutionException upon shutdown of the mExecutor */
             mExecutor.setRejectedExecutionHandler(
-                    (r, executor) -> {
+                    (r, exe) -> {
                         Log.d(TAG, "Runnable " + r + " rejected by the mExecutor");
                     });
         }
@@ -2884,7 +2903,6 @@ public class Vpn {
                     mConfig.dnsServers.addAll(dnsAddrStrings);
 
                     mConfig.underlyingNetworks = new Network[] {network};
-
                     mConfig.disallowedApplications = getAppExclusionList(mPackage);
 
                     networkAgent = mNetworkAgent;
@@ -2900,6 +2918,10 @@ public class Vpn {
                     } else {
                         // Underlying networks also set in agentConnect()
                         networkAgent.setUnderlyingNetworks(Collections.singletonList(network));
+                        mNetworkCapabilities =
+                                new NetworkCapabilities.Builder(mNetworkCapabilities)
+                                        .setUnderlyingNetworks(Collections.singletonList(network))
+                                        .build();
                     }
 
                     lp = makeLinkProperties(); // Accesses VPN instance fields; must be locked
@@ -4015,7 +4037,9 @@ public class Vpn {
                 case VpnProfile.TYPE_IKEV2_IPSEC_RSA:
                 case VpnProfile.TYPE_IKEV2_FROM_IKE_TUN_CONN_PARAMS:
                     mVpnRunner =
-                            new IkeV2VpnRunner(Ikev2VpnProfile.fromVpnProfile(profile));
+                            new IkeV2VpnRunner(
+                                    Ikev2VpnProfile.fromVpnProfile(profile),
+                                    mDeps.newScheduledThreadPoolExecutor());
                     mVpnRunner.start();
                     break;
                 default:
@@ -4185,6 +4209,31 @@ public class Vpn {
         return isCurrentIkev2VpnLocked(packageName) ? makeVpnProfileStateLocked() : null;
     }
 
+    /**
+     * Proxy to allow testing
+     *
+     * @hide
+     */
+    @VisibleForTesting
+    public static class IkeSessionWrapper {
+        private final IkeSession mImpl;
+
+        /** Create an IkeSessionWrapper */
+        public IkeSessionWrapper(IkeSession session) {
+            mImpl = session;
+        }
+
+        /** Update the underlying network of the IKE Session */
+        public void setNetwork(@NonNull Network network) {
+            mImpl.setNetwork(network);
+        }
+
+        /** Forcibly terminate the IKE Session */
+        public void kill() {
+            mImpl.kill();
+        }
+    }
+
     /**
      * Proxy to allow testing
      *
@@ -4193,20 +4242,21 @@ public class Vpn {
     @VisibleForTesting
     public static class Ikev2SessionCreator {
         /** Creates a IKE session */
-        public IkeSession createIkeSession(
+        public IkeSessionWrapper createIkeSession(
                 @NonNull Context context,
                 @NonNull IkeSessionParams ikeSessionParams,
                 @NonNull ChildSessionParams firstChildSessionParams,
                 @NonNull Executor userCbExecutor,
                 @NonNull IkeSessionCallback ikeSessionCallback,
                 @NonNull ChildSessionCallback firstChildSessionCallback) {
-            return new IkeSession(
-                    context,
-                    ikeSessionParams,
-                    firstChildSessionParams,
-                    userCbExecutor,
-                    ikeSessionCallback,
-                    firstChildSessionCallback);
+            return new IkeSessionWrapper(
+                    new IkeSession(
+                            context,
+                            ikeSessionParams,
+                            firstChildSessionParams,
+                            userCbExecutor,
+                            ikeSessionCallback,
+                            firstChildSessionCallback));
         }
     }