Merge changes If8097d58,Ic8ec4587,If37b26fa am: fbd9f328ab

Original change: https://android-review.googlesource.com/c/platform/frameworks/base/+/2347143

Change-Id: Iaae6438c1ce5e98501c25d3ca9eb4caf6b885dac
Signed-off-by: Automerger Merge Worker <android-build-automerger-merge-worker@system.gserviceaccount.com>

core/api/current.txt

@@ -27303,13 +27303,17 @@ package android.net.vcn {
     method @IntRange(from=0x500) public int getMaxMtu();
     method @NonNull public long[] getRetryIntervalsMillis();
     method @NonNull public java.util.List<android.net.vcn.VcnUnderlyingNetworkTemplate> getVcnUnderlyingNetworkPriorities();
+    method public boolean hasGatewayOption(int);
+    field public static final int VCN_GATEWAY_OPTION_ENABLE_DATA_STALL_RECOVERY_WITH_MOBILITY = 0; // 0x0
   }
 
   public static final class VcnGatewayConnectionConfig.Builder {
     ctor public VcnGatewayConnectionConfig.Builder(@NonNull String, @NonNull android.net.ipsec.ike.IkeTunnelConnectionParams);
     method @NonNull public android.net.vcn.VcnGatewayConnectionConfig.Builder addExposedCapability(int);
+    method @NonNull public android.net.vcn.VcnGatewayConnectionConfig.Builder addGatewayOption(int);
     method @NonNull public android.net.vcn.VcnGatewayConnectionConfig build();
     method @NonNull public android.net.vcn.VcnGatewayConnectionConfig.Builder removeExposedCapability(int);
+    method @NonNull public android.net.vcn.VcnGatewayConnectionConfig.Builder removeGatewayOption(int);
     method @NonNull public android.net.vcn.VcnGatewayConnectionConfig.Builder setMaxMtu(@IntRange(from=0x500) int);
     method @NonNull public android.net.vcn.VcnGatewayConnectionConfig.Builder setRetryIntervalsMillis(@NonNull long[]);
     method @NonNull public android.net.vcn.VcnGatewayConnectionConfig.Builder setVcnUnderlyingNetworkPriorities(@NonNull java.util.List<android.net.vcn.VcnUnderlyingNetworkTemplate>);

core/java/android/net/vcn/VcnGatewayConnectionConfig.java

@@ -42,6 +42,7 @@ import java.lang.annotation.RetentionPolicy;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Objects;
 import java.util.Set;
@@ -130,6 +131,30 @@ public final class VcnGatewayConnectionConfig {
             })
     public @interface VcnSupportedCapability {}
 
