Merge "Expose API for configuring NATT Keepalive Delay."
diff --git a/api/current.txt b/api/current.txt
index 1e41011..ce64d38 100644
--- a/api/current.txt
+++ b/api/current.txt
@@ -7,6 +7,7 @@
     method @NonNull public byte[] getEapIdentity();
     method @Nullable public android.net.eap.EapSessionConfig.EapMsChapV2Config getEapMsChapV2Config();
     method @Nullable public android.net.eap.EapSessionConfig.EapSimConfig getEapSimConfig();
+    method @Nullable public android.net.eap.EapSessionConfig.EapTtlsConfig getEapTtlsConfig();
   }
 
   public static final class EapSessionConfig.Builder {
@@ -17,6 +18,7 @@
     method @NonNull public android.net.eap.EapSessionConfig.Builder setEapIdentity(@NonNull byte[]);
     method @NonNull public android.net.eap.EapSessionConfig.Builder setEapMsChapV2Config(@NonNull String, @NonNull String);
     method @NonNull public android.net.eap.EapSessionConfig.Builder setEapSimConfig(int, int);
+    method @NonNull public android.net.eap.EapSessionConfig.Builder setEapTtlsConfig(@Nullable java.security.cert.X509Certificate, @NonNull android.net.eap.EapSessionConfig);
   }
 
   public static class EapSessionConfig.EapAkaConfig extends android.net.eap.EapSessionConfig.EapMethodConfig {
@@ -43,6 +45,11 @@
     method public int getSubId();
   }
 
+  public static class EapSessionConfig.EapTtlsConfig extends android.net.eap.EapSessionConfig.EapMethodConfig {
+    method @NonNull public android.net.eap.EapSessionConfig getInnerEapSessionConfig();
+    method @Nullable public java.security.cert.X509Certificate getServerCaCert();
+  }
+
 }
 
 package android.net.ipsec.ike {
@@ -63,6 +70,7 @@
     method public void onClosedExceptionally(@NonNull android.net.ipsec.ike.exceptions.IkeException);
     method public void onIpSecTransformCreated(@NonNull android.net.IpSecTransform, int);
     method public void onIpSecTransformDeleted(@NonNull android.net.IpSecTransform, int);
+    method public default void onIpSecTransformsMigrated(@NonNull android.net.IpSecTransform, @NonNull android.net.IpSecTransform);
     method public void onOpened(@NonNull android.net.ipsec.ike.ChildSessionConfiguration);
   }
 
@@ -136,11 +144,14 @@
     method public void finalize();
     method public void kill();
     method public void openChildSession(@NonNull android.net.ipsec.ike.ChildSessionParams, @NonNull android.net.ipsec.ike.ChildSessionCallback);
+    method public void setNetwork(@NonNull android.net.Network);
   }
 
   public interface IkeSessionCallback {
     method public void onClosed();
     method public void onClosedExceptionally(@NonNull android.net.ipsec.ike.exceptions.IkeException);
+    method public default void onError(@NonNull android.net.ipsec.ike.exceptions.IkeException);
+    method public default void onIkeSessionConnectionInfoChanged(@NonNull android.net.ipsec.ike.IkeSessionConnectionInfo);
     method public void onOpened(@NonNull android.net.ipsec.ike.IkeSessionConfiguration);
   }
 
@@ -177,6 +188,7 @@
     method public boolean hasIkeOption(int);
     field public static final int IKE_OPTION_ACCEPT_ANY_REMOTE_ID = 0; // 0x0
     field public static final int IKE_OPTION_EAP_ONLY_AUTH = 1; // 0x1