+    /**
+     * Perform mobility update to attempt recovery from suspected data stalls.
+     *
+     * <p>If set, the gatway connection will monitor the data stall detection of the VCN network.
+     * When there is a suspected data stall, the gateway connection will attempt recovery by
+     * performing a mobility update on the underlying IKE session.
+     */
+    public static final int VCN_GATEWAY_OPTION_ENABLE_DATA_STALL_RECOVERY_WITH_MOBILITY = 0;
+
+    /** @hide */
+    @Retention(RetentionPolicy.SOURCE)
+    @IntDef(
+            prefix = {"VCN_GATEWAY_OPTION_"},
+            value = {
+                VCN_GATEWAY_OPTION_ENABLE_DATA_STALL_RECOVERY_WITH_MOBILITY,
+            })
+    public @interface VcnGatewayOption {}
+
+    private static final Set<Integer> ALLOWED_GATEWAY_OPTIONS = new ArraySet<>();
+
+    static {
+        ALLOWED_GATEWAY_OPTIONS.add(VCN_GATEWAY_OPTION_ENABLE_DATA_STALL_RECOVERY_WITH_MOBILITY);
+    }
+
     private static final int DEFAULT_MAX_MTU = 1500;
 
     /**
@@ -201,6 +226,9 @@ public final class VcnGatewayConnectionConfig {
     private static final String RETRY_INTERVAL_MS_KEY = "mRetryIntervalsMs";
     @NonNull private final long[] mRetryIntervalsMs;
 
+    private static final String GATEWAY_OPTIONS_KEY = "mGatewayOptions";
+    @NonNull private final Set<Integer> mGatewayOptions;
+
     /** Builds a VcnGatewayConnectionConfig with the specified parameters. */
     private VcnGatewayConnectionConfig(
             @NonNull String gatewayConnectionName,
@@ -208,12 +236,14 @@ public final class VcnGatewayConnectionConfig {
             @NonNull Set<Integer> exposedCapabilities,
             @NonNull List<VcnUnderlyingNetworkTemplate> underlyingNetworkTemplates,
             @NonNull long[] retryIntervalsMs,
-            @IntRange(from = MIN_MTU_V6) int maxMtu) {
+            @IntRange(from = MIN_MTU_V6) int maxMtu,
+            @NonNull Set<Integer> gatewayOptions) {
         mGatewayConnectionName = gatewayConnectionName;
         mTunnelConnectionParams = tunnelConnectionParams;
         mExposedCapabilities = new TreeSet(exposedCapabilities);
         mRetryIntervalsMs = retryIntervalsMs;
         mMaxMtu = maxMtu;
+        mGatewayOptions = Collections.unmodifiableSet(new HashSet(gatewayOptions));
 
         mUnderlyingNetworkTemplates = new ArrayList<>(underlyingNetworkTemplates);
         if (mUnderlyingNetworkTemplates.isEmpty()) {
@@ -256,6 +286,20 @@ public final class VcnGatewayConnectionConfig {
                             VcnUnderlyingNetworkTemplate::fromPersistableBundle);
         }
 
+        final PersistableBundle gatewayOptionsBundle = in.getPersistableBundle(GATEWAY_OPTIONS_KEY);
+
+        if (gatewayOptionsBundle == null) {
+            // GATEWAY_OPTIONS_KEY was added in Android U. Thus VcnGatewayConnectionConfig created
+            // on old platforms will not have this data and will be assigned with the default value
+            mGatewayOptions = Collections.emptySet();
+        } else {
+            mGatewayOptions =
+                    new HashSet<>(
+                            PersistableBundleUtils.toList(
+                                    gatewayOptionsBundle,
+                                    PersistableBundleUtils.INTEGER_DESERIALIZER));
+        }
+
         mRetryIntervalsMs = in.getLongArray(RETRY_INTERVAL_MS_KEY);
         mMaxMtu = in.getInt(MAX_MTU_KEY);
 
@@ -279,6 +323,10 @@ public final class VcnGatewayConnectionConfig {
 
         Preconditions.checkArgument(
                 mMaxMtu >= MIN_MTU_V6, "maxMtu must be at least IPv6 min MTU (1280)");
+
+        for (int option : mGatewayOptions) {
+            validateGatewayOption(option);
+        }
     }
 
     private static void checkValidCapability(int capability) {
@@ -315,6 +363,12 @@ public final class VcnGatewayConnectionConfig {
         }
     }
 
+    private static void validateGatewayOption(int option) {
+        if (!ALLOWED_GATEWAY_OPTIONS.contains(option)) {
+            throw new IllegalArgumentException("Invalid vcn gateway option: " + option);
+        }
+    }
+
     /**
      * Returns the configured Gateway Connection name.
      *
@@ -398,6 +452,19 @@ public final class VcnGatewayConnectionConfig {
         return mMaxMtu;
     }
 
+    /**
+     * Checks if the given VCN gateway option is enabled.
+     *
+     * @param option the option to check.
+     * @throws IllegalArgumentException if the provided option is invalid.
+     * @see Builder#addGatewayOption(int)
+     * @see Builder#removeGatewayOption(int)
+     */
+    public boolean hasGatewayOption(@VcnGatewayOption int option) {
+        validateGatewayOption(option);
+        return mGatewayOptions.contains(option);
+    }
+
     /**
      * Converts this config to a PersistableBundle.
      *
@@ -418,11 +485,16 @@ public final class VcnGatewayConnectionConfig {
                 PersistableBundleUtils.fromList(
                         mUnderlyingNetworkTemplates,
                         VcnUnderlyingNetworkTemplate::toPersistableBundle);
+        final PersistableBundle gatewayOptionsBundle =
+                PersistableBundleUtils.fromList(
+                        new ArrayList<>(mGatewayOptions),
+                        PersistableBundleUtils.INTEGER_SERIALIZER);
 
         result.putString(GATEWAY_CONNECTION_NAME_KEY, mGatewayConnectionName);
         result.putPersistableBundle(TUNNEL_CONNECTION_PARAMS_KEY, tunnelConnectionParamsBundle);
         result.putPersistableBundle(EXPOSED_CAPABILITIES_KEY, exposedCapsBundle);
         result.putPersistableBundle(UNDERLYING_NETWORK_TEMPLATES_KEY, networkTemplatesBundle);
+        result.putPersistableBundle(GATEWAY_OPTIONS_KEY, gatewayOptionsBundle);
         result.putLongArray(RETRY_INTERVAL_MS_KEY, mRetryIntervalsMs);
         result.putInt(MAX_MTU_KEY, mMaxMtu);
 
@@ -437,7 +509,8 @@ public final class VcnGatewayConnectionConfig {
                 mExposedCapabilities,
                 mUnderlyingNetworkTemplates,
                 Arrays.hashCode(mRetryIntervalsMs),
-                mMaxMtu);
+                mMaxMtu,
+                mGatewayOptions);
     }
 
     @Override
@@ -452,7 +525,8 @@ public final class VcnGatewayConnectionConfig {
                 && mExposedCapabilities.equals(rhs.mExposedCapabilities)
                 && mUnderlyingNetworkTemplates.equals(rhs.mUnderlyingNetworkTemplates)
                 && Arrays.equals(mRetryIntervalsMs, rhs.mRetryIntervalsMs)
-                && mMaxMtu == rhs.mMaxMtu;
+                && mMaxMtu == rhs.mMaxMtu
+                && mGatewayOptions.equals(rhs.mGatewayOptions);
     }
 
     /**
@@ -470,6 +544,8 @@ public final class VcnGatewayConnectionConfig {
         @NonNull private long[] mRetryIntervalsMs = DEFAULT_RETRY_INTERVALS_MS;
         private int mMaxMtu = DEFAULT_MAX_MTU;
 
+        @NonNull private final Set<Integer> mGatewayOptions = new ArraySet<>();
+
         // TODO: (b/175829816) Consider VCN-exposed capabilities that may be transport dependent.
         //       Consider the case where the VCN might only expose MMS on WiFi, but defer to MMS
         //       when on Cell.
@@ -627,6 +703,34 @@ public final class VcnGatewayConnectionConfig {
             return this;
         }
 
+        /**
+         * Enables the specified VCN gateway option.
+         *
+         * @param option the option to be enabled
+         * @return this {@link Builder} instance, for chaining
+         * @throws IllegalArgumentException if the provided option is invalid
+         */
+        @NonNull
+        public Builder addGatewayOption(@VcnGatewayOption int option) {
+            validateGatewayOption(option);
+            mGatewayOptions.add(option);
+            return this;
+        }
+
+        /**
+         * Resets (disables) the specified VCN gateway option.
+         *
+         * @param option the option to be disabled
+         * @return this {@link Builder} instance, for chaining
+         * @throws IllegalArgumentException if the provided option is invalid
+         */
+        @NonNull
+        public Builder removeGatewayOption(@VcnGatewayOption int option) {
+            validateGatewayOption(option);
+            mGatewayOptions.remove(option);
+            return this;
+        }
+
         /**
          * Builds and validates the VcnGatewayConnectionConfig.
          *
@@ -640,7 +744,8 @@ public final class VcnGatewayConnectionConfig {
                     mExposedCapabilities,
                     mUnderlyingNetworkTemplates,
                     mRetryIntervalsMs,
-                    mMaxMtu);
+                    mMaxMtu,
+                    mGatewayOptions);
         }
     }
 }

services/core/java/com/android/server/vcn/VcnGatewayConnection.java

@@ -26,6 +26,7 @@ import static android.net.NetworkCapabilities.NET_CAPABILITY_NOT_VCN_MANAGED;
 import static android.net.NetworkCapabilities.TRANSPORT_CELLULAR;
 import static android.net.NetworkCapabilities.TRANSPORT_WIFI;
 import static android.net.ipsec.ike.exceptions.IkeProtocolException.ERROR_TYPE_AUTHENTICATION_FAILED;
+import static android.net.vcn.VcnGatewayConnectionConfig.VCN_GATEWAY_OPTION_ENABLE_DATA_STALL_RECOVERY_WITH_MOBILITY;
 import static android.net.vcn.VcnManager.VCN_ERROR_CODE_CONFIG_ERROR;
 import static android.net.vcn.VcnManager.VCN_ERROR_CODE_INTERNAL_ERROR;
 import static android.net.vcn.VcnManager.VCN_ERROR_CODE_NETWORK_ERROR;
@@ -36,6 +37,8 @@ import static com.android.server.VcnManagementService.VDBG;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.content.Context;
+import android.net.ConnectivityDiagnosticsManager;
+import android.net.ConnectivityDiagnosticsManager.ConnectivityDiagnosticsCallback;
 import android.net.ConnectivityManager;
 import android.net.InetAddresses;
 import android.net.IpPrefix;
@@ -50,6 +53,7 @@ import android.net.NetworkAgent;
 import android.net.NetworkAgentConfig;
 import android.net.NetworkCapabilities;
 import android.net.NetworkProvider;
+import android.net.NetworkRequest;
 import android.net.NetworkScore;
 import android.net.RouteInfo;
 import android.net.TelephonyNetworkSpecifier;
@@ -546,6 +550,39 @@ public class VcnGatewayConnection extends StateMachine {
         }
     }
 
+    /**
+     * Sent when there is a suspected data stall on a network
+     *
+     * <p>Only relevant in the Connected state.
+     *
+     * @param arg1 The "all" token; this signal is always honored.
+     * @param obj @NonNull An EventDataStallSuspectedInfo instance with relevant data.
+     */
+    private static final int EVENT_DATA_STALL_SUSPECTED = 13;
+
+    private static class EventDataStallSuspectedInfo implements EventInfo {
+        @NonNull public final Network network;
+
+        EventDataStallSuspectedInfo(@NonNull Network network) {
+            this.network = network;
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(network);
+        }
+
+        @Override
+        public boolean equals(@Nullable Object other) {
+            if (!(other instanceof EventDataStallSuspectedInfo)) {
+                return false;
+            }
+
+            final EventDataStallSuspectedInfo rhs = (EventDataStallSuspectedInfo) other;
+            return Objects.equals(network, rhs.network);
+        }
+    }
+
     @VisibleForTesting(visibility = Visibility.PRIVATE)
     @NonNull
     final DisconnectedState mDisconnectedState = new DisconnectedState();
@@ -578,10 +615,13 @@ public class VcnGatewayConnection extends StateMachine {
     @NonNull
     private final VcnUnderlyingNetworkControllerCallback mUnderlyingNetworkControllerCallback;
 
+    @NonNull private final VcnConnectivityDiagnosticsCallback mConnectivityDiagnosticsCallback;
+
     private final boolean mIsMobileDataEnabled;
 
     @NonNull private final IpSecManager mIpSecManager;
     @NonNull private final ConnectivityManager mConnectivityManager;
+    @NonNull private final ConnectivityDiagnosticsManager mConnectivityDiagnosticsManager;
 
     @Nullable private IpSecTunnelInterface mTunnelIface = null;
 
@@ -748,6 +788,20 @@ public class VcnGatewayConnection extends StateMachine {
                         mUnderlyingNetworkControllerCallback);
         mIpSecManager = mVcnContext.getContext().getSystemService(IpSecManager.class);
         mConnectivityManager = mVcnContext.getContext().getSystemService(ConnectivityManager.class);
+        mConnectivityDiagnosticsManager =
+                mVcnContext.getContext().getSystemService(ConnectivityDiagnosticsManager.class);
+
+        mConnectivityDiagnosticsCallback = new VcnConnectivityDiagnosticsCallback();
+
+        if (mConnectionConfig.hasGatewayOption(
+                VCN_GATEWAY_OPTION_ENABLE_DATA_STALL_RECOVERY_WITH_MOBILITY)) {
+            final NetworkRequest diagRequest =
+                    new NetworkRequest.Builder().addTransportType(TRANSPORT_CELLULAR).build();
+            mConnectivityDiagnosticsManager.registerConnectivityDiagnosticsCallback(
+                    diagRequest,
+                    new HandlerExecutor(new Handler(vcnContext.getLooper())),
+                    mConnectivityDiagnosticsCallback);
+        }
 
         addState(mDisconnectedState);
         addState(mDisconnectingState);
@@ -810,6 +864,9 @@ public class VcnGatewayConnection extends StateMachine {
         mUnderlyingNetworkController.teardown();
 
         mGatewayStatusCallback.onQuit();
+
+        mConnectivityDiagnosticsManager.unregisterConnectivityDiagnosticsCallback(
+                mConnectivityDiagnosticsCallback);
     }
 
     /**
@@ -828,6 +885,20 @@ public class VcnGatewayConnection extends StateMachine {
         sendMessageAndAcquireWakeLock(EVENT_SUBSCRIPTIONS_CHANGED, TOKEN_ALL);
     }
 
+    private class VcnConnectivityDiagnosticsCallback extends ConnectivityDiagnosticsCallback {
+        @Override
+        public void onDataStallSuspected(ConnectivityDiagnosticsManager.DataStallReport report) {
+            mVcnContext.ensureRunningOnLooperThread();
+
+            final Network network = report.getNetwork();
+            logInfo("Data stall suspected on " + network);
+            sendMessageAndAcquireWakeLock(
+                    EVENT_DATA_STALL_SUSPECTED,
+                    TOKEN_ALL,
+                    new EventDataStallSuspectedInfo(network));
+        }
+    }
+
     private class VcnUnderlyingNetworkControllerCallback
             implements UnderlyingNetworkControllerCallback {
         @Override
@@ -1367,7 +1438,8 @@ public class VcnGatewayConnection extends StateMachine {
                 case EVENT_SUBSCRIPTIONS_CHANGED: // Fallthrough
                 case EVENT_SAFE_MODE_TIMEOUT_EXCEEDED: // Fallthrough
                 case EVENT_MIGRATION_COMPLETED: // Fallthrough
-                case EVENT_IKE_CONNECTION_INFO_CHANGED:
+                case EVENT_IKE_CONNECTION_INFO_CHANGED: // Fallthrough
+                case EVENT_DATA_STALL_SUSPECTED:
                     logUnexpectedEvent(msg.what);
                     break;
                 default:
@@ -1925,6 +1997,11 @@ public class VcnGatewayConnection extends StateMachine {
                     mIkeConnectionInfo =
                             ((EventIkeConnectionInfoChangedInfo) msg.obj).ikeConnectionInfo;
                     break;
+                case EVENT_DATA_STALL_SUSPECTED:
+                    final Network networkWithDataStall =
+                            ((EventDataStallSuspectedInfo) msg.obj).network;
+                    handleDataStallSuspected(networkWithDataStall);
+                    break;
                 default:
                     logUnhandledMessage(msg);
                     break;
@@ -1985,6 +2062,15 @@ public class VcnGatewayConnection extends StateMachine {
             }
         }
 
+        private void handleDataStallSuspected(Network networkWithDataStall) {
+            if (mUnderlying != null
+                    && mNetworkAgent != null
+                    && mNetworkAgent.getNetwork().equals(networkWithDataStall)) {
+                logInfo("Perform Mobility update to recover from suspected data stall");
+                mIkeSession.setNetwork(mUnderlying.network);
+            }
+        }
+
         protected void setupInterfaceAndNetworkAgent(
                 int token,
                 @NonNull IpSecTunnelInterface tunnelIface,
@@ -2423,6 +2509,11 @@ public class VcnGatewayConnection extends StateMachine {
         return mUnderlyingNetworkControllerCallback;
     }
 
+    @VisibleForTesting(visibility = Visibility.PRIVATE)
+    ConnectivityDiagnosticsCallback getConnectivityDiagnosticsCallback() {
+        return mConnectivityDiagnosticsCallback;
+    }
+
     @VisibleForTesting(visibility = Visibility.PRIVATE)
     UnderlyingNetworkRecord getUnderlyingNetwork() {
         return mUnderlying;

tests/vcn/java/android/net/vcn/VcnGatewayConnectionConfigTest.java

@@ -19,9 +19,11 @@ package android.net.vcn;
 import static android.net.ipsec.ike.IkeSessionParams.IKE_OPTION_MOBIKE;
 import static android.net.vcn.VcnGatewayConnectionConfig.DEFAULT_UNDERLYING_NETWORK_TEMPLATES;
 import static android.net.vcn.VcnGatewayConnectionConfig.UNDERLYING_NETWORK_TEMPLATES_KEY;
+import static android.net.vcn.VcnGatewayConnectionConfig.VCN_GATEWAY_OPTION_ENABLE_DATA_STALL_RECOVERY_WITH_MOBILITY;
 
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotEquals;
 import static org.junit.Assert.assertNotSame;
 import static org.junit.Assert.assertTrue;
@@ -42,7 +44,9 @@ import org.junit.runner.RunWith;
 
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.List;
+import java.util.Set;
 import java.util.concurrent.TimeUnit;
 
 @RunWith(AndroidJUnit4.class)
@@ -79,6 +83,9 @@ public class VcnGatewayConnectionConfigTest {
             };
     public static final int MAX_MTU = 1360;
 
+    private static final Set<Integer> GATEWAY_OPTIONS =
+            Collections.singleton(VCN_GATEWAY_OPTION_ENABLE_DATA_STALL_RECOVERY_WITH_MOBILITY);
+
     public static final IkeTunnelConnectionParams TUNNEL_CONNECTION_PARAMS =
             TunnelConnectionParamsUtilsTest.buildTestParams();
 
@@ -109,10 +116,16 @@ public class VcnGatewayConnectionConfigTest {
                 TUNNEL_CONNECTION_PARAMS);
     }
 
-    private static VcnGatewayConnectionConfig buildTestConfigWithExposedCaps(
+    private static VcnGatewayConnectionConfig buildTestConfigWithExposedCapsAndOptions(
-            VcnGatewayConnectionConfig.Builder builder, int... exposedCaps) {
+            VcnGatewayConnectionConfig.Builder builder,
+            Set<Integer> gatewayOptions,
+            int... exposedCaps) {
         builder.setRetryIntervalsMillis(RETRY_INTERVALS_MS).setMaxMtu(MAX_MTU);
 
+        for (int option : gatewayOptions) {
+            builder.addGatewayOption(option);
+        }
+
         for (int caps : exposedCaps) {
             builder.addExposedCapability(caps);
         }
@@ -120,11 +133,28 @@ public class VcnGatewayConnectionConfigTest {
         return builder.build();
     }
 
+    private static VcnGatewayConnectionConfig buildTestConfigWithExposedCaps(
+            VcnGatewayConnectionConfig.Builder builder, int... exposedCaps) {
+        return buildTestConfigWithExposedCapsAndOptions(
+                builder, Collections.emptySet(), exposedCaps);
+    }
+
     // Public for use in VcnGatewayConnectionTest
     public static VcnGatewayConnectionConfig buildTestConfigWithExposedCaps(int... exposedCaps) {
         return buildTestConfigWithExposedCaps(newBuilder(), exposedCaps);
     }
 
+    private static VcnGatewayConnectionConfig buildTestConfigWithGatewayOptions(
+            VcnGatewayConnectionConfig.Builder builder, Set<Integer> gatewayOptions) {
+        return buildTestConfigWithExposedCapsAndOptions(builder, gatewayOptions, EXPOSED_CAPS);
+    }
+
+    // Public for use in VcnGatewayConnectionTest
+    public static VcnGatewayConnectionConfig buildTestConfigWithGatewayOptions(
+            Set<Integer> gatewayOptions) {
+        return buildTestConfigWithExposedCapsAndOptions(newBuilder(), gatewayOptions, EXPOSED_CAPS);
+    }
+
     @Test
     public void testBuilderRequiresNonNullGatewayConnectionName() {
         try {
@@ -210,6 +240,15 @@ public class VcnGatewayConnectionConfigTest {
         }
     }
 
+    @Test
+    public void testBuilderRequiresValidOption() {
+        try {
+            newBuilder().addGatewayOption(-1);
+            fail("Expected exception due to the invalid VCN gateway option");
+        } catch (IllegalArgumentException e) {
+        }
+    }
+
     @Test
     public void testBuilderAndGetters() {
         final VcnGatewayConnectionConfig config = buildTestConfig();
@@ -225,6 +264,20 @@ public class VcnGatewayConnectionConfigTest {
 
         assertArrayEquals(RETRY_INTERVALS_MS, config.getRetryIntervalsMillis());
         assertEquals(MAX_MTU, config.getMaxMtu());
+
+        assertFalse(
+                config.hasGatewayOption(
+                        VCN_GATEWAY_OPTION_ENABLE_DATA_STALL_RECOVERY_WITH_MOBILITY));
+    }
+
+    @Test
+    public void testBuilderAndGettersWithOptions() {
+        final VcnGatewayConnectionConfig config =
+                buildTestConfigWithGatewayOptions(GATEWAY_OPTIONS);
+
+        for (int option : GATEWAY_OPTIONS) {
+            assertTrue(config.hasGatewayOption(option));
+        }
     }
 
     @Test
@@ -234,6 +287,14 @@ public class VcnGatewayConnectionConfigTest {
         assertEquals(config, new VcnGatewayConnectionConfig(config.toPersistableBundle()));
     }
 
+    @Test
+    public void testPersistableBundleWithOptions() {
+        final VcnGatewayConnectionConfig config =
+                buildTestConfigWithGatewayOptions(GATEWAY_OPTIONS);
+
+        assertEquals(config, new VcnGatewayConnectionConfig(config.toPersistableBundle()));
+    }
+
     @Test
     public void testParsePersistableBundleWithoutVcnUnderlyingNetworkTemplates() {
         PersistableBundle configBundle = buildTestConfig().toPersistableBundle();
@@ -318,4 +379,27 @@ public class VcnGatewayConnectionConfigTest {
         assertNotEquals(UNDERLYING_NETWORK_TEMPLATES, networkTemplatesNotEqual);
         assertNotEquals(config, configNotEqual);
     }
+
+    private static VcnGatewayConnectionConfig buildConfigWithGatewayOptionsForEqualityTest(
+            Set<Integer> gatewayOptions) {
+        return buildTestConfigWithGatewayOptions(
+                new VcnGatewayConnectionConfig.Builder(
+                        "buildConfigWithGatewayOptionsForEqualityTest", TUNNEL_CONNECTION_PARAMS),
+                gatewayOptions);
+    }
+
+    @Test
+    public void testVcnGatewayOptionsEquality() throws Exception {
+        final VcnGatewayConnectionConfig config =
+                buildConfigWithGatewayOptionsForEqualityTest(GATEWAY_OPTIONS);
+
+        final VcnGatewayConnectionConfig configEqual =
+                buildConfigWithGatewayOptionsForEqualityTest(GATEWAY_OPTIONS);
+
+        final VcnGatewayConnectionConfig configNotEqual =
+                buildConfigWithGatewayOptionsForEqualityTest(Collections.emptySet());
+
+        assertEquals(config, configEqual);
+        assertNotEquals(config, configNotEqual);
+    }
 }

tests/vcn/java/com/android/server/vcn/VcnGatewayConnectionConnectedStateTest.java

@@ -50,9 +50,11 @@ import static org.mockito.Mockito.when;
 
 import static java.util.Collections.singletonList;
 
+import android.net.ConnectivityDiagnosticsManager.DataStallReport;
 import android.net.ConnectivityManager;
 import android.net.LinkAddress;
 import android.net.LinkProperties;
+import android.net.Network;
 import android.net.NetworkAgent;
 import android.net.NetworkCapabilities;
 import android.net.ipsec.ike.ChildSaProposal;
@@ -63,10 +65,12 @@ import android.net.ipsec.ike.exceptions.IkeProtocolException;
 import android.net.vcn.VcnGatewayConnectionConfig;
 import android.net.vcn.VcnGatewayConnectionConfigTest;
 import android.net.vcn.VcnManager.VcnErrorCode;
+import android.os.PersistableBundle;
 
 import androidx.test.filters.SmallTest;
 import androidx.test.runner.AndroidJUnit4;
 
+import com.android.server.vcn.routeselection.UnderlyingNetworkRecord;
 import com.android.server.vcn.util.MtuUtils;
 
 import org.junit.Before;
@@ -88,6 +92,7 @@ import java.util.function.Consumer;
 public class VcnGatewayConnectionConnectedStateTest extends VcnGatewayConnectionTestBase {
     private VcnIkeSession mIkeSession;
     private VcnNetworkAgent mNetworkAgent;
+    private Network mVcnNetwork;
 
     @Before
     public void setUp() throws Exception {
@@ -98,6 +103,9 @@ public class VcnGatewayConnectionConnectedStateTest extends VcnGatewayConnection
                 .when(mDeps)
                 .newNetworkAgent(any(), any(), any(), any(), any(), any(), any(), any(), any());
 
+        mVcnNetwork = mock(Network.class);
+        doReturn(mVcnNetwork).when(mNetworkAgent).getNetwork();
+
         mGatewayConnection.setUnderlyingNetwork(TEST_UNDERLYING_NETWORK_RECORD_1);
 
         mIkeSession = mGatewayConnection.buildIkeSession(TEST_UNDERLYING_NETWORK_RECORD_1.network);
@@ -166,6 +174,56 @@ public class VcnGatewayConnectionConnectedStateTest extends VcnGatewayConnection
         assertEquals(mGatewayConnection.mConnectedState, mGatewayConnection.getCurrentState());
     }
 
+    private void verifyDataStallTriggersMigration(
+            UnderlyingNetworkRecord networkRecord,
+            Network networkWithDataStall,
+            boolean expectMobilityUpdate)
+            throws Exception {
+        mGatewayConnection.setUnderlyingNetwork(networkRecord);
+        triggerChildOpened();
+        mTestLooper.dispatchAll();
+
+        final DataStallReport report =
+                new DataStallReport(
+                        networkWithDataStall,
+                        1234 /* reportTimestamp */,
+                        1 /* detectionMethod */,
+                        new LinkProperties(),
+                        new NetworkCapabilities(),
+                        new PersistableBundle());
+
+        mGatewayConnection.getConnectivityDiagnosticsCallback().onDataStallSuspected(report);
+        mTestLooper.dispatchAll();
+
+        assertEquals(mGatewayConnection.mConnectedState, mGatewayConnection.getCurrentState());
+
+        if (expectMobilityUpdate) {
+            verify(mIkeSession).setNetwork(networkRecord.network);
+        } else {
+            verify(mIkeSession, never()).setNetwork(any(Network.class));
+        }
+    }
+
+    @Test
+    public void testDataStallTriggersMigration() throws Exception {
+        verifyDataStallTriggersMigration(
+                TEST_UNDERLYING_NETWORK_RECORD_1, mVcnNetwork, true /* expectMobilityUpdate */);
+    }
+
+    @Test
+    public void testDataStallWontTriggerMigrationWhenOnOtherNetwork() throws Exception {
+        verifyDataStallTriggersMigration(
+                TEST_UNDERLYING_NETWORK_RECORD_1,
+                mock(Network.class),
+                false /* expectMobilityUpdate */);
+    }
+
+    @Test
+    public void testDataStallWontTriggerMigrationWhenUnderlyingNetworkLost() throws Exception {
+        verifyDataStallTriggersMigration(
+                null /* networkRecord */, mock(Network.class), false /* expectMobilityUpdate */);
+    }
+
     private void verifyVcnTransformsApplied(
             VcnGatewayConnection vcnGatewayConnection, boolean expectForwardTransform)
             throws Exception {

tests/vcn/java/com/android/server/vcn/VcnGatewayConnectionTest.java

@@ -24,6 +24,7 @@ import static android.net.NetworkCapabilities.NET_CAPABILITY_NOT_ROAMING;
 import static android.net.NetworkCapabilities.NET_CAPABILITY_NOT_VCN_MANAGED;
 import static android.net.NetworkCapabilities.TRANSPORT_CELLULAR;
 import static android.net.NetworkCapabilities.TRANSPORT_WIFI;
+import static android.net.vcn.VcnGatewayConnectionConfig.VCN_GATEWAY_OPTION_ENABLE_DATA_STALL_RECOVERY_WITH_MOBILITY;
 
 import static com.android.server.vcn.VcnGatewayConnection.DUMMY_ADDR;
 import static com.android.server.vcn.VcnGatewayConnection.VcnChildSessionConfiguration;
@@ -34,20 +35,25 @@ import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyLong;
-import static org.mockito.Matchers.eq;
+import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.CALLS_REAL_METHODS;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 
+import android.net.ConnectivityDiagnosticsManager.ConnectivityDiagnosticsCallback;
 import android.net.IpSecManager;
 import android.net.LinkAddress;
 import android.net.LinkProperties;
 import android.net.Network;
 import android.net.NetworkCapabilities;
+import android.net.NetworkRequest;
 import android.net.TelephonyNetworkSpecifier;
+import android.net.vcn.VcnGatewayConnectionConfig;
 import android.net.vcn.VcnGatewayConnectionConfigTest;
 import android.net.vcn.VcnTransportInfo;
 import android.net.wifi.WifiInfo;
@@ -64,6 +70,7 @@ import com.android.server.vcn.routeselection.UnderlyingNetworkRecord;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
 
 import java.net.InetAddress;
 import java.util.Arrays;
@@ -71,7 +78,9 @@ import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.UUID;
+import java.util.concurrent.Executor;
 
 /** Tests for TelephonySubscriptionTracker */
 @RunWith(AndroidJUnit4.class)
@@ -287,5 +296,60 @@ public class VcnGatewayConnectionTest extends VcnGatewayConnectionTestBase {
         verify(vcnNetworkAgent).unregister();
 
         verifyWakeLockReleased();
+
+        verify(mConnDiagMgr)
+                .unregisterConnectivityDiagnosticsCallback(
+                        mGatewayConnection.getConnectivityDiagnosticsCallback());
+    }
+
+    private VcnGatewayConnection buildConnectionWithDataStallHandling(
+            boolean datatStallHandlingEnabled) throws Exception {
+        Set<Integer> options =
+                datatStallHandlingEnabled
+                        ? Collections.singleton(
+                                VCN_GATEWAY_OPTION_ENABLE_DATA_STALL_RECOVERY_WITH_MOBILITY)
+                        : Collections.emptySet();
+        final VcnGatewayConnectionConfig gatewayConfig =
+                VcnGatewayConnectionConfigTest.buildTestConfigWithGatewayOptions(options);
+        final VcnGatewayConnection gatewayConnection =
+                new VcnGatewayConnection(
+                        mVcnContext,
+                        TEST_SUB_GRP,
+                        TEST_SUBSCRIPTION_SNAPSHOT,
+                        gatewayConfig,
+                        mGatewayStatusCallback,
+                        true /* isMobileDataEnabled */,
+                        mDeps);
+        return gatewayConnection;
+    }
+
+    @Test
+    public void testDataStallHandlingEnabled() throws Exception {
+        final VcnGatewayConnection gatewayConnection =
+                buildConnectionWithDataStallHandling(true /* datatStallHandlingEnabled */);
+
+        final ArgumentCaptor<NetworkRequest> networkRequestCaptor =
+                ArgumentCaptor.forClass(NetworkRequest.class);
+        verify(mConnDiagMgr)
+                .registerConnectivityDiagnosticsCallback(
+                        networkRequestCaptor.capture(),
+                        any(Executor.class),
+                        eq(gatewayConnection.getConnectivityDiagnosticsCallback()));
+
+        final NetworkRequest nr = networkRequestCaptor.getValue();
+        final NetworkRequest expected =
+                new NetworkRequest.Builder().addTransportType(TRANSPORT_CELLULAR).build();
+        assertEquals(expected, nr);
+    }
+
+    @Test
+    public void testDataStallHandlingDisabled() throws Exception {
+        buildConnectionWithDataStallHandling(false /* datatStallHandlingEnabled */);
+
+        verify(mConnDiagMgr, never())
+                .registerConnectivityDiagnosticsCallback(
+                        any(NetworkRequest.class),
+                        any(Executor.class),
+                        any(ConnectivityDiagnosticsCallback.class));
     }
 }

tests/vcn/java/com/android/server/vcn/VcnGatewayConnectionTestBase.java

@@ -35,6 +35,7 @@ import static org.mockito.Mockito.verifyNoMoreInteractions;
 
 import android.annotation.NonNull;
 import android.content.Context;
+import android.net.ConnectivityDiagnosticsManager;
 import android.net.ConnectivityManager;
 import android.net.InetAddresses;
 import android.net.IpSecConfig;
@@ -157,6 +158,7 @@ public class VcnGatewayConnectionTestBase {
 
     @NonNull protected final IpSecService mIpSecSvc;
     @NonNull protected final ConnectivityManager mConnMgr;
+    @NonNull protected final ConnectivityDiagnosticsManager mConnDiagMgr;
 
     @NonNull protected final IkeSessionConnectionInfo mIkeConnectionInfo;
     @NonNull protected final IkeSessionConfiguration mIkeSessionConfiguration;
@@ -186,6 +188,13 @@ public class VcnGatewayConnectionTestBase {
         VcnTestUtils.setupSystemService(
                 mContext, mConnMgr, Context.CONNECTIVITY_SERVICE, ConnectivityManager.class);
 
+        mConnDiagMgr = mock(ConnectivityDiagnosticsManager.class);
+        VcnTestUtils.setupSystemService(
+                mContext,
+                mConnDiagMgr,
+                Context.CONNECTIVITY_DIAGNOSTICS_SERVICE,
+                ConnectivityDiagnosticsManager.class);
+
         mIkeConnectionInfo =
                 new IkeSessionConnectionInfo(TEST_ADDR, TEST_ADDR_2, mock(Network.class));
         mIkeSessionConfiguration = new IkeSessionConfiguration.Builder(mIkeConnectionInfo).build();