+    field public static final int IKE_OPTION_MOBIKE = 2; // 0x2
   }
 
   public static final class IkeSessionParams.Builder {
@@ -350,11 +362,19 @@
   public abstract class IkeException extends java.lang.Exception {
   }
 
-  public final class IkeInternalException extends android.net.ipsec.ike.exceptions.IkeException {
+  public final class IkeInternalException extends android.net.ipsec.ike.exceptions.IkeNonProtocolException {
     ctor public IkeInternalException(@NonNull Throwable);
     ctor public IkeInternalException(@NonNull String, @NonNull Throwable);
   }
 
+  public final class IkeNetworkLostException extends android.net.ipsec.ike.exceptions.IkeNonProtocolException {
+    ctor public IkeNetworkLostException(@NonNull android.net.Network);
+    method @NonNull public android.net.Network getNetwork();
+  }
+
+  public abstract class IkeNonProtocolException extends android.net.ipsec.ike.exceptions.IkeException {
+  }
+
   public abstract class IkeProtocolException extends android.net.ipsec.ike.exceptions.IkeException {
     method public int getErrorType();
     field public static final int ERROR_TYPE_AUTHENTICATION_FAILED = 24; // 0x18
diff --git a/api/system-current.txt b/api/system-current.txt
index 684c496..095de74 100644
--- a/api/system-current.txt
+++ b/api/system-current.txt
@@ -15,7 +15,7 @@
 package android.net.ipsec.ike {
 
   public interface IkeSessionCallback {
-    method public void onError(@NonNull android.net.ipsec.ike.exceptions.IkeProtocolException);
+    method @Deprecated public default void onError(@NonNull android.net.ipsec.ike.exceptions.IkeProtocolException);
   }
 
   public final class IkeSessionParams {
@@ -40,34 +40,32 @@
 
 package android.net.ipsec.ike.ike3gpp {
 
-  public final class Ike3gppBackoffTimer extends android.net.ipsec.ike.ike3gpp.Ike3gppInfo {
+  public final class Ike3gppBackoffTimer extends android.net.ipsec.ike.ike3gpp.Ike3gppData {
     method public int getBackoffCause();
     method public byte getBackoffTimer();
-    method public int getInfoType();
+    method public int getDataType();
     field public static final int ERROR_TYPE_NETWORK_FAILURE = 10500; // 0x2904
     field public static final int ERROR_TYPE_NO_APN_SUBSCRIPTION = 9002; // 0x232a
   }
 
+  public abstract class Ike3gppData {
+    method public abstract int getDataType();
+    field public static final int DATA_TYPE_NOTIFY_BACKOFF_TIMER = 2; // 0x2
+    field public static final int DATA_TYPE_NOTIFY_N1_MODE_INFORMATION = 1; // 0x1
+  }
+
   public final class Ike3gppExtension {
-    ctor public Ike3gppExtension(@NonNull android.net.ipsec.ike.ike3gpp.Ike3gppParams, @NonNull android.net.ipsec.ike.ike3gpp.Ike3gppExtension.Ike3gppCallback);
-    method @NonNull public android.net.ipsec.ike.ike3gpp.Ike3gppExtension.Ike3gppCallback getIke3gppCallback();
+    ctor public Ike3gppExtension(@NonNull android.net.ipsec.ike.ike3gpp.Ike3gppParams, @NonNull android.net.ipsec.ike.ike3gpp.Ike3gppExtension.Ike3gppDataListener);
+    method @NonNull public android.net.ipsec.ike.ike3gpp.Ike3gppExtension.Ike3gppDataListener getIke3gppDataListener();
     method @NonNull public android.net.ipsec.ike.ike3gpp.Ike3gppParams getIke3gppParams();
   }
 
-  public abstract static class Ike3gppExtension.Ike3gppCallback {
-    ctor public Ike3gppExtension.Ike3gppCallback();
-    method public abstract void onIke3gppPayloadsReceived(@NonNull java.util.List<android.net.ipsec.ike.ike3gpp.Ike3gppInfo>);
+  public static interface Ike3gppExtension.Ike3gppDataListener {
+    method public void onIke3gppDataReceived(@NonNull java.util.List<android.net.ipsec.ike.ike3gpp.Ike3gppData>);
   }
 
-  public abstract class Ike3gppInfo {
-    ctor public Ike3gppInfo();
-    method public abstract int getInfoType();
-    field public static final int INFO_TYPE_NOTIFY_BACKOFF_TIMER = 2; // 0x2
-    field public static final int INFO_TYPE_NOTIFY_N1_MODE_INFORMATION = 1; // 0x1
-  }
-
-  public final class Ike3gppN1ModeInformation extends android.net.ipsec.ike.ike3gpp.Ike3gppInfo {
-    method public int getInfoType();
+  public final class Ike3gppN1ModeInformation extends android.net.ipsec.ike.ike3gpp.Ike3gppData {
+    method public int getDataType();
     method @NonNull public byte[] getSnssai();
   }
 
diff --git a/src/java/android/net/eap/EapSessionConfig.java b/src/java/android/net/eap/EapSessionConfig.java
index cfa2fcc..7342dd1 100644
--- a/src/java/android/net/eap/EapSessionConfig.java
+++ b/src/java/android/net/eap/EapSessionConfig.java
@@ -194,7 +194,6 @@
      * Retrieves configuration for EAP-TTLS
      *
      * @return the configuration for EAP-TTLS, or null if it was not set
-     * @hide
      */
     @Nullable
     public EapTtlsConfig getEapTtlsConfig() {
@@ -308,9 +307,11 @@
         }
 
         /**
-         * Sets the configuration for EAP-TTLS
+         * Sets the configuration for EAP-TTLS.
          *
-         * <p>Nested tunnel authentications are disallowed.
+         * <p>Tunneled EAP-TTLS authentications are disallowed, as running multiple layers of
+         * EAP-TTLS increases the data footprint but has no discernible benefits over a single
+         * EAP-TTLS session with a non EAP-TTLS method nested inside it.
          *
          * @param serverCaCert the CA certificate for validating the received server certificate(s).
          *     If a certificate is provided, it MUST be the root CA used by the server, or
@@ -318,7 +319,6 @@
          *     truststore is considered acceptable.
          * @param innerEapSessionConfig represents the configuration for the inner EAP instance
          * @return Builder this, to facilitate chaining
-         * @hide
          */
         @NonNull
         public Builder setEapTtlsConfig(
@@ -762,8 +762,6 @@
 
     /**
      * EapTtlsConfig represents the configs needed for an EAP-TTLS session.
-     *
-     * @hide
      */
     public static class EapTtlsConfig extends EapMethodConfig {
         private static final String TRUST_CERT_KEY = "TRUST_CERT_KEY";
@@ -854,7 +852,6 @@
          *
          * @return the CA certificate for validating the received server certificate or null if the
          *     system default is preferred
-         * @hide
          */
         @Nullable
         public X509Certificate getServerCaCert() {
@@ -865,7 +862,6 @@
          * Retrieves the inner EAP session config
          *
          * @return an EapSessionConfig representing the config for tunneled EAP authentication
-         * @hide
          */
         @NonNull
         public EapSessionConfig getInnerEapSessionConfig() {
diff --git a/src/java/android/net/ipsec/ike/ChildSessionCallback.java b/src/java/android/net/ipsec/ike/ChildSessionCallback.java
index 9b67a00..1dd4530 100644
--- a/src/java/android/net/ipsec/ike/ChildSessionCallback.java
+++ b/src/java/android/net/ipsec/ike/ChildSessionCallback.java
@@ -108,7 +108,6 @@
      *     {@link IpSecManager#DIRECTION_IN}
      * @param outIpSecTransform IpSecTransform to be used for traffic with {@link PolicyDirection}
      *     {@link IpSecManager#DIRECTION_OUT}
-     * @hide
      */
     default void onIpSecTransformsMigrated(
             @NonNull IpSecTransform inIpSecTransform, @NonNull IpSecTransform outIpSecTransform) {}
diff --git a/src/java/android/net/ipsec/ike/IkeSession.java b/src/java/android/net/ipsec/ike/IkeSession.java
index cab99cc..5dab372 100644
--- a/src/java/android/net/ipsec/ike/IkeSession.java
+++ b/src/java/android/net/ipsec/ike/IkeSession.java
@@ -296,7 +296,6 @@
      * @param network the Network to use for this IkeSession
      * @throws IllegalStateException if MOBIKE is not configured in IkeSessionParams, MOBIKE is not
      *     active for this IkeSession, or if the Network was not specified in IkeSessionParams.
-     * @hide
      */
     public void setNetwork(@NonNull Network network) {
         mIkeSessionStateMachine.setNetwork(network);
diff --git a/src/java/android/net/ipsec/ike/IkeSessionCallback.java b/src/java/android/net/ipsec/ike/IkeSessionCallback.java
index 21fac18..9468011 100644
--- a/src/java/android/net/ipsec/ike/IkeSessionCallback.java
+++ b/src/java/android/net/ipsec/ike/IkeSessionCallback.java
@@ -67,12 +67,13 @@
      * INVALID_MESSAGE_ID.
      *
      * @param exception the detailed error information.
+     * @deprecated Implementers should override {@link #onError(IkeException)} to handle {@link
+     *     IkeProtocolException}s instead of using this method.
      * @hide
      */
-    // TODO: b/158033037 Deprecate this method and add a public API that takes an IkeException when
-    // exposing MOBIKE APIs
     @SystemApi
-    void onError(@NonNull IkeProtocolException exception);
+    @Deprecated
+    default void onError(@NonNull IkeProtocolException exception) {}
 
     /**
      * Called if a recoverable error is encountered in an established {@link IkeSession}.
@@ -81,7 +82,6 @@
      * non-protocol errors such as the underlying {@link android.net.Network} dying.
      *
      * @param exception the detailed error information.
-     * @hide
      */
     default void onError(@NonNull IkeException exception) {
         if (exception instanceof IkeProtocolException) {
@@ -108,7 +108,6 @@
      * </ul>
      *
      * @param connectionInfo the updated IkeSessionConnectionInfo for the Session.
-     * @hide
      */
     default void onIkeSessionConnectionInfoChanged(
             @NonNull IkeSessionConnectionInfo connectionInfo) {}
diff --git a/src/java/android/net/ipsec/ike/IkeSessionParams.java b/src/java/android/net/ipsec/ike/IkeSessionParams.java
index 4148f04..0cd23d6 100644
--- a/src/java/android/net/ipsec/ike/IkeSessionParams.java
+++ b/src/java/android/net/ipsec/ike/IkeSessionParams.java
@@ -117,27 +117,39 @@
     /**
      * If set, the IKE library will attempt to enable MOBIKE for the resulting IKE Session.
      *
+     * <p>To support MOBIKE, callers must implement:
+     *
+     * <ul>
+     *   <li>{@link IkeSessionCallback#onIkeSessionConnectionInfoChanged(IkeSessionConnectionInfo)}:
+     *       this MUST migrate all IpSecTunnelInterface instances associated with this IkeSession.
+     *   <li>{@link ChildSessionCallback#onIpSecTransformsMigrated(android.net.IpSecTransform,
+     *       android.net.IpSecTransform)}: this MUST re-apply the migrated transforms to the
+     *       IpSecTunnelInterface associated with this ChildSessionCallback, via {@link
+     *       android.net.IpSecManager#applyTunnelModeTransform(
+     *       android.net.IpSecManager.IpSecTunnelInterface, int, android.net.IpSecTransform)}.
+     * </ul>
+     *
      * <p>MOBIKE support is compatible with two Network modes:
      *
      * <ul>
      *   <li><b>Caller managed:</b> The caller controls the underlying Network for the IKE Session
      *       at all times. The IKE Session will only change underlying Networks if the caller
      *       initiates it through {@link IkeSession#setNetwork(Network)}. If the caller-specified
-     *       Network dies, they will be notified via {@link
+     *       Network is lost, they will be notified via {@link
      *       IkeSessionCallback#onError(android.net.ipsec.ike.exceptions.IkeException)} with an
-     *       {@link android.net.ipsec.ike.exceptions.IkeNetworkDiedException} specifying the Network
-     *       that died.
-     *   <li><b>Platform Default:</b> The IKE Session will always track the Platform default
-     *       Network. The IKE Session will start on the Platform default Network, and any subsequent
-     *       changes to the default Network (after the IKE_AUTH exchange completes) will cause the
-     *       IKE Session's underlying Network to change. If the default Network dies with no
-     *       replacements, the caller will be notified via {@link
+     *       {@link android.net.ipsec.ike.exceptions.IkeNetworkLostException} specifying the Network
+     *       that was lost.
+     *   <li><b>Platform Default:</b> The IKE Session will always track the application default
+     *       Network. The IKE Session will start on the application default Network, and any
+     *       subsequent changes to the default Network (after the IKE_AUTH exchange completes) will
+     *       cause the IKE Session's underlying Network to change. If the default Network is lost
+     *       with no replacements, the caller will be notified via {@link
      *       IkeSessionCallback#onError(android.net.ipsec.ike.exceptions.IkeException)} with an
-     *       {@link android.net.ipsec.ike.exceptions.IkeNetworkDiedException}. The caller can either
+     *       {@link android.net.ipsec.ike.exceptions.IkeNetworkLostException}. The caller can either
      *       wait until for a new default Network to become available or they may close the Session
      *       manually via {@link IkeSession#close()}. Note that the IKE Session's maximum
      *       retransmissions may expire while waiting for a new default Network, in which case the
-     *       Session will automatically close
+     *       Session will automatically close.
      * </ul>
      *
      * <p>Use of MOBIKE in the IKE Session requires the peer to also support MOBIKE.
@@ -147,9 +159,8 @@
      *
      * <p>Checking for MOBIKE use in an IKE Session is done via {@link
      * IkeSessionConfiguration#isIkeExtensionEnabled(int)}.
-     *
-     * @hide
      */
+    // TODO(b/175416035): update docs to @link to API for migrating IpSecTunnelInterfaces
     public static final int IKE_OPTION_MOBIKE = 2;
 
     private static final int MIN_IKE_OPTION = IKE_OPTION_ACCEPT_ANY_REMOTE_ID;
diff --git a/src/java/android/net/ipsec/ike/exceptions/IkeInternalException.java b/src/java/android/net/ipsec/ike/exceptions/IkeInternalException.java
index aaa69e4..37f7293 100644
--- a/src/java/android/net/ipsec/ike/exceptions/IkeInternalException.java
+++ b/src/java/android/net/ipsec/ike/exceptions/IkeInternalException.java
@@ -23,7 +23,7 @@
  * <p>Causes may include exceptions such as {@link IpSecManager.SpiUnavailableException} when the
  * requested SPI resources failed to be allocated.
  */
-public final class IkeInternalException extends IkeException {
+public final class IkeInternalException extends IkeNonProtocolException {
     /**
      * Constructs a new exception with the specified cause.
      *
diff --git a/src/java/android/net/ipsec/ike/exceptions/IkeNetworkDiedException.java b/src/java/android/net/ipsec/ike/exceptions/IkeNetworkLostException.java
similarity index 75%
rename from src/java/android/net/ipsec/ike/exceptions/IkeNetworkDiedException.java
rename to src/java/android/net/ipsec/ike/exceptions/IkeNetworkLostException.java
index 5fdae70..d638502 100644
--- a/src/java/android/net/ipsec/ike/exceptions/IkeNetworkDiedException.java
+++ b/src/java/android/net/ipsec/ike/exceptions/IkeNetworkLostException.java
@@ -16,15 +16,20 @@
 
 package android.net.ipsec.ike.exceptions;
 
+import android.annotation.NonNull;
 import android.net.Network;
 import android.net.ipsec.ike.IkeSessionCallback;
 
 import java.util.Objects;
 
 /**
- * IkeNetworkDiedException is returned to the caller via {@link
+ * IkeNetworkLostException is returned to the caller via {@link
  * IkeSessionCallback#onError(IkeException)} if the underlying Network for the {@link IkeSession}
- * dies with no alternatives.
+ * was lost with no alternatives.
+ *
+ * <p>This Exception corresponds to {@link
+ * android.net.ConnectivityManager.NetworkCallback#onLost(android.net.Network)} being invoked for
+ * the specified underlying Network.
  *
  * <p>When the caller receives this Exception, they must either:
  *
@@ -40,21 +45,20 @@
  *       </ul>
  *   <li>close the corresponding IkeSession.
  * </ul>
- *
- * @hide
  */
-public final class IkeNetworkDiedException extends IkeNonProtocolException {
+public final class IkeNetworkLostException extends IkeNonProtocolException {
     private final Network mNetwork;
 
-    /** Constructs an IkeNetworkDiedException to indicate the specified Network died. */
-    public IkeNetworkDiedException(Network network) {
+    /** Constructs an IkeNetworkLostException to indicate the specified Network was lost. */
+    public IkeNetworkLostException(@NonNull Network network) {
         super();
         Objects.requireNonNull(network, "network is null");
 
         mNetwork = network;
     }
 
-    /** Returns the IkeSession's underlying Network that died. */
+    /** Returns the IkeSession's underlying Network that was lost. */
+    @NonNull
     public Network getNetwork() {
         return mNetwork;
     }
diff --git a/src/java/android/net/ipsec/ike/exceptions/IkeNonProtocolException.java b/src/java/android/net/ipsec/ike/exceptions/IkeNonProtocolException.java
index d8a3162..f005504 100644
--- a/src/java/android/net/ipsec/ike/exceptions/IkeNonProtocolException.java
+++ b/src/java/android/net/ipsec/ike/exceptions/IkeNonProtocolException.java
@@ -18,8 +18,6 @@
 
 /**
  * IkeNonProtocolException encapsulates all implementation-specific non-protocol IKE errors.
- *
- * @hide
  */
 public abstract class IkeNonProtocolException extends IkeException {
     /** @hide */
diff --git a/src/java/android/net/ipsec/ike/ike3gpp/Ike3gppBackoffTimer.java b/src/java/android/net/ipsec/ike/ike3gpp/Ike3gppBackoffTimer.java
index bdf5b46..30a20d5 100644
--- a/src/java/android/net/ipsec/ike/ike3gpp/Ike3gppBackoffTimer.java
+++ b/src/java/android/net/ipsec/ike/ike3gpp/Ike3gppBackoffTimer.java
@@ -33,7 +33,7 @@
  * @hide
  */
 @SystemApi
-public final class Ike3gppBackoffTimer extends Ike3gppInfo {
+public final class Ike3gppBackoffTimer extends Ike3gppData {
     /**
      * Error-Notify indicating that access is not authorized because no subscription was found for
      * the specified APN.
@@ -83,8 +83,8 @@
     }
 
     @Override
-    public @InfoType int getInfoType() {
-        return INFO_TYPE_NOTIFY_BACKOFF_TIMER;
+    public @DataType int getDataType() {
+        return DATA_TYPE_NOTIFY_BACKOFF_TIMER;
     }
 
     /**
diff --git a/src/java/android/net/ipsec/ike/ike3gpp/Ike3gppData.java b/src/java/android/net/ipsec/ike/ike3gpp/Ike3gppData.java
new file mode 100644
index 0000000..3aae560
--- /dev/null
+++ b/src/java/android/net/ipsec/ike/ike3gpp/Ike3gppData.java
@@ -0,0 +1,56 @@
+/*
+ * Copyright (C) 2020 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package android.net.ipsec.ike.ike3gpp;
+
+import android.annotation.IntDef;
+import android.annotation.SystemApi;
+
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+
+/**
+ * Ike3gppData represents 3GPP-specific data sent by the peer/remote endpoint.
+ *
+ * @see 3GPP ETSI TS 24.302: Access to the 3GPP Evolved Packet Core (EPC) via non-3GPP access
+ *     networks
+ * @hide
+ */
+@SystemApi
+public abstract class Ike3gppData {
+    private static final int DATA_TYPE_SHARED_BASE = 0;
+    private static final int DATA_TYPE_CATEGORY_SIZE = 100;
+
+    private static final int DATA_TYPE_PAYLOAD_NOTIFY_BASE = DATA_TYPE_SHARED_BASE;
+
+    /** Data Type representing an {@link Ike3gppN1ModeInformation}. */
+    public static final int DATA_TYPE_NOTIFY_N1_MODE_INFORMATION =
+            DATA_TYPE_PAYLOAD_NOTIFY_BASE + 1;
+
+    /** Data Type representing an {@link Ike3gppBackoffTimer}. */
+    public static final int DATA_TYPE_NOTIFY_BACKOFF_TIMER = DATA_TYPE_PAYLOAD_NOTIFY_BASE + 2;
+
+    /** @hide */
+    @Retention(RetentionPolicy.SOURCE)
+    @IntDef({DATA_TYPE_NOTIFY_N1_MODE_INFORMATION, DATA_TYPE_NOTIFY_BACKOFF_TIMER})
+    public @interface DataType {}
+
+    /** @hide */
+    protected Ike3gppData() {}
+
+    /** Returns the DataType that this Ike3gppData represents. */
+    public abstract @DataType int getDataType();
+}
diff --git a/src/java/android/net/ipsec/ike/ike3gpp/Ike3gppExtension.java b/src/java/android/net/ipsec/ike/ike3gpp/Ike3gppExtension.java
index 03274db..2165e3f 100644
--- a/src/java/android/net/ipsec/ike/ike3gpp/Ike3gppExtension.java
+++ b/src/java/android/net/ipsec/ike/ike3gpp/Ike3gppExtension.java
@@ -36,33 +36,34 @@
 @SystemApi
 public final class Ike3gppExtension {
     @NonNull private final Ike3gppParams mIke3gppParams;
-    @NonNull private final Ike3gppCallback mIke3gppCallback;
+    @NonNull private final Ike3gppDataListener mIke3gppDataListener;
 
     /**
-     * Constructs an Ike3gppExtension instance with the given Ike3gppCallback and Ike3gppParams
+     * Constructs an Ike3gppExtension instance with the given Ike3gppDataListener and Ike3gppParams
      * instances.
      *
      * @param ike3gppParams Ike3gppParams used to configure the 3GPP-support for an IKE Session.
-     * @param ike3gppCallback Ike3gppCallback used to notify the caller of 3GPP-specific payloads
-     *     received during an IKE Session.
+     * @param ike3gppDataListener Ike3gppDataListener used to notify the caller of 3GPP-specific
+     *     data received during an IKE Session.
      */
-    // ExecutorRegistration: Not necessary to take an Executor for invoking the callback here, as
-    // this is not actually where the callback is registered. The caller's Executor provided in the
-    // IkeSession constructor will be used to invoke the Ike3gppCallback.
+    // ExecutorRegistration: Not necessary to take an Executor for invoking the listener here, as
+    // this is not actually where the listener is registered. The caller's Executor provided in the
+    // IkeSession constructor will be used to invoke the Ike3gppDataListener.
     @SuppressLint("ExecutorRegistration")
     public Ike3gppExtension(
-            @NonNull Ike3gppParams ike3gppParams, @NonNull Ike3gppCallback ike3gppCallback) {
+            @NonNull Ike3gppParams ike3gppParams,
+            @NonNull Ike3gppDataListener ike3gppDataListener) {
         Objects.requireNonNull(ike3gppParams, "ike3gppParams must not be null");
-        Objects.requireNonNull(ike3gppCallback, "ike3gppCallback must not be null");
+        Objects.requireNonNull(ike3gppDataListener, "ike3gppDataListener must not be null");
 
         mIke3gppParams = ike3gppParams;
-        mIke3gppCallback = ike3gppCallback;
+        mIke3gppDataListener = ike3gppDataListener;
     }
 
-    /** Retrieves the configured Ike3gppCallback. */
+    /** Retrieves the configured Ike3gppDataListener. */
     @NonNull
-    public Ike3gppCallback getIke3gppCallback() {
-        return mIke3gppCallback;
+    public Ike3gppDataListener getIke3gppDataListener() {
+        return mIke3gppDataListener;
     }
 
     /** Retrieves the configured Ike3gppParams. */
@@ -73,7 +74,7 @@
 
     @Override
     public int hashCode() {
-        return Objects.hash(mIke3gppParams, mIke3gppCallback);
+        return Objects.hash(mIke3gppParams, mIke3gppDataListener);
     }
 
     @Override
@@ -85,26 +86,26 @@
         Ike3gppExtension other = (Ike3gppExtension) o;
 
         return mIke3gppParams.equals(other.mIke3gppParams)
-                && mIke3gppCallback.equals(other.mIke3gppCallback);
+                && mIke3gppDataListener.equals(other.mIke3gppDataListener);
     }
 
     /**
-     * Callback for receiving 3GPP-specific payloads.
+     * Listener for receiving 3GPP-specific data.
      *
      * <p>MUST be unique to each IKE Session.
      *
-     * <p>All Ike3gppCallback calls will be invoked on the Executor provided in the IkeSession
+     * <p>All Ike3gppDataListener calls will be invoked on the Executor provided in the IkeSession
      * constructor.
      */
-    public abstract static class Ike3gppCallback {
+    public interface Ike3gppDataListener {
         /**
-         * Invoked when the IKE Session receives one or more 3GPP-specific payloads.
+         * Invoked when the IKE Session receives 3GPP-specific data.
          *
          * <p>This function will be invoked at most once for each IKE Message received by the IKEv2
          * library.
          *
-         * @param payloads List<Ike3gppInfo> the 3GPP-payloads received
+         * @param ike3gppDataList List<Ike3gppData> the 3GPP-data received
          */
-        public abstract void onIke3gppPayloadsReceived(@NonNull List<Ike3gppInfo> payloads);
+        void onIke3gppDataReceived(@NonNull List<Ike3gppData> ike3gppDataList);
     }
 }
diff --git a/src/java/android/net/ipsec/ike/ike3gpp/Ike3gppInfo.java b/src/java/android/net/ipsec/ike/ike3gpp/Ike3gppInfo.java
deleted file mode 100644
index 2379472..0000000
--- a/src/java/android/net/ipsec/ike/ike3gpp/Ike3gppInfo.java
+++ /dev/null
@@ -1,53 +0,0 @@
-/*
- * Copyright (C) 2020 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package android.net.ipsec.ike.ike3gpp;
-
-import android.annotation.IntDef;
-import android.annotation.SystemApi;
-
-import java.lang.annotation.Retention;
-import java.lang.annotation.RetentionPolicy;
-
-/**
- * Ike3gppInfo represents a 3GPP-specific payload sent by the peer/remote endpoint.
- *
- * @see 3GPP ETSI TS 24.302: Access to the 3GPP Evolved Packet Core (EPC) via non-3GPP access
- *     networks
- * @hide
- */
-@SystemApi
-public abstract class Ike3gppInfo {
-    private static final int INFO_TYPE_SHARED_BASE = 0;
-    private static final int INFO_TYPE_CATEGORY_SIZE = 100;
-
-    private static final int INFO_TYPE_PAYLOAD_NOTIFY_BASE = INFO_TYPE_SHARED_BASE;
-
-    /** Info Type representing an {@link Ike3gppN1ModeInformation}. */
-    public static final int INFO_TYPE_NOTIFY_N1_MODE_INFORMATION =
-            INFO_TYPE_PAYLOAD_NOTIFY_BASE + 1;
-
-    /** Info Type representing an {@link Ike3gppBackoffTimer}. */
-    public static final int INFO_TYPE_NOTIFY_BACKOFF_TIMER = INFO_TYPE_PAYLOAD_NOTIFY_BASE + 2;
-
-    /** @hide */
-    @Retention(RetentionPolicy.SOURCE)
-    @IntDef({INFO_TYPE_NOTIFY_N1_MODE_INFORMATION, INFO_TYPE_NOTIFY_BACKOFF_TIMER})
-    public @interface InfoType {}
-
-    /** Returns the InfoType that this Ike3gppInfo represents. */
-    public abstract @InfoType int getInfoType();
-}
diff --git a/src/java/android/net/ipsec/ike/ike3gpp/Ike3gppN1ModeInformation.java b/src/java/android/net/ipsec/ike/ike3gpp/Ike3gppN1ModeInformation.java
index 8c538ab..bfa9e44 100644
--- a/src/java/android/net/ipsec/ike/ike3gpp/Ike3gppN1ModeInformation.java
+++ b/src/java/android/net/ipsec/ike/ike3gpp/Ike3gppN1ModeInformation.java
@@ -29,7 +29,7 @@
  * @hide
  */
 @SystemApi
-public final class Ike3gppN1ModeInformation extends Ike3gppInfo {
+public final class Ike3gppN1ModeInformation extends Ike3gppData {
     private final byte[] mSnssai;
 
     /** @hide */
@@ -39,8 +39,8 @@
     }
 
     @Override
-    public @InfoType int getInfoType() {
-        return INFO_TYPE_NOTIFY_N1_MODE_INFORMATION;
+    public @DataType int getDataType() {
+        return DATA_TYPE_NOTIFY_N1_MODE_INFORMATION;
     }
 
     /**
diff --git a/src/java/com/android/internal/net/ipsec/ike/IkeSessionStateMachine.java b/src/java/com/android/internal/net/ipsec/ike/IkeSessionStateMachine.java
index 1df255f..132c519 100644
--- a/src/java/com/android/internal/net/ipsec/ike/IkeSessionStateMachine.java
+++ b/src/java/com/android/internal/net/ipsec/ike/IkeSessionStateMachine.java
@@ -32,6 +32,7 @@
 import static com.android.internal.net.ipsec.ike.message.IkeMessage.DECODE_STATUS_PARTIAL;
 import static com.android.internal.net.ipsec.ike.message.IkeMessage.DECODE_STATUS_PROTECTED_ERROR;
 import static com.android.internal.net.ipsec.ike.message.IkeMessage.DECODE_STATUS_UNPROTECTED_ERROR;
+import static com.android.internal.net.ipsec.ike.message.IkeNotifyPayload.NOTIFY_TYPE_COOKIE;
 import static com.android.internal.net.ipsec.ike.message.IkeNotifyPayload.NOTIFY_TYPE_COOKIE2;
 import static com.android.internal.net.ipsec.ike.message.IkeNotifyPayload.NOTIFY_TYPE_EAP_ONLY_AUTHENTICATION;
 import static com.android.internal.net.ipsec.ike.message.IkeNotifyPayload.NOTIFY_TYPE_IKEV2_FRAGMENTATION_SUPPORTED;
@@ -68,6 +69,7 @@
 import android.net.IpSecManager.ResourceUnavailableException;
 import android.net.IpSecManager.SpiUnavailableException;
 import android.net.IpSecManager.UdpEncapsulationSocket;
+import android.net.LinkProperties;
 import android.net.Network;
 import android.net.NetworkCapabilities;
 import android.net.NetworkRequest;
@@ -86,7 +88,7 @@
 import android.net.ipsec.ike.exceptions.AuthenticationFailedException;
 import android.net.ipsec.ike.exceptions.IkeException;
 import android.net.ipsec.ike.exceptions.IkeInternalException;
-import android.net.ipsec.ike.exceptions.IkeNetworkDiedException;
+import android.net.ipsec.ike.exceptions.IkeNetworkLostException;
 import android.net.ipsec.ike.exceptions.IkeProtocolException;
 import android.net.ipsec.ike.exceptions.InvalidKeException;
 import android.net.ipsec.ike.exceptions.InvalidSyntaxException;
@@ -161,6 +163,7 @@
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.net.Inet4Address;
+import java.net.Inet6Address;
 import java.net.InetAddress;
 import java.nio.ByteBuffer;
 import java.security.GeneralSecurityException;
@@ -434,6 +437,11 @@
     /** Local port assigned on device. Initialized in Initial State. */
     @VisibleForTesting int mLocalPort;
 
+    /** Available remote addresses that are v4. Resolved in Initial State. */
+    @VisibleForTesting final List<Inet4Address> mRemoteAddressesV4 = new ArrayList<>();
+    /** Available remote addresses that are v6. Resolved in Initial State. */
+    @VisibleForTesting final List<Inet6Address> mRemoteAddressesV6 = new ArrayList<>();
+
     /** Indicates if both sides support NAT traversal. Set in IKE INIT. */
     @VisibleForTesting boolean mSupportNatTraversal;
     /** Indicates if local node is behind a NAT. */
@@ -1203,9 +1211,21 @@
         @Override
         public void enterState() {
             try {
-                // TODO(b/149954916): Do DNS resolution asynchronously and support resolving
-                // multiple addresses.
-                mRemoteAddress = mNetwork.getByName(mIkeSessionParams.getServerHostname());
+                // TODO(b/149954916): Do DNS resolution asynchronously
+                InetAddress[] allRemoteAddresses =
+                        mNetwork.getAllByName(mIkeSessionParams.getServerHostname());
+
+                logd("Resolved addresses for peer: " + Arrays.toString(allRemoteAddresses));
+
+                for (InetAddress remoteAddress : allRemoteAddresses) {
+                    if (remoteAddress instanceof Inet4Address) {
+                        mRemoteAddressesV4.add((Inet4Address) remoteAddress);
+                    } else {
+                        mRemoteAddressesV6.add((Inet6Address) remoteAddress);
+                    }
+                }
+
+                setRemoteAddress();
 
                 boolean isIpv4 = mRemoteAddress instanceof Inet4Address;
                 if (isIpv4) {
@@ -1243,6 +1263,33 @@
     }
 
     /**
+     * Set the remote address for the peer.
+     *
+     * <p>Prefers IPv6 addresses if:
+     *
+     * <ul>
+     *   <li>an IPv6 address is known for the peer, and
+     *   <li>the current underlying Network has a global (non-link local) IPv6 address available
+     * </ul>
+     *
+     * Otherwise, an IPv4 address will be used.
+     */
+    private void setRemoteAddress() {
+        LinkProperties linkProperties = mConnectivityManager.getLinkProperties(mNetwork);
+        if (!mRemoteAddressesV6.isEmpty() && linkProperties.hasGlobalIpv6Address()) {
+            // TODO(b/175348096): randomly choose from available addresses
+            mRemoteAddress = mRemoteAddressesV6.get(0);
+        } else {
+            if (mRemoteAddressesV4.isEmpty()) {
+                throw new IllegalArgumentException("No valid IPv4 or IPv6 addresses for peer");
+            }
+
+            // TODO(b/175348096): randomly choose from available addresses
+            mRemoteAddress = mRemoteAddressesV4.get(0);
+        }
+    }
+
+    /**
      * Idle represents a state when there is no ongoing IKE exchange affecting established IKE SA.
      */
     class Idle extends LocalRequestQueuer {
@@ -1444,7 +1491,7 @@
         intent.putExtras(bundle);
 
         return PendingIntent.getBroadcast(
-                context, 0 /*requestCode unused*/, intent, 0 /*default flags*/);
+                context, 0 /* requestCode; unused */, intent, PendingIntent.FLAG_IMMUTABLE);
     }
 
     /**
@@ -2026,7 +2073,7 @@
                             IkeNotifyPayload notify = (IkeNotifyPayload) payload;
                             if (notify.notifyType == NOTIFY_TYPE_COOKIE2) {
                                 infoPayloadList.add(
-                                        IkeNotifyPayload.handleCookie2AndGenerateResponse(notify));
+                                        IkeNotifyPayload.handleCookie2AndGenerateCopy(notify));
                             }
 
                             // No action for other notifications
@@ -2884,23 +2931,27 @@
         @Override
         public void enterState() {
             try {
-                IkeMessage request = buildIkeInitReq();
-
-                // Register local SPI to receive the IKE INIT response.
-                mIkeSocket.registerIke(
-                        request.ikeHeader.ikeInitiatorSpi, IkeSessionStateMachine.this);
-
-                mIkeInitRequestBytes = request.encode();
-                mIkeInitNoncePayload =
-                        request.getPayloadForType(
-                                IkePayload.PAYLOAD_TYPE_NONCE, IkeNoncePayload.class);
-                mRetransmitter = new UnencryptedRetransmitter(request);
+                sendRequest(buildIkeInitReq());
             } catch (IOException e) {
                 // Fail to assign IKE SPI
                 handleIkeFatalError(e);
             }
         }
 
+        private void sendRequest(IkeMessage request) {
+            // Register local SPI to receive the IKE INIT response.
+            mIkeSocket.registerIke(request.ikeHeader.ikeInitiatorSpi, IkeSessionStateMachine.this);
+
+            mIkeInitRequestBytes = request.encode();
+            mIkeInitNoncePayload =
+                    request.getPayloadForType(IkePayload.PAYLOAD_TYPE_NONCE, IkeNoncePayload.class);
+
+            if (mRetransmitter != null) {
+                mRetransmitter.stopRetransmitting();
+            }
+            mRetransmitter = new UnencryptedRetransmitter(request);
+        }
+
         @Override
         protected void triggerRetransmit() {
             mRetransmitter.retransmit();
@@ -2977,10 +3028,41 @@
             }
         }
 
+        /** Returns the Notify-Cookie payload, or null if it does not exist */
+        private IkeNotifyPayload getNotifyCookie(IkeMessage ikeMessage) {
+            List<IkeNotifyPayload> notifyPayloads =
+                    ikeMessage.getPayloadListForType(PAYLOAD_TYPE_NOTIFY, IkeNotifyPayload.class);
+            for (IkeNotifyPayload notify : notifyPayloads) {
+                if (notify.notifyType == NOTIFY_TYPE_COOKIE) {
+                    return notify;
+                }
+            }
+            return null;
+        }
+
         @Override
         protected void handleResponseIkeMessage(IkeMessage ikeMessage) {
             boolean ikeInitSuccess = false;
             try {
+                int exchangeType = ikeMessage.ikeHeader.exchangeType;
+                if (exchangeType != IkeHeader.EXCHANGE_TYPE_IKE_SA_INIT) {
+                    throw new InvalidSyntaxException(
+                            "Expected EXCHANGE_TYPE_IKE_SA_INIT but received: " + exchangeType);
+                }
+
+                // Retry IKE INIT if there is Notify-Cookie
+                IkeNotifyPayload inCookiePayload = getNotifyCookie(ikeMessage);
+                if (inCookiePayload != null) {
+                    IkeNotifyPayload outCookiePayload =
+                            IkeNotifyPayload.handleCookieAndGenerateCopy(inCookiePayload);
+                    IkeMessage initReq =
+                            buildReqWithCookie(mRetransmitter.getMessage(), outCookiePayload);
+
+                    sendRequest(initReq);
+                    return;
+                }
+
+                // Negotiate IKE SA
                 validateIkeInitResp(mRetransmitter.getMessage(), ikeMessage);
 
                 mCurrentIkeSaRecord =
@@ -3094,18 +3176,45 @@
             return new IkeMessage(ikeHeader, payloadList);
         }
 
+        /**
+         * Builds an IKE INIT request that has the same payloads and SPI with the original request,
+         * and with the new Notify-Cookie Payload as the first payload.
+         */
+        private IkeMessage buildReqWithCookie(
+                IkeMessage originalReq, IkeNotifyPayload cookieNotify) {
+            List<IkePayload> payloads = new ArrayList<>();
+
+            // Notify-Cookie MUST be the first payload.
+            payloads.add(cookieNotify);
+
+            for (IkePayload payload : originalReq.ikePayloadList) {
+                // Keep all previous payloads except COOKIEs
+                if (payload instanceof IkeNotifyPayload
+                        && ((IkeNotifyPayload) payload).notifyType == NOTIFY_TYPE_COOKIE) {
+                    continue;
+                }
+                payloads.add(payload);
+            }
+
+            IkeHeader originalHeader = originalReq.ikeHeader;
+            IkeHeader header =
+                    new IkeHeader(
+                            originalHeader.ikeInitiatorSpi,
+                            originalHeader.ikeResponderSpi,
+                            PAYLOAD_TYPE_NOTIFY,
+                            IkeHeader.EXCHANGE_TYPE_IKE_SA_INIT,
+                            false /* isResponseMsg */,
+                            true /* fromIkeInitiator */,
+                            0 /* messageId */);
+            return new IkeMessage(header, payloads);
+        }
+
         private void validateIkeInitResp(IkeMessage reqMsg, IkeMessage respMsg)
                 throws IkeProtocolException, IOException {
             IkeHeader respIkeHeader = respMsg.ikeHeader;
             mRemoteIkeSpiResource =
                     mIkeSpiGenerator.allocateSpi(mRemoteAddress, respIkeHeader.ikeResponderSpi);
 
-            int exchangeType = respIkeHeader.exchangeType;
-            if (exchangeType != IkeHeader.EXCHANGE_TYPE_IKE_SA_INIT) {
-                throw new InvalidSyntaxException(
-                        "Expected EXCHANGE_TYPE_IKE_SA_INIT but received: " + exchangeType);
-            }
-
             IkeSaPayload respSaPayload = null;
             IkeKePayload respKePayload = null;
 
@@ -5406,11 +5515,13 @@
             throw new IllegalStateException("MOBIKE must be enabled to update the Network");
         }
 
-        // TODO(b/172060298): prefer IPv6 once Responder addresses are cached
-        boolean isIpv4 = mRemoteAddress instanceof Inet4Address;
         Network oldNetwork = mNetwork;
         mNetwork = network;
 
+        setRemoteAddress();
+
+        boolean isIpv4 = mRemoteAddress instanceof Inet4Address;
+
         try {
             // Only switch the IkeSocket if the underlying Network actually changes. This may not
             // always happen (ex: the underlying Network loses the current local address)
@@ -5466,6 +5577,6 @@
     @Override
     public void onUnderlyingNetworkDied() {
         executeUserCallback(
-                () -> mIkeSessionCallback.onError(new IkeNetworkDiedException(mNetwork)));
+                () -> mIkeSessionCallback.onError(new IkeNetworkLostException(mNetwork)));
     }
 }
diff --git a/src/java/com/android/internal/net/ipsec/ike/ike3gpp/Ike3gppExchangeBase.java b/src/java/com/android/internal/net/ipsec/ike/ike3gpp/Ike3gppExchangeBase.java
index 1e67ba7..88aa114 100644
--- a/src/java/com/android/internal/net/ipsec/ike/ike3gpp/Ike3gppExchangeBase.java
+++ b/src/java/com/android/internal/net/ipsec/ike/ike3gpp/Ike3gppExchangeBase.java
@@ -18,8 +18,8 @@
 import static android.net.ipsec.ike.IkeManager.getIkeLog;
 
 import android.annotation.NonNull;
+import android.net.ipsec.ike.ike3gpp.Ike3gppData;
 import android.net.ipsec.ike.ike3gpp.Ike3gppExtension;
-import android.net.ipsec.ike.ike3gpp.Ike3gppInfo;
 
 import java.util.List;
 import java.util.Objects;
@@ -44,17 +44,17 @@
         mUserCbExecutor = Objects.requireNonNull(userCbExecutor, "userCbExecutor must not be null");
     }
 
-    void maybeInvokeUserCallback(List<Ike3gppInfo> ike3gppInfos) {
-        if (ike3gppInfos.isEmpty()) return;
+    void maybeInvokeUserCallback(List<Ike3gppData> ike3gppDataList) {
+        if (ike3gppDataList.isEmpty()) return;
 
         try {
             mUserCbExecutor.execute(
                     () ->
                             mIke3gppExtension
-                                    .getIke3gppCallback()
-                                    .onIke3gppPayloadsReceived(ike3gppInfos));
+                                    .getIke3gppDataListener()
+                                    .onIke3gppDataReceived(ike3gppDataList));
         } catch (Exception e) {
-            getIkeLog().d(TAG, "Ike3gppCallback#onIke3gppPayloadsReceived execution failed", e);
+            getIkeLog().d(TAG, "Ike3gppDataListener#onIke3gppDataReceived execution failed", e);
         }
     }
 }
diff --git a/src/java/com/android/internal/net/ipsec/ike/ike3gpp/Ike3gppExtensionExchange.java b/src/java/com/android/internal/net/ipsec/ike/ike3gpp/Ike3gppExtensionExchange.java
index bd14086..152fea9 100644
--- a/src/java/com/android/internal/net/ipsec/ike/ike3gpp/Ike3gppExtensionExchange.java
+++ b/src/java/com/android/internal/net/ipsec/ike/ike3gpp/Ike3gppExtensionExchange.java
@@ -23,7 +23,7 @@
 import android.annotation.Nullable;
 import android.net.ipsec.ike.exceptions.InvalidSyntaxException;
 import android.net.ipsec.ike.ike3gpp.Ike3gppExtension;
-import android.net.ipsec.ike.ike3gpp.Ike3gppExtension.Ike3gppCallback;
+import android.net.ipsec.ike.ike3gpp.Ike3gppExtension.Ike3gppDataListener;
 import android.util.ArraySet;
 
 import com.android.internal.net.ipsec.ike.IkeSessionStateMachine;
@@ -39,9 +39,9 @@
  * Ike3gppExtensionExchange contains the implementation for 3GPP-specific functionality in IKEv2.
  */
 public class Ike3gppExtensionExchange implements AutoCloseable {
-    private static final String TAG = Ike3gppExtension.class.getSimpleName();
+    private static final String TAG = Ike3gppExtensionExchange.class.getSimpleName();
 
-    private static final Set<Ike3gppCallback> REGISTERED_CALLBACKS =
+    private static final Set<Ike3gppDataListener> REGISTERED_LISTENERS =
             Collections.synchronizedSet(new ArraySet<>());
 
     /**
@@ -100,9 +100,9 @@
         if (mIke3gppExtension != null) {
             mIke3gppIkeAuth = new Ike3gppIkeAuth(mIke3gppExtension, mUserCbExecutor);
 
-            if (!REGISTERED_CALLBACKS.add(ike3gppExtension.getIke3gppCallback())) {
+            if (!REGISTERED_LISTENERS.add(ike3gppExtension.getIke3gppDataListener())) {
                 throw new IllegalArgumentException(
-                        "Ike3gppCallback must be unique for each IkeSession");
+                        "Ike3gppDataListener must be unique for each IkeSession");
             }
         } else {
             mIke3gppIkeAuth = null;
@@ -113,7 +113,7 @@
     public void close() {
         if (mIke3gppExtension == null) return;
 
-        REGISTERED_CALLBACKS.remove(mIke3gppExtension.getIke3gppCallback());
+        REGISTERED_LISTENERS.remove(mIke3gppExtension.getIke3gppDataListener());
     }
 
     /** Gets the 3GPP-specific Request IkePayloads for the specified exchangeSubtype. */
@@ -155,8 +155,8 @@
     /**
      * Handles the provided Response IkePayloads for the specified exchangeSubtype.
      *
-     * <p>If the caller needs to be notified of received Ike3gppInfos, the configured
-     * Ike3gppCallback will be invoked.
+     * <p>If the caller needs to be notified of received Ike3gppData, the configured
+     * Ike3gppDataListener will be invoked.
      */
     public void handle3gppResponsePayloads(int exchangeSubtype, List<IkePayload> ike3gppPayloads)
             throws InvalidSyntaxException {
diff --git a/src/java/com/android/internal/net/ipsec/ike/ike3gpp/Ike3gppIkeAuth.java b/src/java/com/android/internal/net/ipsec/ike/ike3gpp/Ike3gppIkeAuth.java
index 1919375..f3672f1 100644
--- a/src/java/com/android/internal/net/ipsec/ike/ike3gpp/Ike3gppIkeAuth.java
+++ b/src/java/com/android/internal/net/ipsec/ike/ike3gpp/Ike3gppIkeAuth.java
@@ -25,8 +25,8 @@
 import android.annotation.NonNull;
 import android.net.ipsec.ike.exceptions.InvalidSyntaxException;
 import android.net.ipsec.ike.ike3gpp.Ike3gppBackoffTimer;
+import android.net.ipsec.ike.ike3gpp.Ike3gppData;
 import android.net.ipsec.ike.ike3gpp.Ike3gppExtension;
-import android.net.ipsec.ike.ike3gpp.Ike3gppInfo;
 import android.net.ipsec.ike.ike3gpp.Ike3gppN1ModeInformation;
 import android.util.ArraySet;
 
@@ -91,7 +91,7 @@
     }
 
     void handleAuthResp(List<IkePayload> ike3gppPayloads) throws InvalidSyntaxException {
-        List<Ike3gppInfo> ike3gppInfos = new ArrayList<>();
+        List<Ike3gppData> ike3gppDataList = new ArrayList<>();
         List<IkeNotifyPayload> notifyPayloads =
                 IkePayload.getPayloadListForTypeInProvidedList(
                         IkePayload.PAYLOAD_TYPE_NOTIFY, IkeNotifyPayload.class, ike3gppPayloads);
@@ -109,7 +109,7 @@
 
                     byte[] snssai =
                             Ike3gppN1ModeUtils.getSnssaiFromNotifyData(notifyPayload.notifyData);
-                    ike3gppInfos.add(new Ike3gppN1ModeInformation(snssai));
+                    ike3gppDataList.add(new Ike3gppN1ModeInformation(snssai));
                     break;
                 case NOTIFY_TYPE_BACKOFF_TIMER:
                     backoffTimerPayload = notifyPayload;
@@ -135,12 +135,13 @@
             byte backoffTimer =
                     Ike3gppBackoffTimerUtils.getBackoffTimerfromNotifyData(
                             backoffTimerPayload.notifyData);
-            ike3gppInfos.add(new Ike3gppBackoffTimer(backoffTimer, backoffTimerCause.notifyType));
+            ike3gppDataList.add(
+                    new Ike3gppBackoffTimer(backoffTimer, backoffTimerCause.notifyType));
         } else if (backoffTimerPayload != null) {
             logw("Received BACKOFF_TIMER payload without an Error-Notify");
         }
 
-        maybeInvokeUserCallback(ike3gppInfos);
+        maybeInvokeUserCallback(ike3gppDataList);
     }
 
     private void logd(String msg) {
diff --git a/src/java/com/android/internal/net/ipsec/ike/message/IkeNotifyPayload.java b/src/java/com/android/internal/net/ipsec/ike/message/IkeNotifyPayload.java
index 4617932..728929b 100644
--- a/src/java/com/android/internal/net/ipsec/ike/message/IkeNotifyPayload.java
+++ b/src/java/com/android/internal/net/ipsec/ike/message/IkeNotifyPayload.java
@@ -120,6 +120,11 @@
      */
     public static final int NOTIFY_TYPE_NAT_DETECTION_DESTINATION_IP = 16389;
     /**
+     * Might be sent by the IKE responder in an IKE_SA_INIT response, to prevent DoS Attacks. If
+     * receiving it, IKE client MUST retry IKE_SA_INIT request with the same associated data.
+     */
+    public static final int NOTIFY_TYPE_COOKIE = 16390;
+    /**
      * Indicates a willingness by its sender to use transport mode rather than tunnel mode on this
      * Child SA. Only allowed in the request/response for negotiating a Child SA.
      */
@@ -171,6 +176,9 @@
 
     private static final String NAT_DETECTION_DIGEST_ALGORITHM = "SHA-1";
 
+    private static final int COOKIE_DATA_LEN_MIN = 1;
+    private static final int COOKIE_DATA_LEN_MAX = 64;
+
     private static final int COOKIE2_DATA_LEN_MIN = 8;
     private static final int COOKIE2_DATA_LEN_MAX = 64;
 
@@ -229,6 +237,7 @@
         NOTIFY_TYPE_TO_STRING.put(NOTIFY_TYPE_NAT_DETECTION_SOURCE_IP, "NAT detection source IP");
         NOTIFY_TYPE_TO_STRING.put(
                 NOTIFY_TYPE_NAT_DETECTION_DESTINATION_IP, "NAT detection destination IP");
+        NOTIFY_TYPE_TO_STRING.put(NOTIFY_TYPE_COOKIE, "COOKIE");
         NOTIFY_TYPE_TO_STRING.put(NOTIFY_TYPE_USE_TRANSPORT_MODE, "Use transport mode");
         NOTIFY_TYPE_TO_STRING.put(NOTIFY_TYPE_REKEY_SA, "Rekey SA");
         NOTIFY_TYPE_TO_STRING.put(
@@ -348,16 +357,33 @@
         }
     }
 
-    /** Validate inbound Cookie2 request and build a response Cookie2 notify payload */
-    public static IkeNotifyPayload handleCookie2AndGenerateResponse(IkeNotifyPayload cookie2Notify)
-            throws InvalidSyntaxException {
+    private static IkeNotifyPayload handleCookieAndGenerateCopy(
+            IkeNotifyPayload cookie2Notify, int minLen, int maxLen) throws InvalidSyntaxException {
         byte[] notifyData = cookie2Notify.notifyData;
-        if (notifyData.length < COOKIE2_DATA_LEN_MIN || notifyData.length > COOKIE2_DATA_LEN_MAX) {
+        if (notifyData.length < minLen || notifyData.length > maxLen) {
+            String cookieType =
+                    cookie2Notify.notifyType == NOTIFY_TYPE_COOKIE2 ? "COOKIE2" : "COOKIE";
             throw new InvalidSyntaxException(
-                    "Invalid COOKIE2 notification data with length " + notifyData.length);
+                    "Invalid "
+                            + cookieType
+                            + " notification data with length "
+                            + notifyData.length);
         }
 
-        return new IkeNotifyPayload(NOTIFY_TYPE_COOKIE2, notifyData);
+        return new IkeNotifyPayload(cookie2Notify.notifyType, notifyData);
+    }
+
+    /** Validate inbound Cookie in IKE_INIT response and build a Cookie notify payload in request */
+    public static IkeNotifyPayload handleCookieAndGenerateCopy(IkeNotifyPayload cookieNotify)
+            throws InvalidSyntaxException {
+        return handleCookieAndGenerateCopy(cookieNotify, COOKIE_DATA_LEN_MIN, COOKIE_DATA_LEN_MAX);
+    }
+
+    /** Validate inbound Cookie2 request and build a response Cookie2 notify payload */
+    public static IkeNotifyPayload handleCookie2AndGenerateCopy(IkeNotifyPayload cookie2Notify)
+            throws InvalidSyntaxException {
+        return handleCookieAndGenerateCopy(
+                cookie2Notify, COOKIE2_DATA_LEN_MIN, COOKIE2_DATA_LEN_MAX);
     }
 
     /**
diff --git a/src/java/com/android/internal/net/ipsec/ike/net/IkeDefaultNetworkCallback.java b/src/java/com/android/internal/net/ipsec/ike/net/IkeDefaultNetworkCallback.java
index f3c7fa0..fa4c6bf 100644
--- a/src/java/com/android/internal/net/ipsec/ike/net/IkeDefaultNetworkCallback.java
+++ b/src/java/com/android/internal/net/ipsec/ike/net/IkeDefaultNetworkCallback.java
@@ -21,7 +21,7 @@
 import java.net.InetAddress;
 
 /**
- * IkeDefaultNetworkCallback is a network callback used to track the platform's default network.
+ * IkeDefaultNetworkCallback is a network callback used to track the application default Network.
  *
  * <p>This NetworkCallback will notify IkeSessionStateMachine if:
  *
@@ -47,7 +47,7 @@
             return;
         }
 
-        logd("Platform default Network changed to " + network);
+        logd("Application default Network changed to " + network);
         mIkeNetworkUpdater.onUnderlyingNetworkUpdated(network);
     }
 }
diff --git a/tests/cts/src/android/eap/cts/EapSessionConfigTest.java b/tests/cts/src/android/eap/cts/EapSessionConfigTest.java
index b0a257a..784e4f3 100644
--- a/tests/cts/src/android/eap/cts/EapSessionConfigTest.java
+++ b/tests/cts/src/android/eap/cts/EapSessionConfigTest.java
@@ -28,18 +28,25 @@
 import android.net.eap.EapSessionConfig.EapAkaPrimeConfig;
 import android.net.eap.EapSessionConfig.EapMsChapV2Config;
 import android.net.eap.EapSessionConfig.EapSimConfig;
+import android.net.eap.EapSessionConfig.EapTtlsConfig;
 import android.net.eap.EapSessionConfig.EapUiccConfig;
 
 import androidx.test.runner.AndroidJUnit4;
 
+import com.android.internal.net.ipsec.ike.testutils.CertUtils;
+
+import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
+import java.security.cert.X509Certificate;
+
 @RunWith(AndroidJUnit4.class)
 public class EapSessionConfigTest {
     // These constants are IANA-defined values and are copies of hidden constants in
     // frameworks/opt/net/ike/src/java/com/android/internal/net/eap/message/EapData.java.
     private static final int EAP_TYPE_SIM = 18;
+    private static final int EAP_TYPE_TTLS = 21;
     private static final int EAP_TYPE_AKA = 23;
     private static final int EAP_TYPE_MSCHAP_V2 = 26;
     private static final int EAP_TYPE_AKA_PRIME = 50;
@@ -50,6 +57,19 @@
     private static final String EAP_MSCHAPV2_USERNAME = "username";
     private static final String EAP_MSCHAPV2_PASSWORD = "password";
 
+    private static final EapSessionConfig INNER_EAP_SESSION_CONFIG =
+            new EapSessionConfig.Builder()
+                    .setEapIdentity(EAP_IDENTITY)
+                    .setEapMsChapV2Config(EAP_MSCHAPV2_USERNAME, EAP_MSCHAPV2_PASSWORD)
+                    .build();
+
+    private X509Certificate mServerCaCert;
+
+    @Before
+    public void setUp() throws Exception {
+        mServerCaCert = CertUtils.createCertFromPemFile("server-a-self-signed-ca.pem");
+    }
+
     @Test
     public void testBuildWithAllEapMethods() {
         EapSessionConfig result =
@@ -63,6 +83,7 @@
                                 NETWORK_NAME,
                                 true /* allowMismatchedNetworkNames */)
                         .setEapMsChapV2Config(EAP_MSCHAPV2_USERNAME, EAP_MSCHAPV2_PASSWORD)
+                        .setEapTtlsConfig(mServerCaCert, INNER_EAP_SESSION_CONFIG)
                         .build();
 
         assertArrayEquals(EAP_IDENTITY, result.getEapIdentity());
@@ -89,6 +110,12 @@
         assertEquals(EAP_TYPE_MSCHAP_V2, eapMsChapV2Config.getMethodType());
         assertEquals(EAP_MSCHAPV2_USERNAME, eapMsChapV2Config.getUsername());
         assertEquals(EAP_MSCHAPV2_PASSWORD, eapMsChapV2Config.getPassword());
+
+        EapTtlsConfig eapTtlsConfig = result.getEapTtlsConfig();
+        assertNotNull(eapTtlsConfig);
+        assertEquals(EAP_TYPE_TTLS, eapTtlsConfig.getMethodType());
+        assertEquals(mServerCaCert, eapTtlsConfig.getServerCaCert());
+        assertEquals(INNER_EAP_SESSION_CONFIG, eapTtlsConfig.getInnerEapSessionConfig());
     }
 
     private void verifyEapUiccConfigCommon(EapUiccConfig config) {
diff --git a/tests/iketests/src/java/android/net/ipsec/ike/IkeSessionParamsTest.java b/tests/iketests/src/java/android/net/ipsec/ike/IkeSessionParamsTest.java
index f08bdfb..7ad524a 100644
--- a/tests/iketests/src/java/android/net/ipsec/ike/IkeSessionParamsTest.java
+++ b/tests/iketests/src/java/android/net/ipsec/ike/IkeSessionParamsTest.java
@@ -55,7 +55,7 @@
 import android.net.SocketKeepalive;
 import android.net.eap.EapSessionConfig;
 import android.net.ipsec.ike.ike3gpp.Ike3gppExtension;
-import android.net.ipsec.ike.ike3gpp.Ike3gppExtension.Ike3gppCallback;
+import android.net.ipsec.ike.ike3gpp.Ike3gppExtension.Ike3gppDataListener;
 import android.net.ipsec.ike.ike3gpp.Ike3gppParams;
 import android.os.PersistableBundle;
 import android.telephony.TelephonyManager;
@@ -284,7 +284,7 @@
     public void testEncodeIkeSessionParamsWith3gppExtension() throws Exception {
         Ike3gppExtension ike3gppExtension =
                 new Ike3gppExtension(
-                        new Ike3gppParams.Builder().build(), mock(Ike3gppCallback.class));
+                        new Ike3gppParams.Builder().build(), mock(Ike3gppDataListener.class));
 
         IkeSessionParams sessionParams =
                 buildWithPskCommon(REMOTE_IPV4_HOST_ADDRESS)
@@ -728,7 +728,7 @@
     public void testBuildWithIke3gppExtension() throws Exception {
         Ike3gppExtension ike3gppExtension =
                 new Ike3gppExtension(
-                        new Ike3gppParams.Builder().build(), mock(Ike3gppCallback.class));
+                        new Ike3gppParams.Builder().build(), mock(Ike3gppDataListener.class));
 
         IkeSessionParams sessionParams =
                 buildWithPskCommon(REMOTE_IPV4_HOST_ADDRESS)
diff --git a/tests/iketests/src/java/android/net/ipsec/ike/ike3gpp/Ike3gppBackoffTimerTest.java b/tests/iketests/src/java/android/net/ipsec/ike/ike3gpp/Ike3gppBackoffTimerTest.java
index ab576be..b0f966a 100644
--- a/tests/iketests/src/java/android/net/ipsec/ike/ike3gpp/Ike3gppBackoffTimerTest.java
+++ b/tests/iketests/src/java/android/net/ipsec/ike/ike3gpp/Ike3gppBackoffTimerTest.java
@@ -32,7 +32,7 @@
     public void testIke3gppBackoffTimer() {
         Ike3gppBackoffTimer backoffTimer = new Ike3gppBackoffTimer(BACKOFF_TIMER, BACKOFF_CAUSE);
 
-        assertEquals(Ike3gppInfo.INFO_TYPE_NOTIFY_BACKOFF_TIMER, backoffTimer.getInfoType());
+        assertEquals(Ike3gppData.DATA_TYPE_NOTIFY_BACKOFF_TIMER, backoffTimer.getDataType());
         assertEquals(BACKOFF_TIMER, backoffTimer.getBackoffTimer());
         assertEquals(BACKOFF_CAUSE, backoffTimer.getBackoffCause());
     }
diff --git a/tests/iketests/src/java/android/net/ipsec/ike/ike3gpp/Ike3gppExtensionTest.java b/tests/iketests/src/java/android/net/ipsec/ike/ike3gpp/Ike3gppExtensionTest.java
index 12658b6..62c41a3 100644
--- a/tests/iketests/src/java/android/net/ipsec/ike/ike3gpp/Ike3gppExtensionTest.java
+++ b/tests/iketests/src/java/android/net/ipsec/ike/ike3gpp/Ike3gppExtensionTest.java
@@ -20,7 +20,7 @@
 import static org.junit.Assert.assertNotEquals;
 import static org.mockito.Mockito.mock;
 
-import android.net.ipsec.ike.ike3gpp.Ike3gppExtension.Ike3gppCallback;
+import android.net.ipsec.ike.ike3gpp.Ike3gppExtension.Ike3gppDataListener;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -28,21 +28,21 @@
 public class Ike3gppExtensionTest {
     private static final byte PDU_SESSION_ID = (byte) 0x01;
 
-    private Ike3gppCallback mMockIke3gppCallback;
+    private Ike3gppDataListener mMockIke3gppDataListener;
     private Ike3gppParams mIke3gppParams;
 
     @Before
     public void setUp() {
-        mMockIke3gppCallback = mock(Ike3gppCallback.class);
+        mMockIke3gppDataListener = mock(Ike3gppDataListener.class);
         mIke3gppParams = new Ike3gppParams.Builder().build();
     }
 
     @Test
     public void testIke3gppExtensionConstructor() {
         Ike3gppExtension ike3gppExtension =
-                new Ike3gppExtension(mIke3gppParams, mMockIke3gppCallback);
+                new Ike3gppExtension(mIke3gppParams, mMockIke3gppDataListener);
 
-        assertEquals(mMockIke3gppCallback, ike3gppExtension.getIke3gppCallback());
+        assertEquals(mMockIke3gppDataListener, ike3gppExtension.getIke3gppDataListener());
         assertEquals(mIke3gppParams, ike3gppExtension.getIke3gppParams());
     }
 
@@ -53,7 +53,7 @@
 
     @Test(expected = NullPointerException.class)
     public void testIke3gppExtensionConstructorInvalidParams() {
-        Ike3gppExtension ike3gppExtension = new Ike3gppExtension(null, mMockIke3gppCallback);
+        Ike3gppExtension ike3gppExtension = new Ike3gppExtension(null, mMockIke3gppDataListener);
     }
 
     @Test
@@ -61,12 +61,12 @@
         Ike3gppExtension extensionA =
                 new Ike3gppExtension(
                         new Ike3gppParams.Builder().setPduSessionId(PDU_SESSION_ID).build(),
-                        mMockIke3gppCallback);
+                        mMockIke3gppDataListener);
 
         Ike3gppExtension extensionB =
                 new Ike3gppExtension(
                         new Ike3gppParams.Builder().setPduSessionId(PDU_SESSION_ID).build(),
-                        mMockIke3gppCallback);
+                        mMockIke3gppDataListener);
 
         assertEquals(extensionA, extensionB);
     }
@@ -74,12 +74,12 @@
     @Test
     public void testNotEquals() {
         Ike3gppExtension extensionA =
-                new Ike3gppExtension(new Ike3gppParams.Builder().build(), mMockIke3gppCallback);
+                new Ike3gppExtension(new Ike3gppParams.Builder().build(), mMockIke3gppDataListener);
 
         Ike3gppExtension extensionB =
                 new Ike3gppExtension(
                         new Ike3gppParams.Builder().setPduSessionId(PDU_SESSION_ID).build(),
-                        mMockIke3gppCallback);
+                        mMockIke3gppDataListener);
 
         assertNotEquals(extensionA, extensionB);
     }
diff --git a/tests/iketests/src/java/android/net/ipsec/ike/ike3gpp/Ike3gppN1ModeInformationTest.java b/tests/iketests/src/java/android/net/ipsec/ike/ike3gpp/Ike3gppN1ModeInformationTest.java
index 2f43b7b..f3ea3d2 100644
--- a/tests/iketests/src/java/android/net/ipsec/ike/ike3gpp/Ike3gppN1ModeInformationTest.java
+++ b/tests/iketests/src/java/android/net/ipsec/ike/ike3gpp/Ike3gppN1ModeInformationTest.java
@@ -31,7 +31,7 @@
         Ike3gppN1ModeInformation n1ModeInformation = new Ike3gppN1ModeInformation(SNSSAI);
 
         assertEquals(
-                Ike3gppInfo.INFO_TYPE_NOTIFY_N1_MODE_INFORMATION, n1ModeInformation.getInfoType());
+                Ike3gppData.DATA_TYPE_NOTIFY_N1_MODE_INFORMATION, n1ModeInformation.getDataType());
         assertArrayEquals(SNSSAI, n1ModeInformation.getSnssai());
     }
 
diff --git a/tests/iketests/src/java/com/android/internal/net/ipsec/ike/IkeSessionStateMachineTest.java b/tests/iketests/src/java/com/android/internal/net/ipsec/ike/IkeSessionStateMachineTest.java
index ccd731d..eb68ab0 100644
--- a/tests/iketests/src/java/com/android/internal/net/ipsec/ike/IkeSessionStateMachineTest.java
+++ b/tests/iketests/src/java/com/android/internal/net/ipsec/ike/IkeSessionStateMachineTest.java
@@ -51,6 +51,7 @@
 import static com.android.internal.net.ipsec.ike.message.IkeConfigPayload.CONFIG_ATTR_IP6_PCSCF;
 import static com.android.internal.net.ipsec.ike.message.IkeHeader.EXCHANGE_TYPE_CREATE_CHILD_SA;
 import static com.android.internal.net.ipsec.ike.message.IkeHeader.EXCHANGE_TYPE_INFORMATIONAL;
+import static com.android.internal.net.ipsec.ike.message.IkeNotifyPayload.NOTIFY_TYPE_COOKIE;
 import static com.android.internal.net.ipsec.ike.message.IkeNotifyPayload.NOTIFY_TYPE_COOKIE2;
 import static com.android.internal.net.ipsec.ike.message.IkeNotifyPayload.NOTIFY_TYPE_EAP_ONLY_AUTHENTICATION;
 import static com.android.internal.net.ipsec.ike.message.IkeNotifyPayload.NOTIFY_TYPE_IKEV2_FRAGMENTATION_SUPPORTED;
@@ -96,6 +97,8 @@
 import android.annotation.Nullable;
 import android.app.AlarmManager;
 import android.content.Context;
+import android.net.LinkAddress;
+import android.net.LinkProperties;
 import android.net.Network;
 import android.net.eap.EapSessionConfig;
 import android.net.ipsec.ike.ChildSaProposal;
@@ -116,16 +119,16 @@
 import android.net.ipsec.ike.exceptions.AuthenticationFailedException;
 import android.net.ipsec.ike.exceptions.IkeException;
 import android.net.ipsec.ike.exceptions.IkeInternalException;
-import android.net.ipsec.ike.exceptions.IkeNetworkDiedException;
+import android.net.ipsec.ike.exceptions.IkeNetworkLostException;
 import android.net.ipsec.ike.exceptions.IkeProtocolException;
 import android.net.ipsec.ike.exceptions.InvalidSyntaxException;
 import android.net.ipsec.ike.exceptions.NoValidProposalChosenException;
 import android.net.ipsec.ike.exceptions.UnrecognizedIkeProtocolException;
 import android.net.ipsec.ike.exceptions.UnsupportedCriticalPayloadException;
 import android.net.ipsec.ike.ike3gpp.Ike3gppBackoffTimer;
+import android.net.ipsec.ike.ike3gpp.Ike3gppData;
 import android.net.ipsec.ike.ike3gpp.Ike3gppExtension;
-import android.net.ipsec.ike.ike3gpp.Ike3gppExtension.Ike3gppCallback;
-import android.net.ipsec.ike.ike3gpp.Ike3gppInfo;
+import android.net.ipsec.ike.ike3gpp.Ike3gppExtension.Ike3gppDataListener;
 import android.net.ipsec.ike.ike3gpp.Ike3gppN1ModeInformation;
 import android.net.ipsec.ike.ike3gpp.Ike3gppParams;
 import android.os.Looper;
@@ -203,6 +206,7 @@
 
 import java.io.IOException;
 import java.net.Inet4Address;
+import java.net.Inet6Address;
 import java.net.InetAddress;
 import java.nio.ByteBuffer;
 import java.security.GeneralSecurityException;
@@ -335,13 +339,16 @@
 
     private static final int PAYLOAD_TYPE_UNSUPPORTED = 127;
 
+    private static final int COOKIE_DATA_LEN = 64;
     private static final int COOKIE2_DATA_LEN = 64;
 
+    private static final byte[] COOKIE_DATA = new byte[COOKIE_DATA_LEN];
     private static final byte[] COOKIE2_DATA = new byte[COOKIE2_DATA_LEN];
 
     private static final int NATT_KEEPALIVE_DELAY = 20;
 
     static {
+        new Random().nextBytes(COOKIE_DATA);
         new Random().nextBytes(COOKIE2_DATA);
     }
 
@@ -358,6 +365,8 @@
     private IkeUdp6Socket mSpyIkeUdp6Socket;
     private IkeSocket mSpyCurrentIkeSocket;
 
+    private LinkAddress mMockLinkAddressGlobalV6;
+
     private IkeNattKeepalive mMockIkeNattKeepalive;
 
     private TestLooper mLooper;
@@ -398,7 +407,7 @@
 
     private IkeLocalAddressGenerator mMockIkeLocalAddressGenerator;
 
-    private Ike3gppCallback mMockIke3gppCallback;
+    private Ike3gppDataListener mMockIke3gppDataListener;
     private Ike3gppExtension mIke3gppExtension;
 
     private X509Certificate mRootCertificate;
@@ -417,31 +426,39 @@
     private ArgumentCaptor<List<IkePayload>> mPayloadListCaptor =
             ArgumentCaptor.forClass(List.class);
 
-    private ReceivedIkePacket makeDummyReceivedIkeInitRespPacket(
-            long initiatorSpi,
-            long responderSpi,
-            @IkeHeader.ExchangeType int eType,
-            boolean isResp,
-            boolean fromIkeInit,
-            List<Integer> payloadTypeList,
-            List<String> payloadHexStringList)
+    private ReceivedIkePacket makeDummyReceivedIkeInitRespPacket(List<IkePayload> payloadList)
             throws Exception {
+        long dummyInitSpi = 1L;
+        long dummyRespSpi = 2L;
 
-        List<IkePayload> payloadList =
-                hexStrListToIkePayloadList(payloadTypeList, payloadHexStringList, isResp);
         // Build a remotely generated NAT_DETECTION_SOURCE_IP payload to mock a remote node's
         // network that is not behind NAT.
         IkePayload sourceNatPayload =
                 new IkeNotifyPayload(
                         NOTIFY_TYPE_NAT_DETECTION_SOURCE_IP,
                         IkeNotifyPayload.generateNatDetectionData(
-                                initiatorSpi,
-                                responderSpi,
+                                dummyInitSpi,
+                                dummyRespSpi,
                                 REMOTE_ADDRESS,
                                 IkeSocket.SERVER_PORT_UDP_ENCAPSULATED));
         payloadList.add(sourceNatPayload);
+
         return makeDummyUnencryptedReceivedIkePacket(
-                initiatorSpi, responderSpi, eType, isResp, fromIkeInit, payloadList);
+                dummyInitSpi,
+                dummyRespSpi,
+                IkeHeader.EXCHANGE_TYPE_IKE_SA_INIT,
+                true /*isResp*/,
+                false /*fromIkeInit*/,
+                payloadList);
+    }
+
+    private ReceivedIkePacket makeDummyReceivedIkeInitRespPacket(
+            List<Integer> payloadTypeList, List<String> payloadHexStringList) throws Exception {
+
+        List<IkePayload> payloadList =
+                hexStrListToIkePayloadList(
+                        payloadTypeList, payloadHexStringList, true /* isResp */);
+        return makeDummyReceivedIkeInitRespPacket(payloadList);
     }
 
     private ReceivedIkePacket makeDummyUnencryptedReceivedIkePacket(
@@ -754,6 +771,10 @@
                         eq(mMockDefaultNetwork), eq(false /* isIpv4 */), any(), anyInt()))
                 .thenReturn(LOCAL_ADDRESS_V6);
 
+        mMockLinkAddressGlobalV6 = mock(LinkAddress.class);
+        when(mMockLinkAddressGlobalV6.getAddress()).thenReturn(UPDATED_LOCAL_ADDRESS_V6);
+        when(mMockLinkAddressGlobalV6.isGlobalPreferred()).thenReturn(true);
+
         mMockEapAuthenticatorFactory = mock(IkeEapAuthenticatorFactory.class);
         mMockEapAuthenticator = mock(EapAuthenticator.class);
         doReturn(mMockEapAuthenticator)
@@ -818,7 +839,7 @@
         mExpectedCurrentSaLocalReqMsgId = 0;
         mExpectedCurrentSaRemoteReqMsgId = 0;
 
-        mMockIke3gppCallback = mock(Ike3gppCallback.class);
+        mMockIke3gppDataListener = mock(Ike3gppDataListener.class);
 
         mMockIkeNattKeepalive = mock(IkeNattKeepalive.class);
     }
@@ -938,7 +959,7 @@
         Ike3gppExtension ike3gppExtension =
                 new Ike3gppExtension(
                         new Ike3gppParams.Builder().setPduSessionId(pduSessionId).build(),
-                        mMockIke3gppCallback);
+                        mMockIke3gppDataListener);
         return buildIkeSessionParamsCommon()
                 .setAuthPsk(mPsk)
                 .setIke3gppExtension(ike3gppExtension)
@@ -997,16 +1018,7 @@
         payloadHexStringList.add(NONCE_RESP_PAYLOAD_HEX_STRING);
         payloadHexStringList.addAll(optionalPayloadHexStrings);
 
-        // In each test assign different IKE responder SPI in IKE INIT response to avoid remote SPI
-        // collision during response validation.
-        // STOPSHIP: b/131617794 allow #mockIkeSetup to be independent in each test after we can
-        // support IkeSession cleanup.
         return makeDummyReceivedIkeInitRespPacket(
-                1L /*initiator SPI*/,
-                2L /*responder SPI*/,
-                IkeHeader.EXCHANGE_TYPE_IKE_SA_INIT,
-                true /*isResp*/,
-                false /*fromIkeInit*/,
                 payloadTypeList,
                 payloadHexStringList);
     }
@@ -1416,7 +1428,7 @@
                         .build();
         mIkeSessionStateMachine = makeAndStartIkeSession(ikeParams);
 
-        verify(mMockDefaultNetwork).getByName(REMOTE_HOSTNAME);
+        verify(mMockDefaultNetwork).getAllByName(REMOTE_HOSTNAME);
     }
 
     @Test
@@ -1457,11 +1469,6 @@
         // Send back a INVALID_KE_PAYLOAD, and verify that the selected DH group changes
         ReceivedIkePacket resp =
                 makeDummyReceivedIkeInitRespPacket(
-                        1L /*initiator SPI*/,
-                        2L /*responder SPI*/,
-                        IkeHeader.EXCHANGE_TYPE_IKE_SA_INIT,
-                        true /*isResp*/,
-                        false /*fromIkeInit*/,
                         Arrays.asList(IkePayload.PAYLOAD_TYPE_NOTIFY),
                         Arrays.asList(INVALID_KE_PAYLOAD_HEX_STRING));
         mIkeSessionStateMachine.sendMessage(IkeSessionStateMachine.CMD_RECEIVE_IKE_PACKET, resp);
@@ -1472,6 +1479,45 @@
     }
 
     @Test
+    public void testCreateIkeLocalIkeInitReceivesCookie() throws Exception {
+        setupFirstIkeSa();
+
+        mIkeSessionStateMachine.sendMessage(IkeSessionStateMachine.CMD_LOCAL_REQUEST_CREATE_IKE);
+        mLooper.dispatchAll();
+
+        // Encode 2 times: one for mIkeInitRequestBytes and one for sending packets
+        verify(mMockIkeMessageHelper, times(2)).encode(mIkeMessageCaptor.capture());
+        IkeMessage originalReqMsg = mIkeMessageCaptor.getValue();
+        List<IkePayload> originalPayloadList = originalReqMsg.ikePayloadList;
+
+        // Reset to forget sending original IKE INIT request
+        resetMockIkeMessageHelper();
+
+        // Send back a Notify-Cookie
+        IkeNotifyPayload inCookieNotify = new IkeNotifyPayload(NOTIFY_TYPE_COOKIE, COOKIE_DATA);
+        List<IkePayload> payloads = new ArrayList<>();
+        payloads.add(inCookieNotify);
+        ReceivedIkePacket resp = makeDummyReceivedIkeInitRespPacket(payloads);
+        mIkeSessionStateMachine.sendMessage(IkeSessionStateMachine.CMD_RECEIVE_IKE_PACKET, resp);
+        mLooper.dispatchAll();
+
+        // Verify retry IKE INIT request
+        verify(mMockIkeMessageHelper, times(2)).encode(mIkeMessageCaptor.capture());
+        IkeMessage ikeInitReqMessage = mIkeMessageCaptor.getValue();
+        List<IkePayload> payloadList = ikeInitReqMessage.ikePayloadList;
+
+        IkeNotifyPayload outCookieNotify = (IkeNotifyPayload) payloadList.get(0);
+        assertEquals(NOTIFY_TYPE_COOKIE, outCookieNotify.notifyType);
+        assertArrayEquals(COOKIE_DATA, outCookieNotify.notifyData);
+
+        assertEquals(originalPayloadList, payloadList.subList(1, payloadList.size()));
+
+        assertTrue(
+                mIkeSessionStateMachine.getCurrentState()
+                        instanceof IkeSessionStateMachine.CreateIkeLocalIkeInit);
+    }
+
+    @Test
     public void testCreateIkeLocalIkeInitSwitchesToEncapPorts() throws Exception {
         setupFirstIkeSa();
         mIkeSessionStateMachine.sendMessage(IkeSessionStateMachine.CMD_LOCAL_REQUEST_CREATE_IKE);
@@ -2354,27 +2400,8 @@
                 hasChildPayloads,
                 hasConfigPayloadInResp,
                 false /* isMobikeEnabled */,
-                0 /* ike3gppCallbackInvocations */);
-    }
-
-    private IkeMessage verifySharedKeyAuthentication(
-            IkeAuthPskPayload spyAuthPayload,
-            IkeIdPayload respIdPayload,
-            List<IkePayload> authRelatedPayloads,
-            boolean hasChildPayloads,
-            boolean hasConfigPayloadInResp,
-            boolean isMobikeEnabled,
-            int ike3gppCallbackInvocations)
-            throws Exception {
-        return verifySharedKeyAuthentication(
-                spyAuthPayload,
-                respIdPayload,
-                authRelatedPayloads,
-                hasChildPayloads,
-                hasConfigPayloadInResp,
-                isMobikeEnabled,
                 true /* isIpv4 */,
-                ike3gppCallbackInvocations);
+                0 /* ike3gppDataListenerInvocations */);
     }
 
     private IkeMessage verifySharedKeyAuthentication(
@@ -2385,7 +2412,7 @@
             boolean hasConfigPayloadInResp,
             boolean isMobikeEnabled,
             boolean isIpv4,
-            int ike3gppCallbackInvocations)
+            int ike3gppDataListenerInvocations)
             throws Exception {
         IkeMessage ikeAuthReqMessage =
                 verifyAuthenticationCommonAndGetIkeMessage(
@@ -2395,7 +2422,7 @@
                         hasConfigPayloadInResp,
                         isMobikeEnabled,
                         isIpv4,
-                        ike3gppCallbackInvocations);
+                        ike3gppDataListenerInvocations);
 
         // Validate authentication is done. Cannot use matchers because IkeAuthPskPayload is final.
         verify(spyAuthPayload)
@@ -2421,7 +2448,7 @@
             boolean hasConfigPayloadInResp,
             boolean isMobikeEnabled,
             boolean isIpv4,
-            int ike3gppCallbackInvocations)
+            int ike3gppDataListenerInvocations)
             throws Exception {
         // Send IKE AUTH response to IKE state machine
         ReceivedIkePacket authResp = makeIkeAuthRespWithChildPayloads(authRelatedPayloads);
@@ -2451,9 +2478,9 @@
         verifyDecodeEncryptedMessage(mSpyCurrentIkeSaRecord, authResp);
 
         // Validate that user has been notified. Expect one invocation for
-        // IkeSessionCallback#onOpened and 'ike3gppCallbackInvocations' invocations for
-        // Ike3gppCallback#onIke3gppPayloadsReceived
-        verify(mSpyUserCbExecutor, times(1 + ike3gppCallbackInvocations))
+        // IkeSessionCallback#onOpened and 'ike3gppDataListenerInvocations' invocations for
+        // Ike3gppDataListener#onIke3gppDataReceived
+        verify(mSpyUserCbExecutor, times(1 + ike3gppDataListenerInvocations))
                 .execute(any(Runnable.class));
 
         // Verify IkeSessionConfiguration
@@ -2570,8 +2597,8 @@
             }
         }
 
-        // Only expect a N1_MODE_CAPABILITY payload if an Ike3gppExention and PDU Session ID are
-        // specified.
+        // Only expect a N1_MODE_CAPABILITY payload if an Ike3gppExtension and PDU Session ID
+        // are specified.
         Ike3gppExtension ike3gppExtension =
                 mIkeSessionStateMachine.mIkeSessionParams.getIke3gppExtension();
         if (ike3gppExtension == null || !ike3gppExtension.getIke3gppParams().hasPduSessionId()) {
@@ -3159,7 +3186,8 @@
                 false /*hasChildPayloads*/,
                 false /*hasConfigPayloadInResp*/,
                 isMobikeEnabled,
-                0 /* ike3gppCallbackInvocations */);
+                true /* isIpv4 */,
+                0 /* ike3gppDataListenerInvocations */);
         verifyRetransmissionStopped();
     }
 
@@ -4802,7 +4830,6 @@
 
         IkeSessionParams mockSessionParams = mock(IkeSessionParams.class);
         when(mockSessionParams.getServerHostname()).thenReturn(REMOTE_HOSTNAME);
-        when(mMockDefaultNetwork.getByName(REMOTE_HOSTNAME)).thenReturn(REMOTE_ADDRESS);
 
         RuntimeException cause = new RuntimeException();
         when(mockSessionParams.getSaProposalsInternal()).thenThrow(cause);
@@ -5172,13 +5199,13 @@
     @Test
     public void testIkeAuthWithN1Mode() throws Exception {
         verifyIkeAuthWith3gppEnabled(
-                makeN1ModeInformationPayload(), 1 /* ike3gppCallbackInvocations */);
+                makeN1ModeInformationPayload(), 1 /* ike3gppDataListenerInvocations */);
 
         verifyN1ModeReceived();
     }
 
     private void verifyIkeAuthWith3gppEnabled(
-            IkePayload ike3gppPayload, int ike3gppCallbackInvocations) throws Exception {
+            IkePayload ike3gppPayload, int ike3gppDataListenerInvocations) throws Exception {
         // Quit and restart IKE Session with N1 Mode Capability params
         mIkeSessionStateMachine.quitNow();
         reset(mMockChildSessionFactoryHelper);
@@ -5204,7 +5231,8 @@
                 true /*hasChildPayloads*/,
                 true /*hasConfigPayloadInResp*/,
                 false /* isMobikeEnabled */,
-                ike3gppCallbackInvocations);
+                true /* isIpv4 */,
+                ike3gppDataListenerInvocations);
         verifyRetransmissionStopped();
     }
 
@@ -5217,12 +5245,12 @@
     }
 
     private void verifyN1ModeReceived() {
-        ArgumentCaptor<List<Ike3gppInfo>> ike3gppInfoCaptor = ArgumentCaptor.forClass(List.class);
-        verify(mMockIke3gppCallback).onIke3gppPayloadsReceived(ike3gppInfoCaptor.capture());
+        ArgumentCaptor<List<Ike3gppData>> ike3gppDataCaptor = ArgumentCaptor.forClass(List.class);
+        verify(mMockIke3gppDataListener).onIke3gppDataReceived(ike3gppDataCaptor.capture());
 
         Ike3gppN1ModeInformation n1ModeInformation = null;
-        for (Ike3gppInfo payload : ike3gppInfoCaptor.getValue()) {
-            if (payload.getInfoType() == Ike3gppInfo.INFO_TYPE_NOTIFY_N1_MODE_INFORMATION) {
+        for (Ike3gppData payload : ike3gppDataCaptor.getValue()) {
+            if (payload.getDataType() == Ike3gppData.DATA_TYPE_NOTIFY_N1_MODE_INFORMATION) {
                 n1ModeInformation = (Ike3gppN1ModeInformation) payload;
             }
         }
@@ -5270,12 +5298,12 @@
     }
 
     private void verifyBackoffTimer(int expectedNotifyErrorCause) {
-        ArgumentCaptor<List<Ike3gppInfo>> ike3gppInfoCaptor = ArgumentCaptor.forClass(List.class);
-        verify(mMockIke3gppCallback).onIke3gppPayloadsReceived(ike3gppInfoCaptor.capture());
+        ArgumentCaptor<List<Ike3gppData>> ike3gppDataCaptor = ArgumentCaptor.forClass(List.class);
+        verify(mMockIke3gppDataListener).onIke3gppDataReceived(ike3gppDataCaptor.capture());
 
         Ike3gppBackoffTimer backoffTimer = null;
-        for (Ike3gppInfo payload : ike3gppInfoCaptor.getValue()) {
-            if (payload.getInfoType() == Ike3gppInfo.INFO_TYPE_NOTIFY_BACKOFF_TIMER) {
+        for (Ike3gppData payload : ike3gppDataCaptor.getValue()) {
+            if (payload.getDataType() == Ike3gppData.DATA_TYPE_NOTIFY_BACKOFF_TIMER) {
                 backoffTimer = (Ike3gppBackoffTimer) payload;
             }
         }
@@ -5289,10 +5317,10 @@
     public void testIkeAuthWithBackoffTimerWithoutError() throws Exception {
         verifyIkeAuthWith3gppEnabled(
                 new IkeNotifyPayload(NOTIFY_TYPE_BACKOFF_TIMER, BACKOFF_TIMER_DATA),
-                0 /* ike3gppCallbackInvocations */);
+                0 /* ike3gppDataListenerInvocations */);
 
         // BackoffTimer should be ignored
-        verify(mMockIke3gppCallback, never()).onIke3gppPayloadsReceived(any());
+        verify(mMockIke3gppDataListener, never()).onIke3gppDataReceived(any());
     }
 
     @Test(expected = IllegalArgumentException.class)
@@ -5425,8 +5453,15 @@
 
         // IKE client always supports NAT-T. So the peer decides if both sides support NAT-T.
         mIkeSessionStateMachine.mSupportNatTraversal = doesPeerSupportNatt;
-        mIkeSessionStateMachine.mLocalAddress = isIpv4 ? LOCAL_ADDRESS : LOCAL_ADDRESS_V6;
-        mIkeSessionStateMachine.mRemoteAddress = isIpv4 ? REMOTE_ADDRESS : REMOTE_ADDRESS_V6;
+        if (isIpv4) {
+            mIkeSessionStateMachine.mLocalAddress = LOCAL_ADDRESS;
+            mIkeSessionStateMachine.mRemoteAddress = REMOTE_ADDRESS;
+            mIkeSessionStateMachine.mRemoteAddressesV4.add(REMOTE_ADDRESS);
+        } else {
+            mIkeSessionStateMachine.mLocalAddress = LOCAL_ADDRESS_V6;
+            mIkeSessionStateMachine.mRemoteAddress = REMOTE_ADDRESS_V6;
+            mIkeSessionStateMachine.mRemoteAddressesV6.add(REMOTE_ADDRESS_V6);
+        }
 
         if (doesPeerSupportNatt && isIpv4) {
             // Assume NATs are detected on both sides
@@ -5467,7 +5502,7 @@
                         true /* hasConfigPayloadInResp */,
                         doesPeerSupportMobike,
                         isIpv4,
-                        0 /* ike3gppCallbackInvocations */);
+                        0 /* ike3gppDataListenerInvocations */);
         verifyRetransmissionStopped();
 
         boolean isMobikeSupportIndicated = false;
@@ -5576,7 +5611,7 @@
 
         ArgumentCaptor<IkeException> exceptionCaptor = ArgumentCaptor.forClass(IkeException.class);
         verify(mMockIkeSessionCallback).onError(exceptionCaptor.capture());
-        IkeNetworkDiedException cause = (IkeNetworkDiedException) exceptionCaptor.getValue();
+        IkeNetworkLostException cause = (IkeNetworkLostException) exceptionCaptor.getValue();
         assertEquals(mMockDefaultNetwork, cause.getNetwork());
     }
 
@@ -5598,9 +5633,23 @@
     private Network mockNewNetworkAndAddress(boolean isIpv4) throws Exception {
         Network newNetwork = mock(Network.class);
 
-        InetAddress expectedRemoteAddress = isIpv4 ? REMOTE_ADDRESS : REMOTE_ADDRESS_V6;
-        InetAddress injectedLocalAddress =
-                isIpv4 ? UPDATED_LOCAL_ADDRESS : UPDATED_LOCAL_ADDRESS_V6;
+        InetAddress expectedRemoteAddress;
+        InetAddress injectedLocalAddress;
+        if (isIpv4) {
+            expectedRemoteAddress = REMOTE_ADDRESS;
+            injectedLocalAddress = UPDATED_LOCAL_ADDRESS;
+
+            mIkeSessionStateMachine.mRemoteAddressesV4.add((Inet4Address) expectedRemoteAddress);
+        } else {
+            expectedRemoteAddress = REMOTE_ADDRESS_V6;
+            injectedLocalAddress = UPDATED_LOCAL_ADDRESS_V6;
+            mIkeSessionStateMachine.mRemoteAddressesV6.add((Inet6Address) expectedRemoteAddress);
+
+            LinkProperties linkProperties = new LinkProperties();
+            linkProperties.addLinkAddress(mMockLinkAddressGlobalV6);
+            when(mMockConnectManager.getLinkProperties(eq(newNetwork))).thenReturn(linkProperties);
+        }
+
         when(mMockIkeLocalAddressGenerator.generateLocalAddress(
                         eq(newNetwork), eq(isIpv4), eq(expectedRemoteAddress), anyInt()))
                 .thenReturn(injectedLocalAddress);
diff --git a/tests/iketests/src/java/com/android/internal/net/ipsec/ike/IkeSessionTestBase.java b/tests/iketests/src/java/com/android/internal/net/ipsec/ike/IkeSessionTestBase.java
index 5d14a04..7528c4c 100644
--- a/tests/iketests/src/java/com/android/internal/net/ipsec/ike/IkeSessionTestBase.java
+++ b/tests/iketests/src/java/com/android/internal/net/ipsec/ike/IkeSessionTestBase.java
@@ -50,6 +50,7 @@
 
 import java.net.Inet4Address;
 import java.net.Inet6Address;
+import java.net.InetAddress;
 import java.util.concurrent.Executor;
 
 public abstract class IkeSessionTestBase {
@@ -109,10 +110,12 @@
                 .newWakeLock(anyInt(), argThat(tag -> tag.contains(LOCAL_REQUEST_WAKE_LOCK_TAG)));
 
         mMockDefaultNetwork = mock(Network.class);
-        doReturn(REMOTE_ADDRESS).when(mMockDefaultNetwork).getByName(REMOTE_HOSTNAME);
-        doReturn(REMOTE_ADDRESS)
+        doReturn(new InetAddress[] {REMOTE_ADDRESS})
                 .when(mMockDefaultNetwork)
-                .getByName(REMOTE_ADDRESS.getHostAddress());
+                .getAllByName(REMOTE_HOSTNAME);
+        doReturn(new InetAddress[] {REMOTE_ADDRESS})
+                .when(mMockDefaultNetwork)
+                .getAllByName(REMOTE_ADDRESS.getHostAddress());
 
         mMockSocketKeepalive = mock(SocketKeepalive.class);
 
diff --git a/tests/iketests/src/java/com/android/internal/net/ipsec/ike/ike3gpp/Ike3gppExtensionExchangeTest.java b/tests/iketests/src/java/com/android/internal/net/ipsec/ike/ike3gpp/Ike3gppExtensionExchangeTest.java
index 08200b2..0d46bc1 100644
--- a/tests/iketests/src/java/com/android/internal/net/ipsec/ike/ike3gpp/Ike3gppExtensionExchangeTest.java
+++ b/tests/iketests/src/java/com/android/internal/net/ipsec/ike/ike3gpp/Ike3gppExtensionExchangeTest.java
@@ -28,7 +28,7 @@
 import static org.mockito.Mockito.mock;
 
 import android.net.ipsec.ike.ike3gpp.Ike3gppExtension;
-import android.net.ipsec.ike.ike3gpp.Ike3gppExtension.Ike3gppCallback;
+import android.net.ipsec.ike.ike3gpp.Ike3gppExtension.Ike3gppDataListener;
 import android.net.ipsec.ike.ike3gpp.Ike3gppParams;
 
 import com.android.internal.net.ipsec.ike.message.IkeNotifyPayload;
@@ -59,19 +59,19 @@
 
     private static final Executor INLINE_EXECUTOR = Runnable::run;
 
-    private Ike3gppCallback mMockIke3gppCallback;
+    private Ike3gppDataListener mMockIke3gppDataListener;
 
     private Ike3gppParams mIke3gppParams;
     private Ike3gppExtensionExchange mIke3gppExtensionExchange;
 
     @Before
     public void setUp() {
-        mMockIke3gppCallback = mock(Ike3gppCallback.class);
+        mMockIke3gppDataListener = mock(Ike3gppDataListener.class);
 
         mIke3gppParams = new Ike3gppParams.Builder().setPduSessionId(PDU_SESSION_ID).build();
         mIke3gppExtensionExchange =
                 new Ike3gppExtensionExchange(
-                        new Ike3gppExtension(mIke3gppParams, mMockIke3gppCallback),
+                        new Ike3gppExtension(mIke3gppParams, mMockIke3gppDataListener),
                         INLINE_EXECUTOR);
     }
 
diff --git a/tests/iketests/src/java/com/android/internal/net/ipsec/ike/message/IkeNotifyPayloadTest.java b/tests/iketests/src/java/com/android/internal/net/ipsec/ike/message/IkeNotifyPayloadTest.java
index fd30297..26ed661 100644
--- a/tests/iketests/src/java/com/android/internal/net/ipsec/ike/message/IkeNotifyPayloadTest.java
+++ b/tests/iketests/src/java/com/android/internal/net/ipsec/ike/message/IkeNotifyPayloadTest.java
@@ -31,6 +31,7 @@
 import static android.net.ipsec.ike.exceptions.IkeProtocolException.ERROR_TYPE_TS_UNACCEPTABLE;
 import static android.net.ipsec.ike.exceptions.IkeProtocolException.ERROR_TYPE_UNSUPPORTED_CRITICAL_PAYLOAD;
 
+import static com.android.internal.net.ipsec.ike.message.IkeNotifyPayload.NOTIFY_TYPE_COOKIE;
 import static com.android.internal.net.ipsec.ike.message.IkeNotifyPayload.NOTIFY_TYPE_COOKIE2;
 
 import static org.junit.Assert.assertArrayEquals;
@@ -74,14 +75,13 @@
     private static final String PACKET_INFO_HEX_STRING =
             "4500009cafcd4000403208adc0a80064c0a800012ad4c0a200000001";
 
-    private static final int COOKIE_INVALID_DATA_LEN_SMALL = 7;
+    private static final int COOKIE_INVALID_DATA_LEN_SMALL = 0;
     private static final int COOKIE_INVALID_DATA_LEN_LARGE = 65;
     private static final int COOKIE_DATA_LEN = 64;
-    private static final byte[] COOKIE2_DATA_BYTES = new byte[COOKIE_DATA_LEN];
 
-    static {
-        new Random().nextBytes(COOKIE2_DATA_BYTES);
-    }
+    private static final int COOKIE2_INVALID_DATA_LEN_SMALL = 7;
+    private static final int COOKIE2_INVALID_DATA_LEN_LARGE = 65;
+    private static final int COOKIE2_DATA_LEN = 64;
 
     private static final String NOTIFY_REKEY_PAYLOAD_BODY_HEX_STRING = "030440092ad4c0a2";
     private static final int CHILD_SPI = 0x2ad4c0a2;
@@ -133,30 +133,50 @@
         assertArrayEquals(expectedBytes, netDetectionData);
     }
 
+    private void verifyHandleCookieAndGenerateCopy(boolean isCookie2, int dataLen)
+            throws Exception {
+        final byte[] cookieData = new byte[dataLen];
+        new Random().nextBytes(cookieData);
+        int cookieType = isCookie2 ? NOTIFY_TYPE_COOKIE2 : NOTIFY_TYPE_COOKIE;
+        IkeNotifyPayload inboundCookieNotify = new IkeNotifyPayload(cookieType, cookieData);
+
+        IkeNotifyPayload outboundCookieNotify =
+                isCookie2
+                        ? IkeNotifyPayload.handleCookie2AndGenerateCopy(inboundCookieNotify)
+                        : IkeNotifyPayload.handleCookieAndGenerateCopy(inboundCookieNotify);
+
+        assertArrayEquals(cookieData, outboundCookieNotify.notifyData);
+        assertEquals(cookieType, outboundCookieNotify.notifyType);
+    }
+
     @Test
-    public void testHandleAndReplyCookie2Request() throws Exception {
-        IkeNotifyPayload cookie2Req = new IkeNotifyPayload(NOTIFY_TYPE_COOKIE2, COOKIE2_DATA_BYTES);
-        IkeNotifyPayload cookie2Resp =
-                IkeNotifyPayload.handleCookie2AndGenerateResponse(cookie2Req);
-        assertArrayEquals(COOKIE2_DATA_BYTES, cookie2Resp.notifyData);
-        assertEquals(NOTIFY_TYPE_COOKIE2, cookie2Resp.notifyType);
-    }
-
-    private void checkHandleCookie2Request(int dataLen) throws Exception {
-        final byte[] invalidCookie2Data = new byte[dataLen];
-        new Random().nextBytes(invalidCookie2Data);
-        IkeNotifyPayload cookie2Req = new IkeNotifyPayload(NOTIFY_TYPE_COOKIE2, invalidCookie2Data);
-        IkeNotifyPayload.handleCookie2AndGenerateResponse(cookie2Req);
+    public void testHandleCookieAndGenerateCopy() throws Exception {
+        verifyHandleCookieAndGenerateCopy(false /* isCookie2 */, COOKIE_DATA_LEN);
     }
 
     @Test(expected = InvalidSyntaxException.class)
-    public void testHandleCookie2RequestWithTooSmallLengthOfData() throws Exception {
-        checkHandleCookie2Request(COOKIE_INVALID_DATA_LEN_SMALL);
+    public void testHandleCookieWithTooSmallLengthOfData() throws Exception {
+        verifyHandleCookieAndGenerateCopy(false /* isCookie2 */, COOKIE_INVALID_DATA_LEN_SMALL);
     }
 
     @Test(expected = InvalidSyntaxException.class)
-    public void testHandleCookie2RequestWithTooLargeLengthOfData() throws Exception {
-        checkHandleCookie2Request(COOKIE_INVALID_DATA_LEN_LARGE);
+    public void testHandleCookieWithTooLargeLengthOfData() throws Exception {
+        verifyHandleCookieAndGenerateCopy(false /* isCookie2 */, COOKIE_INVALID_DATA_LEN_SMALL);
+    }
+
+    @Test
+    public void testHandleCookie2AndGenerateCopy() throws Exception {
+        verifyHandleCookieAndGenerateCopy(true /* isCookie2 */, COOKIE2_DATA_LEN);
+    }
+
+    @Test(expected = InvalidSyntaxException.class)
+    public void testHandleCookie2WithTooSmallLengthOfData() throws Exception {
+        verifyHandleCookieAndGenerateCopy(true /* isCookie2 */, COOKIE2_INVALID_DATA_LEN_SMALL);
+    }
+
+    @Test(expected = InvalidSyntaxException.class)
+    public void testHandleCookie2WithTooLargeLengthOfData() throws Exception {
+        verifyHandleCookieAndGenerateCopy(true /* isCookie2 */, COOKIE2_INVALID_DATA_LEN_SMALL);
     }
 
     @Test