Merge "Bump Mainline Module versionNumber"
diff --git a/api/current.txt b/api/current.txt
index 9d138d1..81a5d61 100644
--- a/api/current.txt
+++ b/api/current.txt
@@ -1,4 +1,51 @@
 // Signature format: 2.0
+package android.net.eap {
+
+  public final class EapSessionConfig {
+    method @Nullable public android.net.eap.EapSessionConfig.EapAkaConfig getEapAkaConfig();
+    method @Nullable public android.net.eap.EapSessionConfig.EapAkaPrimeConfig getEapAkaPrimeConfig();
+    method @NonNull public byte[] getEapIdentity();
+    method @Nullable public android.net.eap.EapSessionConfig.EapMsChapV2Config getEapMsChapV2Config();
+    method @Nullable public android.net.eap.EapSessionConfig.EapSimConfig getEapSimConfig();
+  }
+
+  public static final class EapSessionConfig.Builder {
+    ctor public EapSessionConfig.Builder();
+    method @NonNull public android.net.eap.EapSessionConfig build();
+    method @NonNull public android.net.eap.EapSessionConfig.Builder setEapAkaConfig(int, int);
+    method @NonNull public android.net.eap.EapSessionConfig.Builder setEapAkaPrimeConfig(int, int, @NonNull String, boolean);
+    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);
+  }
+
+  public static class EapSessionConfig.EapAkaConfig extends android.net.eap.EapSessionConfig.EapUiccConfig {
+  }
+
+  public static class EapSessionConfig.EapAkaPrimeConfig extends android.net.eap.EapSessionConfig.EapAkaConfig {
+    method public boolean allowsMismatchedNetworkNames();
+    method @NonNull public String getNetworkName();
+  }
+
+  public abstract static class EapSessionConfig.EapMethodConfig {
+    method public int getMethodType();
+  }
+
+  public static class EapSessionConfig.EapMsChapV2Config extends android.net.eap.EapSessionConfig.EapMethodConfig {
+    method @NonNull public String getPassword();
+    method @NonNull public String getUsername();
+  }
+
+  public static class EapSessionConfig.EapSimConfig extends android.net.eap.EapSessionConfig.EapUiccConfig {
+  }
+
+  public abstract static class EapSessionConfig.EapUiccConfig extends android.net.eap.EapSessionConfig.EapMethodConfig {
+    method public int getAppType();
+    method public int getSubId();
+  }
+
+}
+
 package android.net.ipsec.ike {
 
   public final class ChildSaProposal extends android.net.ipsec.ike.SaProposal {
@@ -83,6 +130,15 @@
     method @NonNull public android.net.ipsec.ike.IkeSaProposal build();
   }
 
+  public final class IkeSession implements java.lang.AutoCloseable {
+    ctor public IkeSession(@NonNull android.content.Context, @NonNull android.net.ipsec.ike.IkeSessionParams, @NonNull android.net.ipsec.ike.ChildSessionParams, @NonNull java.util.concurrent.Executor, @NonNull android.net.ipsec.ike.IkeSessionCallback, @NonNull android.net.ipsec.ike.ChildSessionCallback);
+    method public void close();
+    method public void closeChildSession(@NonNull android.net.ipsec.ike.ChildSessionCallback);
+    method public void finalize();
+    method public void kill();
+    method public void openChildSession(@NonNull android.net.ipsec.ike.ChildSessionParams, @NonNull android.net.ipsec.ike.ChildSessionCallback);
+  }
+
   public interface IkeSessionCallback {
     method public void onClosed();
     method public void onClosedExceptionally(@NonNull android.net.ipsec.ike.exceptions.IkeException);
@@ -105,6 +161,77 @@
     method @NonNull public java.net.InetAddress getRemoteAddress();
   }
 
+  public final class IkeSessionParams {
+    method @NonNull public java.util.List<android.net.ipsec.ike.IkeSessionParams.IkeConfigRequest> getConfigurationRequests();
+    method @Nullable public android.net.Network getConfiguredNetwork();
+    method @IntRange(from=0x14, to=0x708) public int getDpdDelaySeconds();
+    method @IntRange(from=0x12c, to=0x15180) public int getHardLifetimeSeconds();
+    method @NonNull public android.net.ipsec.ike.IkeSessionParams.IkeAuthConfig getLocalAuthConfig();
+    method @NonNull public android.net.ipsec.ike.IkeIdentification getLocalIdentification();
+    method @NonNull public android.net.ipsec.ike.IkeSessionParams.IkeAuthConfig getRemoteAuthConfig();
+    method @NonNull public android.net.ipsec.ike.IkeIdentification getRemoteIdentification();
+    method @NonNull public int[] getRetransmissionTimeoutsMillis();
+    method @NonNull public java.util.List<android.net.ipsec.ike.IkeSaProposal> getSaProposals();
+    method @NonNull public String getServerHostname();
+    method @IntRange(from=0x78, to=0x15180) public int getSoftLifetimeSeconds();
+    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
+  }
+
+  public static final class IkeSessionParams.Builder {
+    ctor public IkeSessionParams.Builder(@NonNull android.content.Context);
+    method @NonNull public android.net.ipsec.ike.IkeSessionParams.Builder addIkeOption(int);
+    method @NonNull public android.net.ipsec.ike.IkeSessionParams.Builder addPcscfServerRequest(@NonNull java.net.InetAddress);
+    method @NonNull public android.net.ipsec.ike.IkeSessionParams.Builder addPcscfServerRequest(int);
+    method @NonNull public android.net.ipsec.ike.IkeSessionParams.Builder addSaProposal(@NonNull android.net.ipsec.ike.IkeSaProposal);
+    method @NonNull public android.net.ipsec.ike.IkeSessionParams build();
+    method @NonNull public android.net.ipsec.ike.IkeSessionParams.Builder removeIkeOption(int);
+    method @NonNull public android.net.ipsec.ike.IkeSessionParams.Builder setAuthDigitalSignature(@Nullable java.security.cert.X509Certificate, @NonNull java.security.cert.X509Certificate, @NonNull java.security.PrivateKey);
+    method @NonNull public android.net.ipsec.ike.IkeSessionParams.Builder setAuthDigitalSignature(@Nullable java.security.cert.X509Certificate, @NonNull java.security.cert.X509Certificate, @NonNull java.util.List<java.security.cert.X509Certificate>, @NonNull java.security.PrivateKey);
+    method @NonNull public android.net.ipsec.ike.IkeSessionParams.Builder setAuthEap(@Nullable java.security.cert.X509Certificate, @NonNull android.net.eap.EapSessionConfig);
+    method @NonNull public android.net.ipsec.ike.IkeSessionParams.Builder setAuthPsk(@NonNull byte[]);
+    method @NonNull public android.net.ipsec.ike.IkeSessionParams.Builder setConfiguredNetwork(@NonNull android.net.Network);
+    method @NonNull public android.net.ipsec.ike.IkeSessionParams.Builder setDpdDelaySeconds(@IntRange(from=0x14, to=0x708) int);
+    method @NonNull public android.net.ipsec.ike.IkeSessionParams.Builder setLifetimeSeconds(@IntRange(from=0x12c, to=0x15180) int, @IntRange(from=0x78, to=0x15180) int);
+    method @NonNull public android.net.ipsec.ike.IkeSessionParams.Builder setLocalIdentification(@NonNull android.net.ipsec.ike.IkeIdentification);
+    method @NonNull public android.net.ipsec.ike.IkeSessionParams.Builder setRemoteIdentification(@NonNull android.net.ipsec.ike.IkeIdentification);
+    method @NonNull public android.net.ipsec.ike.IkeSessionParams.Builder setRetransmissionTimeoutsMillis(@NonNull int[]);
+    method @NonNull public android.net.ipsec.ike.IkeSessionParams.Builder setServerHostname(@NonNull String);
+  }
+
+  public static interface IkeSessionParams.ConfigRequestIpv4PcscfServer extends android.net.ipsec.ike.IkeSessionParams.IkeConfigRequest {
+    method @Nullable public java.net.Inet4Address getAddress();
+  }
+
+  public static interface IkeSessionParams.ConfigRequestIpv6PcscfServer extends android.net.ipsec.ike.IkeSessionParams.IkeConfigRequest {
+    method @Nullable public java.net.Inet6Address getAddress();
+  }
+
+  public abstract static class IkeSessionParams.IkeAuthConfig {
+  }
+
+  public static class IkeSessionParams.IkeAuthDigitalSignLocalConfig extends android.net.ipsec.ike.IkeSessionParams.IkeAuthConfig {
+    method @NonNull public java.security.cert.X509Certificate getClientEndCertificate();
+    method @NonNull public java.util.List<java.security.cert.X509Certificate> getIntermediateCertificates();
+    method @NonNull public java.security.PrivateKey getPrivateKey();
+  }
+
+  public static class IkeSessionParams.IkeAuthDigitalSignRemoteConfig extends android.net.ipsec.ike.IkeSessionParams.IkeAuthConfig {
+    method @Nullable public java.security.cert.X509Certificate getRemoteCaCert();
+  }
+
+  public static class IkeSessionParams.IkeAuthEapConfig extends android.net.ipsec.ike.IkeSessionParams.IkeAuthConfig {
+    method @NonNull public android.net.eap.EapSessionConfig getEapConfig();
+  }
+
+  public static class IkeSessionParams.IkeAuthPskConfig extends android.net.ipsec.ike.IkeSessionParams.IkeAuthConfig {
+    method @NonNull public byte[] getPsk();
+  }
+
+  public static interface IkeSessionParams.IkeConfigRequest {
+  }
+
   public final class IkeTrafficSelector {
     ctor public IkeTrafficSelector(int, int, @NonNull java.net.InetAddress, @NonNull java.net.InetAddress);
     field public final int endPort;
diff --git a/api/system-current.txt b/api/system-current.txt
index eb826d3..342858d 100644
--- a/api/system-current.txt
+++ b/api/system-current.txt
@@ -2,135 +2,25 @@
 package android.net.eap {
 
   public final class EapSessionConfig {
-    method @Nullable public android.net.eap.EapSessionConfig.EapAkaConfig getEapAkaConfig();
-    method @Nullable public android.net.eap.EapSessionConfig.EapAkaPrimeConfig getEapAkaPrimeConfig();
-    method @NonNull public byte[] getEapIdentity();
-    method @Nullable public android.net.eap.EapSessionConfig.EapMsChapV2Config getEapMsChapV2onfig();
-    method @Nullable public android.net.eap.EapSessionConfig.EapSimConfig getEapSimConfig();
-  }
-
-  public static final class EapSessionConfig.Builder {
-    ctor public EapSessionConfig.Builder();
-    method @NonNull public android.net.eap.EapSessionConfig build();
-    method @NonNull public android.net.eap.EapSessionConfig.Builder setEapAkaConfig(int, int);
-    method @NonNull public android.net.eap.EapSessionConfig.Builder setEapAkaPrimeConfig(int, int, @NonNull String, boolean);
-    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);
-  }
-
-  public static class EapSessionConfig.EapAkaConfig extends android.net.eap.EapSessionConfig.EapUiccConfig {
-  }
-
-  public static class EapSessionConfig.EapAkaPrimeConfig extends android.net.eap.EapSessionConfig.EapAkaConfig {
-    method public boolean allowsMismatchedNetworkNames();
-    method @NonNull public String getNetworkName();
-  }
-
-  public abstract static class EapSessionConfig.EapMethodConfig {
-    method public int getMethodType();
-  }
-
-  public static class EapSessionConfig.EapMsChapV2Config extends android.net.eap.EapSessionConfig.EapMethodConfig {
-    method @NonNull public String getPassword();
-    method @NonNull public String getUsername();
-  }
-
-  public static class EapSessionConfig.EapSimConfig extends android.net.eap.EapSessionConfig.EapUiccConfig {
-  }
-
-  public abstract static class EapSessionConfig.EapUiccConfig extends android.net.eap.EapSessionConfig.EapMethodConfig {
-    method public int getAppType();
-    method public int getSubId();
+    method @Deprecated @Nullable public android.net.eap.EapSessionConfig.EapMsChapV2Config getEapMsChapV2onfig();
   }
 
 }
 
 package android.net.ipsec.ike {
 
-  public final class IkeSession implements java.lang.AutoCloseable {
-    ctor public IkeSession(@NonNull android.content.Context, @NonNull android.net.ipsec.ike.IkeSessionParams, @NonNull android.net.ipsec.ike.ChildSessionParams, @NonNull java.util.concurrent.Executor, @NonNull android.net.ipsec.ike.IkeSessionCallback, @NonNull android.net.ipsec.ike.ChildSessionCallback);
-    method public void close();
-    method public void closeChildSession(@NonNull android.net.ipsec.ike.ChildSessionCallback);
-    method public void kill();
-    method public void openChildSession(@NonNull android.net.ipsec.ike.ChildSessionParams, @NonNull android.net.ipsec.ike.ChildSessionCallback);
-  }
-
   public interface IkeSessionCallback {
     method public void onError(@NonNull android.net.ipsec.ike.exceptions.IkeProtocolException);
   }
 
   public final class IkeSessionParams {
-    method @NonNull public java.util.List<android.net.ipsec.ike.IkeSessionParams.IkeConfigRequest> getConfigurationRequests();
-    method @IntRange(from=0x14, to=0x708) public int getDpdDelaySeconds();
-    method @IntRange(from=0x12c, to=0x15180) public int getHardLifetimeSeconds();
     method @Nullable public android.net.ipsec.ike.ike3gpp.Ike3gppExtension getIke3gppExtension();
-    method @NonNull public android.net.ipsec.ike.IkeSessionParams.IkeAuthConfig getLocalAuthConfig();
-    method @NonNull public android.net.ipsec.ike.IkeIdentification getLocalIdentification();
-    method @NonNull public android.net.Network getNetwork();
-    method @NonNull public android.net.ipsec.ike.IkeSessionParams.IkeAuthConfig getRemoteAuthConfig();
-    method @NonNull public android.net.ipsec.ike.IkeIdentification getRemoteIdentification();
-    method public int[] getRetransmissionTimeoutsMillis();
-    method @NonNull public java.util.List<android.net.ipsec.ike.IkeSaProposal> getSaProposals();
-    method @NonNull public String getServerHostname();
-    method @IntRange(from=0x78, to=0x15180) public int getSoftLifetimeSeconds();
-    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
+    method @Deprecated @NonNull public android.net.Network getNetwork();
   }
 
   public static final class IkeSessionParams.Builder {
-    ctor public IkeSessionParams.Builder(@NonNull android.content.Context);
-    method @NonNull public android.net.ipsec.ike.IkeSessionParams.Builder addIkeOption(int);
-    method @NonNull public android.net.ipsec.ike.IkeSessionParams.Builder addPcscfServerRequest(@NonNull java.net.InetAddress);
-    method @NonNull public android.net.ipsec.ike.IkeSessionParams.Builder addPcscfServerRequest(int);
-    method @NonNull public android.net.ipsec.ike.IkeSessionParams.Builder addSaProposal(@NonNull android.net.ipsec.ike.IkeSaProposal);
-    method @NonNull public android.net.ipsec.ike.IkeSessionParams build();
-    method @NonNull public android.net.ipsec.ike.IkeSessionParams.Builder removeIkeOption(int);
-    method @NonNull public android.net.ipsec.ike.IkeSessionParams.Builder setAuthDigitalSignature(@Nullable java.security.cert.X509Certificate, @NonNull java.security.cert.X509Certificate, @NonNull java.security.PrivateKey);
-    method @NonNull public android.net.ipsec.ike.IkeSessionParams.Builder setAuthDigitalSignature(@Nullable java.security.cert.X509Certificate, @NonNull java.security.cert.X509Certificate, @NonNull java.util.List<java.security.cert.X509Certificate>, @NonNull java.security.PrivateKey);
-    method @NonNull public android.net.ipsec.ike.IkeSessionParams.Builder setAuthEap(@Nullable java.security.cert.X509Certificate, @NonNull android.net.eap.EapSessionConfig);
-    method @NonNull public android.net.ipsec.ike.IkeSessionParams.Builder setAuthPsk(@NonNull byte[]);
-    method @NonNull public android.net.ipsec.ike.IkeSessionParams.Builder setDpdDelaySeconds(@IntRange(from=0x14, to=0x708) int);
     method @NonNull public android.net.ipsec.ike.IkeSessionParams.Builder setIke3gppExtension(@NonNull android.net.ipsec.ike.ike3gpp.Ike3gppExtension);
-    method @NonNull public android.net.ipsec.ike.IkeSessionParams.Builder setLifetimeSeconds(@IntRange(from=0x12c, to=0x15180) int, @IntRange(from=0x78, to=0x15180) int);
-    method @NonNull public android.net.ipsec.ike.IkeSessionParams.Builder setLocalIdentification(@NonNull android.net.ipsec.ike.IkeIdentification);
-    method @NonNull public android.net.ipsec.ike.IkeSessionParams.Builder setNetwork(@NonNull android.net.Network);
-    method @NonNull public android.net.ipsec.ike.IkeSessionParams.Builder setRemoteIdentification(@NonNull android.net.ipsec.ike.IkeIdentification);
-    method @NonNull public android.net.ipsec.ike.IkeSessionParams.Builder setRetransmissionTimeoutsMillis(@NonNull int[]);
-    method @NonNull public android.net.ipsec.ike.IkeSessionParams.Builder setServerHostname(@NonNull String);
-  }
-
-  public static interface IkeSessionParams.ConfigRequestIpv4PcscfServer extends android.net.ipsec.ike.IkeSessionParams.IkeConfigRequest {
-    method @Nullable public java.net.Inet4Address getAddress();
-  }
-
-  public static interface IkeSessionParams.ConfigRequestIpv6PcscfServer extends android.net.ipsec.ike.IkeSessionParams.IkeConfigRequest {
-    method @Nullable public java.net.Inet6Address getAddress();
-  }
-
-  public abstract static class IkeSessionParams.IkeAuthConfig {
-  }
-
-  public static class IkeSessionParams.IkeAuthDigitalSignLocalConfig extends android.net.ipsec.ike.IkeSessionParams.IkeAuthConfig {
-    method @NonNull public java.security.cert.X509Certificate getClientEndCertificate();
-    method @NonNull public java.util.List<java.security.cert.X509Certificate> getIntermediateCertificates();
-    method @NonNull public java.security.PrivateKey getPrivateKey();
-  }
-
-  public static class IkeSessionParams.IkeAuthDigitalSignRemoteConfig extends android.net.ipsec.ike.IkeSessionParams.IkeAuthConfig {
-    method @Nullable public java.security.cert.X509Certificate getRemoteCaCert();
-  }
-
-  public static class IkeSessionParams.IkeAuthEapConfig extends android.net.ipsec.ike.IkeSessionParams.IkeAuthConfig {
-    method @NonNull public android.net.eap.EapSessionConfig getEapConfig();
-  }
-
-  public static class IkeSessionParams.IkeAuthPskConfig extends android.net.ipsec.ike.IkeSessionParams.IkeAuthConfig {
-    method @NonNull public byte[] getPsk();
-  }
-
-  public static interface IkeSessionParams.IkeConfigRequest {
+    method @Deprecated @NonNull public android.net.ipsec.ike.IkeSessionParams.Builder setNetwork(@NonNull android.net.Network);
   }
 
 }
diff --git a/src/java/android/net/eap/EapSessionConfig.java b/src/java/android/net/eap/EapSessionConfig.java
index ef30e22..538f2cc 100644
--- a/src/java/android/net/eap/EapSessionConfig.java
+++ b/src/java/android/net/eap/EapSessionConfig.java
@@ -47,10 +47,7 @@
  *
  * <p>The EAP authentication server decides which EAP method is used, so clients are encouraged to
  * provide configs for several EAP methods.
- *
- * @hide
  */
-@SystemApi
 public final class EapSessionConfig {
     private static final String EAP_ID_KEY = "eapIdentity";
     private static final String EAP_METHOD_CONFIGS_KEY = "eapConfigs";
@@ -175,11 +172,25 @@
      * @return the configuration for EAP MSCHAPV2, or null if it was not set
      */
     @Nullable
-    public EapMsChapV2Config getEapMsChapV2onfig() {
+    public EapMsChapV2Config getEapMsChapV2Config() {
         return (EapMsChapV2Config) mEapConfigs.get(EAP_TYPE_MSCHAP_V2);
     }
 
     /**
+     * Retrieves configuration for EAP MSCHAPV2
+     *
+     * @return the configuration for EAP MSCHAPV2, or null if it was not set
+     * @hide
+     * @deprecated Callers should use {@link #getEapMsChapV2Config}
+     */
+    @Deprecated
+    @SystemApi
+    @Nullable
+    public EapMsChapV2Config getEapMsChapV2onfig() {
+        return getEapMsChapV2Config();
+    }
+
+    /**
      * Retrieves configuration for EAP-TTLS
      *
      * @return the configuration for EAP-TTLS, or null if it was not set
diff --git a/src/java/android/net/ipsec/ike/ChildSessionCallback.java b/src/java/android/net/ipsec/ike/ChildSessionCallback.java
index aff56ae..9b67a00 100644
--- a/src/java/android/net/ipsec/ike/ChildSessionCallback.java
+++ b/src/java/android/net/ipsec/ike/ChildSessionCallback.java
@@ -110,8 +110,8 @@
      *     {@link IpSecManager#DIRECTION_OUT}
      * @hide
      */
-    void onIpSecTransformsMigrated(
-            @NonNull IpSecTransform inIpSecTransform, @NonNull IpSecTransform outIpSecTransform);
+    default void onIpSecTransformsMigrated(
+            @NonNull IpSecTransform inIpSecTransform, @NonNull IpSecTransform outIpSecTransform) {}
 
     /**
      * Called when an {@link IpSecTransform} is deleted by this Child Session.
diff --git a/src/java/android/net/ipsec/ike/IkeSession.java b/src/java/android/net/ipsec/ike/IkeSession.java
index 169cbff..2c4e404 100644
--- a/src/java/android/net/ipsec/ike/IkeSession.java
+++ b/src/java/android/net/ipsec/ike/IkeSession.java
@@ -17,7 +17,6 @@
 
 import android.annotation.NonNull;
 import android.annotation.SuppressLint;
-import android.annotation.SystemApi;
 import android.content.Context;
 import android.content.pm.PackageManager;
 import android.net.IpSecManager;
@@ -48,9 +47,7 @@
  *
  * @see <a href="https://tools.ietf.org/html/rfc7296">RFC 7296, Internet Key Exchange Protocol
  *     Version 2 (IKEv2)</a>
- * @hide
  */
-@SystemApi
 public final class IkeSession implements AutoCloseable {
     private final CloseGuard mCloseGuard = new CloseGuard();
     private final Context mContext;
diff --git a/src/java/android/net/ipsec/ike/IkeSessionCallback.java b/src/java/android/net/ipsec/ike/IkeSessionCallback.java
index 2dce4ce..ef51ff0 100644
--- a/src/java/android/net/ipsec/ike/IkeSessionCallback.java
+++ b/src/java/android/net/ipsec/ike/IkeSessionCallback.java
@@ -75,15 +75,22 @@
     void onError(@NonNull IkeProtocolException exception);
 
     /**
-     * Called if a non-protocol, recoverable error is encountered in an established {@link
-     * IkeSession}.
+     * Called if a recoverable error is encountered in an established {@link IkeSession}.
      *
-     * <p>This method can be triggered by non-protocol errors, such as the underlying Network for
-     * this IkeSession dying.
+     * <p>This method may be triggered by protocol errors such as an INVALID_IKE_SPI, or by
+     * non-protocol errors such as the underlying {@link android.net.Network} dying.
      *
+     * @param exception the detailed error information.
      * @hide
      */
-    void onError(@NonNull IkeException exception);
+    default void onError(@NonNull IkeException exception) {
+        if (exception instanceof IkeProtocolException) {
+            onError((IkeProtocolException) exception);
+            return;
+        }
+
+        // do nothing for non-protocol errors by default
+    }
 
     /**
      * Called if the IkeSessionConnectionInfo for an established {@link IkeSession} changes.
@@ -103,5 +110,6 @@
      * @param connectionInfo the updated IkeSessionConnectionInfo for the Session.
      * @hide
      */
-    void onIkeSessionConnectionInfoChanged(@NonNull IkeSessionConnectionInfo connectionInfo);
+    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 8b0f058..724aad5 100644
--- a/src/java/android/net/ipsec/ike/IkeSessionParams.java
+++ b/src/java/android/net/ipsec/ike/IkeSessionParams.java
@@ -66,10 +66,7 @@
  *
  * <p>Note that all negotiated configurations will be reused during rekey including SA Proposal and
  * lifetime.
- *
- * @hide
  */
-@SystemApi
 public final class IkeSessionParams {
     /** @hide */
     @Retention(RetentionPolicy.SOURCE)
@@ -382,16 +379,7 @@
         return mServerHostname;
     }
 
-    // TODO: b/151984042 Keep #getNetwork @SystemApi and make #getConfiguredNetwork public
-
-    /**
-     * Retrieves the configured {@link Network}, or null if not configured
-     *
-     * <p>This is the initially-configured Network (with MOBIKE, the caller may later specify a new
-     * Network and that update will not persist to here).
-     *
-     * @hide
-     */
+    /** Retrieves the configured {@link Network}, or null if was not set */
     @Nullable
     public Network getConfiguredNetwork() {
         return mCallerConfiguredNetwork;
@@ -405,14 +393,17 @@
      * informational because if MOBIKE is enabled, IKE Session may switch to a different default
      * Network.
      *
-     * <p>This is method will be deprecated. Callers should use {@link #getConfiguredNetwork}
+     * @hide
+     * @deprecated Callers should use {@link #getConfiguredNetwork}
      */
+    @Deprecated
+    @SystemApi
     @NonNull
     public Network getNetwork() {
         return mNetwork;
     }
 
-    /** Retrieves all ChildSaProposals configured */
+    /** Retrieves all IkeSaProposals configured */
     @NonNull
     public List<IkeSaProposal> getSaProposals() {
         return Arrays.asList(mSaProposals);
@@ -464,6 +455,8 @@
     }
 
     /** Retrieves the Dead Peer Detection(DPD) delay in seconds */
+    // Use "second" because smaller unit does not make sense to a DPD delay.
+    @SuppressLint("MethodNameUnits")
     @IntRange(from = IKE_DPD_DELAY_SEC_MIN, to = IKE_DPD_DELAY_SEC_MAX)
     public int getDpdDelaySeconds() {
         return mDpdDelaySec;
@@ -474,11 +467,17 @@
      *
      * <p>@see {@link Builder#setRetransmissionTimeoutsMillis(int[])}
      */
+    @NonNull
     public int[] getRetransmissionTimeoutsMillis() {
         return mRetransTimeoutMsList;
     }
 
-    /** Retrieves the configured Ike3gppExtension, or null if it was not set. */
+    /**
+     * Retrieves the configured Ike3gppExtension, or null if it was not set.
+     *
+     * @hide
+     */
+    @SystemApi
     @Nullable
     public Ike3gppExtension getIke3gppExtension() {
         return mIke3gppExtension;
@@ -1127,28 +1126,9 @@
          *
          * @param network the {@link Network} that IKE Session will use.
          * @return Builder this, to facilitate chaining.
-         * @hide
          */
         @NonNull
         public Builder setConfiguredNetwork(@NonNull Network network) {
-            return setNetwork(network);
-        }
-
-        // TODO: b/151984042 Keep #setNetwork @SystemApi and make #setConfiguredNetwork public
-
-        /**
-         * Sets the {@link Network} for the {@link IkeSessionParams} being built.
-         *
-         * <p>If no {@link Network} is provided, the default Network (as per {@link
-         * ConnectivityManager#getActiveNetwork()}) will be used.
-         *
-         * <p>This is method will be deprecated. Callers should use {@link #setConfiguredNetwork}
-         *
-         * @param network the {@link Network} that IKE Session will use.
-         * @return Builder this, to facilitate chaining.
-         */
-        @NonNull
-        public Builder setNetwork(@NonNull Network network) {
             if (network == null) {
                 throw new NullPointerException("Required argument not provided");
             }
@@ -1159,6 +1139,24 @@
         }
 
         /**
+         * Sets the {@link Network} for the {@link IkeSessionParams} being built.
+         *
+         * <p>If no {@link Network} is provided, the default Network (as per {@link
+         * ConnectivityManager#getActiveNetwork()}) will be used.
+         *
+         * @param network the {@link Network} that IKE Session will use.
+         * @return Builder this, to facilitate chaining.
+         * @hide
+         * @deprecated Callers should use {@link #setConfiguredNetwork}
+         */
+        @Deprecated
+        @SystemApi
+        @NonNull
+        public Builder setNetwork(@NonNull Network network) {
+            return setConfiguredNetwork(network);
+        }
+
+        /**
          * Sets local IKE identification for the {@link IkeSessionParams} being built.
          *
          * <p>It is not allowed to use KEY ID together with digital-signature-based authentication
@@ -1240,6 +1238,9 @@
          * @param sharedKey the shared key.
          * @return Builder this, to facilitate chaining.
          */
+        // #getLocalAuthConfig and #getRemoveAuthConfig are defined to retrieve
+        // authentication configurations
+        @SuppressLint("MissingGetterMatchingBuilder")
         @NonNull
         public Builder setAuthPsk(@NonNull byte[] sharedKey) {
             if (sharedKey == null) {
@@ -1283,6 +1284,9 @@
          */
         // TODO(b/151667921): Consider also supporting configuring EAP method that is not accepted
         // by EAP-Only when {@link IKE_OPTION_EAP_ONLY_AUTH} is set
+        // MissingGetterMatchingBuilder: #getLocalAuthConfig and #getRemoveAuthConfig are defined to
+        // retrieve authentication configurations
+        @SuppressLint("MissingGetterMatchingBuilder")
         @NonNull
         public Builder setAuthEap(
                 @Nullable X509Certificate serverCaCert, @NonNull EapSessionConfig eapConfig) {
@@ -1315,6 +1319,9 @@
          *     PrivateKey} MUST be an instance of {@link RSAKey}.
          * @return Builder this, to facilitate chaining.
          */
+        // #getLocalAuthConfig and #getRemoveAuthConfig are defined to retrieve
+        // authentication configurations
+        @SuppressLint("MissingGetterMatchingBuilder")
         @NonNull
         public Builder setAuthDigitalSignature(
                 @Nullable X509Certificate serverCaCert,
@@ -1349,6 +1356,9 @@
          *     PrivateKey} MUST be an instance of {@link RSAKey}.
          * @return Builder this, to facilitate chaining.
          */
+        // #getLocalAuthConfig and #getRemoveAuthConfig are defined to retrieve
+        // authentication configurations
+        @SuppressLint("MissingGetterMatchingBuilder")
         @NonNull
         public Builder setAuthDigitalSignature(
                 @Nullable X509Certificate serverCaCert,
@@ -1391,6 +1401,8 @@
          * @param address the requested P_CSCF address.
          * @return Builder this, to facilitate chaining.
          */
+        // #getConfigurationRequests is defined to retrieve PCSCF server requests
+        @SuppressLint("MissingGetterMatchingBuilder")
         @NonNull
         public Builder addPcscfServerRequest(@NonNull InetAddress address) {
             if (address == null) {
@@ -1413,6 +1425,8 @@
          *     OsConstants.AF_INET6} are allowed.
          * @return Builder this, to facilitate chaining.
          */
+        // #getConfigurationRequests is defined to retrieve PCSCF server requests
+        @SuppressLint("MissingGetterMatchingBuilder")
         @NonNull
         public Builder addPcscfServerRequest(int addressFamily) {
             if (addressFamily == AF_INET) {
@@ -1437,6 +1451,9 @@
          *     least 60 seconds (1 minute) shorter than the hard lifetime.
          * @return Builder this, to facilitate chaining.
          */
+        // #getHardLifetimeSeconds and #getSoftLifetimeSeconds are defined for callers to retrieve
+        // the lifetimes
+        @SuppressLint("MissingGetterMatchingBuilder")
         @NonNull
         public Builder setLifetimeSeconds(
                 @IntRange(from = IKE_HARD_LIFETIME_SEC_MINIMUM, to = IKE_HARD_LIFETIME_SEC_MAXIMUM)
@@ -1519,7 +1536,9 @@
          *     access networks
          * @param ike3gppExtension the Ike3gppExtension to use for this IKE Session.
          * @return Builder this, to facilitate chaining.
+         * @hide
          */
+        @SystemApi
         @NonNull
         public Builder setIke3gppExtension(@NonNull Ike3gppExtension ike3gppExtension) {
             Objects.requireNonNull(ike3gppExtension, "ike3gppExtension must not be null");
@@ -1534,6 +1553,9 @@
          * @param ikeOption the option to be enabled.
          * @return Builder this, to facilitate chaining.
          */
+        // Use #hasIkeOption instead of @getIkeOptions because #hasIkeOption allows callers to check
+        // the presence of one IKE option more easily
+        @SuppressLint("MissingGetterMatchingBuilder")
         @NonNull
         public Builder addIkeOption(@IkeOption int ikeOption) {
             validateIkeOptionOrThrow(ikeOption);
@@ -1547,6 +1569,9 @@
          * @param ikeOption the option to be disabled.
          * @return Builder this, to facilitate chaining.
          */
+        // Use #removeIkeOption instead of #clearIkeOption because "clear" sounds indicating
+        // clearing all enabled IKE options
+        @SuppressLint("BuilderSetStyle")
         @NonNull
         public Builder removeIkeOption(@IkeOption int ikeOption) {
             validateIkeOptionOrThrow(ikeOption);
diff --git a/src/java/android/net/ipsec/ike/exceptions/IkeNetworkDiedException.java b/src/java/android/net/ipsec/ike/exceptions/IkeNetworkDiedException.java
index fc52fba..5fdae70 100644
--- a/src/java/android/net/ipsec/ike/exceptions/IkeNetworkDiedException.java
+++ b/src/java/android/net/ipsec/ike/exceptions/IkeNetworkDiedException.java
@@ -31,14 +31,22 @@
  * <ul>
  *   <li>set a new underlying Network for the corresponding IkeSession (MOBIKE must be enabled and
  *       the IKE Session must have started with a caller-configured Network), or
+ *   <li>wait for a new underlying Network to become available (MOBIKE must be enabled and the IKE
+ *       Session must be tracking the System default Network), or
+ *       <ul>
+ *         <li>Note: if the maximum retransmission time is encountered while waiting, the IKE
+ *             Session will close. If this occurs, the caller will be notified via {@link
+ *             IkeSessionCallback#onClosedExceptionally(IkeException)}.
+ *       </ul>
  *   <li>close the corresponding IkeSession.
  * </ul>
  *
  * @hide
  */
-public final class IkeNetworkDiedException extends IkeException {
+public final class IkeNetworkDiedException extends IkeNonProtocolException {
     private final Network mNetwork;
 
+    /** Constructs an IkeNetworkDiedException to indicate the specified Network died. */
     public IkeNetworkDiedException(Network network) {
         super();
         Objects.requireNonNull(network, "network is null");
@@ -46,13 +54,6 @@
         mNetwork = network;
     }
 
-    public IkeNetworkDiedException(Network network, String message) {
-        super(message);
-        Objects.requireNonNull(network, "network is null");
-
-        mNetwork = network;
-    }
-
     /** Returns the IkeSession's underlying Network that died. */
     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
new file mode 100644
index 0000000..d8a3162
--- /dev/null
+++ b/src/java/android/net/ipsec/ike/exceptions/IkeNonProtocolException.java
@@ -0,0 +1,44 @@
+/*
+ * 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.exceptions;
+
+/**
+ * IkeNonProtocolException encapsulates all implementation-specific non-protocol IKE errors.
+ *
+ * @hide
+ */
+public abstract class IkeNonProtocolException extends IkeException {
+    /** @hide */
+    protected IkeNonProtocolException() {
+        super();
+    }
+
+    /** @hide */
+    protected IkeNonProtocolException(String message) {
+        super(message);
+    }
+
+    /** @hide */
+    protected IkeNonProtocolException(Throwable cause) {
+        super(cause);
+    }
+
+    /** @hide */
+    protected IkeNonProtocolException(String message, Throwable cause) {
+        super(message, cause);
+    }
+}
diff --git a/src/java/com/android/internal/net/eap/message/simaka/EapAkaPrimeTypeData.java b/src/java/com/android/internal/net/eap/message/simaka/EapAkaPrimeTypeData.java
index 23dcd02..2c15996 100644
--- a/src/java/com/android/internal/net/eap/message/simaka/EapAkaPrimeTypeData.java
+++ b/src/java/com/android/internal/net/eap/message/simaka/EapAkaPrimeTypeData.java
@@ -34,11 +34,17 @@
             new EapAkaPrimeTypeDataDecoder();
 
     @VisibleForTesting
-    EapAkaPrimeTypeData(
-            int eapSubType, LinkedHashMap<Integer, EapSimAkaAttribute> attributeMap) {
+    EapAkaPrimeTypeData(int eapSubType, LinkedHashMap<Integer, EapSimAkaAttribute> attributeMap) {
         super(eapSubType, attributeMap);
     }
 
+    private EapAkaPrimeTypeData(
+            int eapSubType,
+            LinkedHashMap<Integer, EapSimAkaAttribute> attributeMap,
+            byte[] reservedBytes) {
+        super(eapSubType, attributeMap, reservedBytes);
+    }
+
     /**
      * Creates and returns an EapAkaPrimeTypeData instance with the given subtype and attributes.
      *
@@ -88,8 +94,10 @@
 
         @Override
         protected EapAkaPrimeTypeData getInstance(
-                int eapSubtype, LinkedHashMap<Integer, EapSimAkaAttribute> attributeMap) {
-            return new EapAkaPrimeTypeData(eapSubtype, attributeMap);
+                int eapSubtype,
+                LinkedHashMap<Integer, EapSimAkaAttribute> attributeMap,
+                byte[] reservedBytes) {
+            return new EapAkaPrimeTypeData(eapSubtype, attributeMap, reservedBytes);
         }
     }
 }
diff --git a/src/java/com/android/internal/net/eap/message/simaka/EapAkaTypeData.java b/src/java/com/android/internal/net/eap/message/simaka/EapAkaTypeData.java
index e3ce817..0fb2aa4 100644
--- a/src/java/com/android/internal/net/eap/message/simaka/EapAkaTypeData.java
+++ b/src/java/com/android/internal/net/eap/message/simaka/EapAkaTypeData.java
@@ -73,6 +73,13 @@
         super(eapSubType, attributeMap);
     }
 
+    protected EapAkaTypeData(
+            int eapSubType,
+            LinkedHashMap<Integer, EapSimAkaAttribute> attributeMap,
+            byte[] reservedBytes) {
+        super(eapSubType, attributeMap, reservedBytes);
+    }
+
     /**
      * Creates and returns an EapAkaTypeData instance with the given subtype and attributes.
      *
@@ -130,8 +137,9 @@
         @Override
         protected EapAkaTypeData getInstance(
                 int eapSubtype,
-                LinkedHashMap<Integer, EapSimAkaAttribute> attributeMap) {
-            return new EapAkaTypeData(eapSubtype, attributeMap);
+                LinkedHashMap<Integer, EapSimAkaAttribute> attributeMap,
+                byte[] reservedBytes) {
+            return new EapAkaTypeData(eapSubtype, attributeMap, reservedBytes);
         }
     }
 }
diff --git a/src/java/com/android/internal/net/eap/message/simaka/EapSimAkaAttribute.java b/src/java/com/android/internal/net/eap/message/simaka/EapSimAkaAttribute.java
index 73d6752..cd41287 100644
--- a/src/java/com/android/internal/net/eap/message/simaka/EapSimAkaAttribute.java
+++ b/src/java/com/android/internal/net/eap/message/simaka/EapSimAkaAttribute.java
@@ -23,6 +23,7 @@
 import com.android.internal.net.eap.exceptions.simaka.EapSimAkaInvalidAttributeException;
 import com.android.internal.net.eap.exceptions.simaka.EapSimInvalidAtRandException;
 
+import java.nio.BufferUnderflowException;
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -132,6 +133,64 @@
      */
     public abstract void encode(ByteBuffer byteBuffer);
 
+    /**
+     * EapSimAkaReservedBytesAttribute represents any EAP-SIM/AKA attribute that is of the format:
+     *
+     * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
+     * |  Attribute Type (1B)  |  Length (1B)  |  Reserved (2B)  |
+     * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
+     * |  Value...
+     * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
+     *
+     * <p>Note: This Attribute type ignores (but preserves) the Reserved bytes. This is needed for
+     * calculating MACs in EAP-SIM/AKA.
+     */
+    protected abstract static class EapSimAkaReservedBytesAttribute extends EapSimAkaAttribute {
+        protected static final int RESERVED_BYTES_LEN = 2;
+
+        @VisibleForTesting public final byte[] reservedBytes = new byte[RESERVED_BYTES_LEN];
+
+        protected EapSimAkaReservedBytesAttribute(
+                int attributeType, int lengthInBytes, ByteBuffer buffer)
+                throws EapSimAkaInvalidAttributeException {
+            super(attributeType, lengthInBytes);
+
+            try {
+                buffer.get(reservedBytes);
+            } catch (BufferUnderflowException e) {
+                throw new EapSimAkaInvalidAttributeException("Invalid attribute length", e);
+            }
+        }
+
+        protected EapSimAkaReservedBytesAttribute(int attributeType, int lengthInBytes)
+                throws EapSimAkaInvalidAttributeException {
+            super(attributeType, lengthInBytes);
+        }
+
+        protected EapSimAkaReservedBytesAttribute(
+                int attributeType, int lengthInBytes, byte[] reservedBytes)
+                throws EapSimAkaInvalidAttributeException {
+            this(attributeType, lengthInBytes);
+
+            if (reservedBytes.length != RESERVED_BYTES_LEN) {
+                throw new EapSimAkaInvalidAttributeException("Invalid attribute length");
+            }
+            System.arraycopy(
+                    reservedBytes,
+                    0 /* srcPos */,
+                    this.reservedBytes,
+                    0 /* destPos */,
+                    RESERVED_BYTES_LEN);
+        }
+
+        @Override
+        public void encode(ByteBuffer buffer) {
+            encodeAttributeHeader(buffer);
+
+            buffer.put(reservedBytes);
+        }
+    }
+
     protected void encodeAttributeHeader(ByteBuffer byteBuffer) {
         byteBuffer.put((byte) attributeType);
         byteBuffer.put((byte) (lengthInBytes / LENGTH_SCALING));
@@ -286,9 +345,8 @@
     /**
      * AtNonceMt represents the AT_NONCE_MT attribute defined in RFC 4186#10.4
      */
-    public static class AtNonceMt extends EapSimAkaAttribute {
+    public static class AtNonceMt extends EapSimAkaReservedBytesAttribute {
         private static final int LENGTH = 5 * LENGTH_SCALING;
-        private static final int RESERVED_BYTES = 2;
 
         public static final int NONCE_MT_LENGTH = 16;
 
@@ -296,13 +354,11 @@
 
         public AtNonceMt(int lengthInBytes, ByteBuffer byteBuffer)
                 throws EapSimAkaInvalidAttributeException {
-            super(EAP_AT_NONCE_MT, LENGTH);
+            super(EAP_AT_NONCE_MT, LENGTH, byteBuffer);
             if (lengthInBytes != LENGTH) {
                 throw new EapSimAkaInvalidAttributeException("Invalid Length specified");
             }
 
-            // next two bytes are reserved (RFC 4186 Section 10.4)
-            byteBuffer.get(new byte[RESERVED_BYTES]);
             byteBuffer.get(nonceMt);
         }
 
@@ -318,38 +374,28 @@
 
         @Override
         public void encode(ByteBuffer byteBuffer) {
-            encodeAttributeHeader(byteBuffer);
-            byteBuffer.put(new byte[RESERVED_BYTES]);
+            super.encode(byteBuffer);
+
             byteBuffer.put(nonceMt);
         }
     }
 
-    private abstract static class AtIdReq extends EapSimAkaAttribute {
+    private abstract static class AtIdReq extends EapSimAkaReservedBytesAttribute {
         private static final int ATTR_LENGTH = LENGTH_SCALING;
-        private static final int RESERVED_BYTES = 2;
 
         protected AtIdReq(int lengthInBytes, int attributeType, ByteBuffer byteBuffer)
                 throws EapSimAkaInvalidAttributeException {
-            super(attributeType, ATTR_LENGTH);
+            super(attributeType, ATTR_LENGTH, byteBuffer);
 
             if (lengthInBytes != ATTR_LENGTH) {
                 throw new EapSimAkaInvalidAttributeException("Invalid Length specified");
             }
-
-            // next two bytes are reserved (RFC 4186 Section 10.5-10.7)
-            byteBuffer.get(new byte[RESERVED_BYTES]);
         }
 
         @VisibleForTesting
         protected AtIdReq(int attributeType) throws EapSimAkaInvalidAttributeException {
             super(attributeType, ATTR_LENGTH);
         }
-
-        @Override
-        public void encode(ByteBuffer byteBuffer) {
-            encodeAttributeHeader(byteBuffer);
-            byteBuffer.put(new byte[RESERVED_BYTES]);
-        }
     }
 
     /**
@@ -454,9 +500,8 @@
     /**
      * AtRandSim represents the AT_RAND attribute for EAP-SIM defined in RFC 4186#10.9
      */
-    public static class AtRandSim extends EapSimAkaAttribute {
+    public static class AtRandSim extends EapSimAkaReservedBytesAttribute {
         private static final int RAND_LENGTH = 16;
-        private static final int RESERVED_BYTES = 2;
         private static final int MIN_RANDS = 2;
         private static final int MAX_RANDS = 3;
 
@@ -464,10 +509,7 @@
 
         public AtRandSim(int lengthInBytes, ByteBuffer byteBuffer)
                 throws EapSimAkaInvalidAttributeException {
-            super(EAP_AT_RAND, lengthInBytes);
-
-            // next two bytes are reserved (RFC 4186 Section 10.9)
-            byteBuffer.get(new byte[RESERVED_BYTES]);
+            super(EAP_AT_RAND, lengthInBytes, byteBuffer);
 
             int numRands = (lengthInBytes - MIN_ATTR_LENGTH) / RAND_LENGTH;
             if (!isValidNumRands(numRands)) {
@@ -510,8 +552,7 @@
 
         @Override
         public void encode(ByteBuffer byteBuffer) {
-            encodeAttributeHeader(byteBuffer);
-            byteBuffer.put(new byte[RESERVED_BYTES]);
+            super.encode(byteBuffer);
 
             for (byte[] rand : rands) {
                 byteBuffer.put(rand);
@@ -522,24 +563,20 @@
     /**
      * AtRandAka represents the AT_RAND attribute for EAP-AKA defined in RFC 4187#10.6
      */
-    public static class AtRandAka extends EapSimAkaAttribute {
+    public static class AtRandAka extends EapSimAkaReservedBytesAttribute {
         private static final int ATTR_LENGTH = 5 * LENGTH_SCALING;
         private static final int RAND_LENGTH = 16;
-        private static final int RESERVED_BYTES = 2;
 
         public final byte[] rand = new byte[RAND_LENGTH];
 
         public AtRandAka(int lengthInBytes, ByteBuffer byteBuffer)
                 throws EapSimAkaInvalidAttributeException {
-            super(EAP_AT_RAND, lengthInBytes);
+            super(EAP_AT_RAND, lengthInBytes, byteBuffer);
 
             if (lengthInBytes != ATTR_LENGTH) {
                 throw new EapSimAkaInvalidAttributeException("Length must be 20B");
             }
 
-            // next two bytes are reserved (RFC 4187#10.6)
-            byteBuffer.get(new byte[RESERVED_BYTES]);
-
             byteBuffer.get(rand);
         }
 
@@ -557,8 +594,8 @@
 
         @Override
         public void encode(ByteBuffer byteBuffer) {
-            encodeAttributeHeader(byteBuffer);
-            byteBuffer.put(new byte[RESERVED_BYTES]);
+            super.encode(byteBuffer);
+
             byteBuffer.put(rand);
         }
     }
@@ -598,9 +635,8 @@
     /**
      * AtMac represents the AT_MAC attribute defined in RFC 4186#10.14 and RFC 4187#10.15
      */
-    public static class AtMac extends EapSimAkaAttribute {
+    public static class AtMac extends EapSimAkaReservedBytesAttribute {
         private static final int ATTR_LENGTH = 5 * LENGTH_SCALING;
-        private static final int RESERVED_BYTES = 2;
 
         public static final int MAC_LENGTH = 4 * LENGTH_SCALING;
 
@@ -608,41 +644,52 @@
 
         public AtMac(int lengthInBytes, ByteBuffer byteBuffer)
                 throws EapSimAkaInvalidAttributeException {
-            super(EAP_AT_MAC, lengthInBytes);
+            super(EAP_AT_MAC, lengthInBytes, byteBuffer);
 
             if (lengthInBytes != ATTR_LENGTH) {
                 throw new EapSimAkaInvalidAttributeException("Invalid Length specified");
             }
 
-            // next two bytes are reserved (RFC 4186 Section 10.14)
-            byteBuffer.get(new byte[RESERVED_BYTES]);
-
             mac = new byte[MAC_LENGTH];
             byteBuffer.get(mac);
         }
 
-        // Used for calculating MACs. Per RFC 4186 Section 10.14, the MAC should be calculated over
-        // the entire packet, with the value field of the MAC attribute set to zero.
+        // Constructs an AtMac with an empty MAC and empty RESERVED bytes. Should only be used for
+        // calculating MACs in outbound messages.
         public AtMac() throws EapSimAkaInvalidAttributeException {
-            super(EAP_AT_MAC, ATTR_LENGTH);
-            mac = new byte[MAC_LENGTH];
+            this(new byte[MAC_LENGTH]);
         }
 
         public AtMac(byte[] mac) throws EapSimAkaInvalidAttributeException {
-            super(EAP_AT_MAC, ATTR_LENGTH);
-            this.mac = mac;
+            this(new byte[RESERVED_BYTES_LEN], mac);
+        }
+
+        @VisibleForTesting
+        public AtMac(byte[] reservedBytes, byte[] mac) throws EapSimAkaInvalidAttributeException {
+            super(EAP_AT_MAC, ATTR_LENGTH, reservedBytes);
 
             if (mac.length != MAC_LENGTH) {
                 throw new EapSimAkaInvalidAttributeException("Invalid length for MAC");
             }
+            this.mac = mac;
         }
 
         @Override
         public void encode(ByteBuffer byteBuffer) {
-            encodeAttributeHeader(byteBuffer);
-            byteBuffer.put(new byte[RESERVED_BYTES]);
+            super.encode(byteBuffer);
+
             byteBuffer.put(mac);
         }
+
+        /**
+         * Returns a copy of this AtMac with the MAC cleared (and the reserved bytes preserved).
+         *
+         * <p>Per RFC 4186 Section 10.14, the MAC should be calculated over the entire packet, with
+         * the value field of the MAC attribute set to zero.
+         */
+        public AtMac getAtMacWithMacCleared() throws EapSimAkaInvalidAttributeException {
+            return new AtMac(reservedBytes, new byte[MAC_LENGTH]);
+        }
     }
 
     /**
@@ -712,23 +759,20 @@
      *
      * <p>This Nonce is generated by the server and used for fast re-authentication only.
      */
-    public static class AtNonceS extends EapSimAkaAttribute {
+    public static class AtNonceS extends EapSimAkaReservedBytesAttribute {
         private static final int ATTR_LENGTH = 5 * LENGTH_SCALING;
         private static final int NONCE_S_LENGTH = 4 * LENGTH_SCALING;
-        private static final int RESERVED_BYTES = 2;
 
         public final byte[] nonceS = new byte[NONCE_S_LENGTH];
 
         public AtNonceS(int lengthInBytes, ByteBuffer byteBuffer)
                 throws EapSimAkaInvalidAttributeException {
-            super(EAP_AT_NONCE_S, lengthInBytes);
+            super(EAP_AT_NONCE_S, lengthInBytes, byteBuffer);
 
             if (lengthInBytes != ATTR_LENGTH) {
                 throw new EapSimAkaInvalidAttributeException("Invalid Length specified");
             }
 
-            // next two bytes are reserved (RFC 4186 Section 10.17)
-            byteBuffer.get(new byte[RESERVED_BYTES]);
             byteBuffer.get(nonceS);
         }
 
@@ -745,8 +789,8 @@
 
         @Override
         public void encode(ByteBuffer byteBuffer) {
-            encodeAttributeHeader(byteBuffer);
-            byteBuffer.put(new byte[RESERVED_BYTES]);
+            super.encode(byteBuffer);
+
             byteBuffer.put(nonceS);
         }
     }
@@ -888,24 +932,20 @@
     /**
      * AtAutn represents the AT_AUTN attribute defined in RFC 4187#10.7
      */
-    public static class AtAutn extends EapSimAkaAttribute {
+    public static class AtAutn extends EapSimAkaReservedBytesAttribute {
         private static final int ATTR_LENGTH = 5 * LENGTH_SCALING;
         private static final int AUTN_LENGTH = 16;
-        private static final int RESERVED_BYTES = 2;
 
         public final byte[] autn = new byte[AUTN_LENGTH];
 
         public AtAutn(int lengthInBytes, ByteBuffer byteBuffer)
                 throws EapSimAkaInvalidAttributeException {
-            super(EAP_AT_AUTN, lengthInBytes);
+            super(EAP_AT_AUTN, lengthInBytes, byteBuffer);
 
             if (lengthInBytes != ATTR_LENGTH) {
                 throw new EapSimAkaInvalidAttributeException("Length must be 20B");
             }
 
-            // next two bytes are reserved (RFC 4187#10.7)
-            byteBuffer.get(new byte[RESERVED_BYTES]);
-
             byteBuffer.get(autn);
         }
 
@@ -922,8 +962,8 @@
 
         @Override
         public void encode(ByteBuffer byteBuffer) {
-            encodeAttributeHeader(byteBuffer);
-            byteBuffer.put(new byte[RESERVED_BYTES]);
+            super.encode(byteBuffer);
+
             byteBuffer.put(autn);
         }
     }
diff --git a/src/java/com/android/internal/net/eap/message/simaka/EapSimAkaTypeData.java b/src/java/com/android/internal/net/eap/message/simaka/EapSimAkaTypeData.java
index 56d2447..b90cd0b 100644
--- a/src/java/com/android/internal/net/eap/message/simaka/EapSimAkaTypeData.java
+++ b/src/java/com/android/internal/net/eap/message/simaka/EapSimAkaTypeData.java
@@ -39,19 +39,30 @@
  */
 public abstract class EapSimAkaTypeData {
     private static final int MIN_LEN_BYTES = 3; // subtype (1B) + reserved bytes (2B)
-    private static final int RESERVED_BYTES = 2; // RFC 4186#8.1, RFC 4187#8.1
+    private static final int RESERVED_BYTES_LEN = 2; // RFC 4186#8.1, RFC 4187#8.1
 
     public final int eapSubtype;
 
+    /** Save (but ignore) the 2B reserved section after EAP-Type and EAP-SIM/AKA subtype. */
+    final byte[] mReservedBytes;
+
     // LinkedHashMap used to preserve encoded ordering of attributes. This is necessary for checking
     // the MAC value for the message
     public final LinkedHashMap<Integer, EapSimAkaAttribute> attributeMap;
 
+    protected EapSimAkaTypeData(
+            int eapSubType, LinkedHashMap<Integer, EapSimAkaAttribute> attributeMap) {
+        this(eapSubType, attributeMap, new byte[RESERVED_BYTES_LEN]);
+    }
+
     public EapSimAkaTypeData(
             int eapSubType,
-            LinkedHashMap<Integer, EapSimAkaAttribute> attributeMap) {
+            LinkedHashMap<Integer, EapSimAkaAttribute> attributeMap,
+            byte[] reservedBytes) {
         this.eapSubtype = eapSubType;
         this.attributeMap = attributeMap;
+
+        mReservedBytes = reservedBytes.clone();
     }
 
     /**
@@ -69,7 +80,7 @@
         output.put((byte) eapSubtype);
 
         // two reserved bytes (RFC 4186#8.1, RFC 4187#8.1)
-        output.put(new byte[RESERVED_BYTES]);
+        output.put(mReservedBytes);
 
         for (EapSimAkaAttribute attribute : attributeMap.values()) {
             attribute.encode(output);
@@ -118,7 +129,8 @@
                 }
 
                 // next two bytes are reserved (RFC 4186#8.1, RFC 4187#8.1)
-                byteBuffer.get(new byte[RESERVED_BYTES]);
+                byte[] reservedBytes = new byte[RESERVED_BYTES_LEN];
+                byteBuffer.get(reservedBytes);
 
                 // read attributes
                 LinkedHashMap<Integer, EapSimAkaAttribute> attributeMap = new LinkedHashMap<>();
@@ -138,7 +150,7 @@
                     attributeMap.put(attribute.attributeType, attribute);
                 }
 
-                T eapSimAkaTypeData = getInstance(eapSubType, attributeMap);
+                T eapSimAkaTypeData = getInstance(eapSubType, attributeMap, reservedBytes);
 
                 logDecodedEapSimAkaTypeData(eapSimAkaTypeData);
 
@@ -157,7 +169,8 @@
 
         protected abstract T getInstance(
                 int eapSubType,
-                LinkedHashMap<Integer, EapSimAkaAttribute> attributeMap);
+                LinkedHashMap<Integer, EapSimAkaAttribute> attributeMap,
+                byte[] reservedBytes);
 
         private void logDecodedEapSimAkaTypeData(EapSimAkaTypeData eapSimAkaTypeData) {
             StringBuilder msg = new StringBuilder();
diff --git a/src/java/com/android/internal/net/eap/message/simaka/EapSimTypeData.java b/src/java/com/android/internal/net/eap/message/simaka/EapSimTypeData.java
index ff1be60..09e03c7 100644
--- a/src/java/com/android/internal/net/eap/message/simaka/EapSimTypeData.java
+++ b/src/java/com/android/internal/net/eap/message/simaka/EapSimTypeData.java
@@ -67,6 +67,13 @@
         super(eapSubType, attributeMap);
     }
 
+    private EapSimTypeData(
+            int eapSubType,
+            LinkedHashMap<Integer, EapSimAkaAttribute> attributeMap,
+            byte[] reservedBytes) {
+        super(eapSubType, attributeMap, reservedBytes);
+    }
+
     public EapSimTypeData(int eapSubtype, List<EapSimAkaAttribute> attributes) {
         super(eapSubtype, new LinkedHashMap<>());
 
@@ -118,8 +125,9 @@
         @Override
         protected EapSimTypeData getInstance(
                 int eapSubtype,
-                LinkedHashMap<Integer, EapSimAkaAttribute> attributeMap) {
-            return new EapSimTypeData(eapSubtype, attributeMap);
+                LinkedHashMap<Integer, EapSimAkaAttribute> attributeMap,
+                byte[] reservedBytes) {
+            return new EapSimTypeData(eapSubtype, attributeMap, reservedBytes);
         }
     }
 }
diff --git a/src/java/com/android/internal/net/eap/statemachine/EapSimAkaMethodStateMachine.java b/src/java/com/android/internal/net/eap/statemachine/EapSimAkaMethodStateMachine.java
index b2459ab..8089ea7 100644
--- a/src/java/com/android/internal/net/eap/statemachine/EapSimAkaMethodStateMachine.java
+++ b/src/java/com/android/internal/net/eap/statemachine/EapSimAkaMethodStateMachine.java
@@ -243,7 +243,7 @@
 
         // cache original Mac so it can be restored after calculating the Mac
         AtMac originalMac = (AtMac) typeData.attributeMap.get(EAP_AT_MAC);
-        typeData.attributeMap.put(EAP_AT_MAC, new AtMac());
+        typeData.attributeMap.put(EAP_AT_MAC, originalMac.getAtMacWithMacCleared());
 
         byte[] typeDataWithEmptyMac = typeData.encode();
         EapData eapData = new EapData(getEapMethod(), typeDataWithEmptyMac);
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 84339c9..1163486 100644
--- a/src/java/com/android/internal/net/ipsec/ike/IkeSessionStateMachine.java
+++ b/src/java/com/android/internal/net/ipsec/ike/IkeSessionStateMachine.java
@@ -884,6 +884,12 @@
 
     /** Switch all IKE SAs to the new IKE socket due to an underlying network change. */
     private void switchToIkeSocket(IkeSocket newSocket) {
+        // Changing IkeSockets - make sure to quit NAT-T keepalive if it's going
+        if (mIkeNattKeepalive != null) {
+            mIkeNattKeepalive.stop();
+            mIkeNattKeepalive = null;
+        }
+
         long currentLocalSpi = mCurrentIkeSaRecord.getLocalSpi();
         migrateSpiToIkeSocket(currentLocalSpi, mIkeSocket, newSocket);
 
@@ -900,6 +906,28 @@
         mIkeSocket = newSocket;
     }
 
+    private void buildAndSwitchToIkeSocketWithPort4500(boolean isIpv4) {
+        try {
+            if (isIpv4) {
+                IkeSocket newSocket =
+                        IkeUdpEncapSocket.getIkeUdpEncapSocket(
+                                mNetwork,
+                                mIpSecManager,
+                                IkeSessionStateMachine.this,
+                                getHandler().getLooper());
+                switchToIkeSocket(newSocket);
+                mIkeNattKeepalive = buildAndStartNattKeepalive();
+            } else {
+                IkeSocket newSocket =
+                        IkeUdp6WithEncapPortSocket.getInstance(
+                                mNetwork, IkeSessionStateMachine.this, getHandler());
+                switchToIkeSocket(newSocket);
+            }
+        } catch (ErrnoException | IOException | ResourceUnavailableException e) {
+            handleIkeFatalError(e);
+        }
+    }
+
     private void migrateSpiToIkeSocket(long localSpi, IkeSocket oldSocket, IkeSocket newSocket) {
         newSocket.registerIke(localSpi, IkeSessionStateMachine.this);
         oldSocket.unregisterIke(localSpi);
@@ -3230,20 +3258,11 @@
                                     IkeSessionStateMachine.this,
                                     getHandler().getLooper());
                     switchToIkeSocket(initIkeSpi, newSocket);
+                    mIkeNattKeepalive = buildAndStartNattKeepalive();
+                    mLocalPort = mIkeSocket.getLocalPort();
                 } catch (ErrnoException | IOException | ResourceUnavailableException e) {
                     handleIkeFatalError(e);
                 }
-
-                mIkeNattKeepalive =
-                        new IkeNattKeepalive(
-                                mContext,
-                                NATT_KEEPALIVE_DELAY_SECONDS,
-                                (Inet4Address) mLocalAddress,
-                                (Inet4Address) mRemoteAddress,
-                                ((IkeUdpEncapSocket) mIkeSocket).getUdpEncapsulationSocket(),
-                                mIkeSocket.getNetwork(),
-                                buildKeepaliveIntent());
-                mIkeNattKeepalive.start();
             }
         }
 
@@ -3254,14 +3273,6 @@
             mIkeSocket = newSocket;
         }
 
-        private PendingIntent buildKeepaliveIntent() {
-            return buildIkeAlarmIntent(
-                    mContext,
-                    ACTION_KEEPALIVE,
-                    getIntentIdentifier(),
-                    obtainMessage(CMD_ALARM_FIRED, mIkeSessionId, CMD_SEND_KEEPALIVE));
-        }
-
         @Override
         public void exitState() {
             super.exitState();
@@ -3339,6 +3350,34 @@
         }
     }
 
+    /** Starts NAT-T keepalive for current IkeUdpEncapSocket */
+    private IkeNattKeepalive buildAndStartNattKeepalive() throws IOException {
+        if (!(mIkeSocket instanceof IkeUdpEncapSocket)) {
+            throw new IllegalStateException(
+                    "Cannot start NAT-T keepalive when IKE Session is not using UDP Encap socket");
+        }
+
+        PendingIntent keepaliveIntent =
+                buildIkeAlarmIntent(
+                        mContext,
+                        ACTION_KEEPALIVE,
+                        getIntentIdentifier(),
+                        obtainMessage(CMD_ALARM_FIRED, mIkeSessionId, CMD_SEND_KEEPALIVE));
+
+        IkeNattKeepalive keepalive =
+                new IkeNattKeepalive(
+                        mContext,
+                        mConnectivityManager,
+                        NATT_KEEPALIVE_DELAY_SECONDS,
+                        (Inet4Address) mLocalAddress,
+                        (Inet4Address) mRemoteAddress,
+                        ((IkeUdpEncapSocket) mIkeSocket).getUdpEncapsulationSocket(),
+                        mIkeSocket.getNetwork(),
+                        keepaliveIntent);
+        keepalive.start();
+        return keepalive;
+    }
+
     /**
      * CreateIkeLocalIkeAuthBase represents the common state and functionality required to perform
      * IKE AUTH exchanges in both the EAP and non-EAP flows.
@@ -3490,8 +3529,6 @@
                 mSupportMobike = true;
                 mEnabledExtensions.add(EXTENSION_TYPE_MOBIKE);
 
-                // TODO(b/173237734): use port 4500 if NAT-T is enabled
-
                 try {
                     if (mIkeSessionParams.getConfiguredNetwork() != null) {
                         // Caller configured a specific Network - track it
@@ -3512,6 +3549,17 @@
                     // Error occurred while registering the NetworkCallback
                     throw new IkeInternalException("Error while registering NetworkCallback", e);
                 }
+
+                // Use port 4500 if NAT-T is enabled
+                if (mSupportNatTraversal && !(mIkeSocket instanceof IkeUdpEncapSocket)) {
+                    buildAndSwitchToIkeSocketWithPort4500(mIkeSocket instanceof IkeUdp4Socket);
+                    try {
+                        mLocalPort = mIkeSocket.getLocalPort();
+                    } catch (ErrnoException e) {
+                        throw new IkeInternalException(e);
+                    }
+                }
+
                 return;
             } else {
                 // Unknown and unexpected status notifications are ignored as per
@@ -5372,27 +5420,24 @@
             // Only switch the IkeSocket if the underlying Network actually changes. This may not
             // always happen (ex: the underlying Network loses the current local address)
             if (!mNetwork.equals(oldNetwork)) {
-                // Changing IkeSockets - make sure to quit NAT-T keepalive if it's going
-                if (mIkeNattKeepalive != null) {
-                    mIkeNattKeepalive.stop();
-                    mIkeNattKeepalive = null;
-                }
-
-                IkeSocket newSocket;
-                // TODO(b/173237734): use port 4500 if NAT-T is enabled
-                if (isIpv4) {
-                    newSocket =
-                            IkeUdp4Socket.getInstance(
-                                    mNetwork, IkeSessionStateMachine.this, getHandler());
+                // Use port 4500 if NAT-T is supported by both sides
+                if (mSupportNatTraversal) {
+                    buildAndSwitchToIkeSocketWithPort4500(isIpv4);
                 } else {
-                    newSocket =
-                            IkeUdp6Socket.getInstance(
-                                    mNetwork, IkeSessionStateMachine.this, getHandler());
+                    IkeSocket newSocket;
+                    if (isIpv4) {
+                        newSocket =
+                                IkeUdp4Socket.getInstance(
+                                        mNetwork, IkeSessionStateMachine.this, getHandler());
+                    } else {
+                        newSocket =
+                                IkeUdp6Socket.getInstance(
+                                        mNetwork, IkeSessionStateMachine.this, getHandler());
+                    }
+                    switchToIkeSocket(newSocket);
                 }
-                switchToIkeSocket(newSocket);
-                mLocalPort = mIkeSocket.getLocalPort();
             }
-
+            mLocalPort = mIkeSocket.getLocalPort();
             mLocalAddress =
                     mIkeLocalAddressGenerator.generateLocalAddress(
                             mNetwork, isIpv4, mRemoteAddress, mIkeSocket.getIkeServerPort());
diff --git a/src/java/com/android/internal/net/ipsec/ike/keepalive/HardwareKeepaliveImpl.java b/src/java/com/android/internal/net/ipsec/ike/keepalive/HardwareKeepaliveImpl.java
index bb0abe5..28ea3b3 100644
--- a/src/java/com/android/internal/net/ipsec/ike/keepalive/HardwareKeepaliveImpl.java
+++ b/src/java/com/android/internal/net/ipsec/ike/keepalive/HardwareKeepaliveImpl.java
@@ -49,6 +49,7 @@
     /** Construct an instance of HardwareKeepaliveImpl */
     public HardwareKeepaliveImpl(
             Context context,
+            ConnectivityManager connectMgr,
             int keepaliveDelaySeconds,
             Inet4Address src,
             Inet4Address dest,
@@ -61,10 +62,8 @@
         mKeepaliveDelaySeconds = keepaliveDelaySeconds;
         mHardwareKeepaliveCb = hardwareKeepaliveCb;
 
-        ConnectivityManager connMgr =
-                (ConnectivityManager) context.getSystemService(Context.CONNECTIVITY_SERVICE);
         mSocketKeepalive =
-                connMgr.createSocketKeepalive(
+                connectMgr.createSocketKeepalive(
                         network,
                         socket,
                         src,
diff --git a/src/java/com/android/internal/net/ipsec/ike/keepalive/IkeNattKeepalive.java b/src/java/com/android/internal/net/ipsec/ike/keepalive/IkeNattKeepalive.java
index 129a1f0..00c84c7 100644
--- a/src/java/com/android/internal/net/ipsec/ike/keepalive/IkeNattKeepalive.java
+++ b/src/java/com/android/internal/net/ipsec/ike/keepalive/IkeNattKeepalive.java
@@ -20,6 +20,7 @@
 
 import android.app.PendingIntent;
 import android.content.Context;
+import android.net.ConnectivityManager;
 import android.net.IpSecManager.UdpEncapsulationSocket;
 import android.net.Network;
 
@@ -40,6 +41,7 @@
     /** Construct an instance of IkeNattKeepalive */
     public IkeNattKeepalive(
             Context context,
+            ConnectivityManager connectMgr,
             int keepaliveDelaySeconds,
             Inet4Address src,
             Inet4Address dest,
@@ -50,6 +52,7 @@
         mNattKeepalive =
                 new HardwareKeepaliveImpl(
                         context,
+                        connectMgr,
                         keepaliveDelaySeconds,
                         src,
                         dest,
diff --git a/tests/iketests/src/java/android/net/ipsec/ike/IkeSessionCallbackTest.java b/tests/iketests/src/java/android/net/ipsec/ike/IkeSessionCallbackTest.java
new file mode 100644
index 0000000..5a37ea4
--- /dev/null
+++ b/tests/iketests/src/java/android/net/ipsec/ike/IkeSessionCallbackTest.java
@@ -0,0 +1,92 @@
+/*
+ * 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;
+
+import static org.junit.Assert.assertEquals;
+
+import android.net.ipsec.ike.exceptions.IkeException;
+import android.net.ipsec.ike.exceptions.IkeProtocolException;
+import android.net.ipsec.ike.exceptions.InvalidIkeSpiException;
+
+import org.junit.Before;
+import org.junit.Test;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+
+public final class IkeSessionCallbackTest {
+    private OldOnErrorIkeSessionCallback mOldCallback;
+    private UpdatedOnErrorIkeSessionCallback mUpdatedCallback;
+    private IkeProtocolException mIkeException;
+
+    @Before
+    public void setUp() {
+        mOldCallback = new OldOnErrorIkeSessionCallback();
+        mUpdatedCallback = new UpdatedOnErrorIkeSessionCallback();
+        mIkeException = new InvalidIkeSpiException();
+    }
+
+    @Test
+    public void testOnErrorIkeExceptionNotOverridden() {
+        mOldCallback.onError((IkeException) mIkeException);
+        assertEquals(Arrays.asList(mIkeException), mOldCallback.mOnErrorIkeProtocolExceptions);
+    }
+
+    @Test
+    public void testOnErrorIkeExceptionOverridden() {
+        mUpdatedCallback.onError((IkeException) mIkeException);
+        assertEquals(Arrays.asList(mIkeException), mUpdatedCallback.mOnErrorIkeExceptions);
+    }
+
+    private abstract class TestIkeSessionCallbackBase implements IkeSessionCallback {
+        @Override
+        public void onOpened(IkeSessionConfiguration sessionConfiguration) {}
+
+        @Override
+        public void onClosed() {}
+
+        @Override
+        public void onClosedExceptionally(IkeException exception) {}
+
+        @Override
+        public void onIkeSessionConnectionInfoChanged(IkeSessionConnectionInfo connectionInfo) {}
+    }
+
+    private final class OldOnErrorIkeSessionCallback extends TestIkeSessionCallbackBase {
+        List<IkeProtocolException> mOnErrorIkeProtocolExceptions = new ArrayList<>();
+
+        @Override
+        public void onError(IkeProtocolException exception) {
+            mOnErrorIkeProtocolExceptions.add(exception);
+        }
+    }
+
+    private final class UpdatedOnErrorIkeSessionCallback extends TestIkeSessionCallbackBase {
+        List<IkeException> mOnErrorIkeExceptions = new ArrayList<>();
+
+        @Override
+        public void onError(IkeProtocolException exception) {
+            throw new UnsupportedOperationException();
+        }
+
+        @Override
+        public void onError(IkeException exception) {
+            mOnErrorIkeExceptions.add(exception);
+        }
+    }
+}
diff --git a/tests/iketests/src/java/com/android/internal/net/eap/message/simaka/EapAkaPrimeTypeDataTest.java b/tests/iketests/src/java/com/android/internal/net/eap/message/simaka/EapAkaPrimeTypeDataTest.java
index 6a9e517..d5aa9d1 100644
--- a/tests/iketests/src/java/com/android/internal/net/eap/message/simaka/EapAkaPrimeTypeDataTest.java
+++ b/tests/iketests/src/java/com/android/internal/net/eap/message/simaka/EapAkaPrimeTypeDataTest.java
@@ -56,14 +56,25 @@
     private static final byte[] AUTN_BYTES = hexStringToByteArray(AUTN);
     private static final String MAC = "95FEB9E70427F34B4FAC8F2C7A65A302";
     private static final byte[] MAC_BYTES = hexStringToByteArray(MAC);
+
     private static final byte[] EAP_AKA_PRIME_CHALLENGE_REQUEST =
             hexStringToByteArray(
+                    "010A0B" // Challenge | 2B padding
+                            + "01051A1B" + RAND // AT_RAND attribute
+                            + "02052A2B" + AUTN // AT_AUTN attribute
+                            + "1704000B" + NETWORK_NAME_HEX + "00" // AT_KDF_INPUT
+                            + "18010001" // AT_KDF
+                            + "0B053A3B" + MAC); // AT_MAC attribute
+
+    private static final byte[] EAP_AKA_PRIME_CHALLENGE_REQUEST_EMPTY_RESERVED =
+            hexStringToByteArray(
                     "010000" // Challenge | 2B padding
                             + "01050000" + RAND // AT_RAND attribute
                             + "02050000" + AUTN // AT_AUTN attribute
                             + "1704000B" + NETWORK_NAME_HEX + "00" // AT_KDF_INPUT
                             + "18010001" // AT_KDF
                             + "0B050000" + MAC); // AT_MAC attribute
+
     private static final byte[] EAP_AKA_PRIME_MULTIPLE_AT_KDF =
             hexStringToByteArray(
                     "010000" // Challenge | 2B padding
@@ -117,6 +128,15 @@
     }
 
     @Test
+    public void testDecodeEncode() {
+        DecodeResult<EapAkaTypeData> result =
+                mTypeDataDecoder.decode(EAP_AKA_PRIME_CHALLENGE_REQUEST);
+        assertTrue(result.isSuccessfulDecode());
+
+        assertArrayEquals(EAP_AKA_PRIME_CHALLENGE_REQUEST, result.eapTypeData.encode());
+    }
+
+    @Test
     public void testDecodeMultipleAtKdfAttributes() {
         DecodeResult<EapAkaTypeData> result =
                 mTypeDataDecoder.decode(EAP_AKA_PRIME_MULTIPLE_AT_KDF);
@@ -137,6 +157,6 @@
                 new EapAkaPrimeTypeData(EAP_AKA_CHALLENGE, attributes);
 
         byte[] result = eapAkaPrimeTypeData.encode();
-        assertArrayEquals(EAP_AKA_PRIME_CHALLENGE_REQUEST, result);
+        assertArrayEquals(EAP_AKA_PRIME_CHALLENGE_REQUEST_EMPTY_RESERVED, result);
     }
 }
diff --git a/tests/iketests/src/java/com/android/internal/net/eap/message/simaka/EapAkaTypeDataTest.java b/tests/iketests/src/java/com/android/internal/net/eap/message/simaka/EapAkaTypeDataTest.java
index b2d89d1..cb27cd1 100644
--- a/tests/iketests/src/java/com/android/internal/net/eap/message/simaka/EapAkaTypeDataTest.java
+++ b/tests/iketests/src/java/com/android/internal/net/eap/message/simaka/EapAkaTypeDataTest.java
@@ -68,12 +68,12 @@
     private static final byte[] MAC_BYTES = hexStringToByteArray(MAC);
     private static final byte[] EAP_AKA_REQUEST =
             hexStringToByteArray(
-                    "010000" // Challenge | 2B padding
-                            + "01050000" + RAND // AT_RAND attribute
-                            + "02050000" + AUTN // AT_AUTN attribute
+                    "010A0B" // Challenge | 2B reserved
+                            + "01051A1B" + RAND // AT_RAND attribute
+                            + "02052A2B" + AUTN // AT_AUTN attribute
                             + "8B010002" // AT_RESULT_IND attribute (TS 124 302#8.2.3.1)
-                            + "0B050000" + MAC // AT_MAC attribute
-                            + "86010000"); // AT_CHECKCODE attribute
+                            + "0B053A3B" + MAC // AT_MAC attribute
+                            + "86014A4B"); // AT_CHECKCODE attribute
 
     private EapAkaTypeDataDecoder mEapAkaTypeDataDecoder;
 
@@ -106,6 +106,14 @@
     }
 
     @Test
+    public void testDecodeEncode() {
+        DecodeResult<EapAkaTypeData> result = mEapAkaTypeDataDecoder.decode(EAP_AKA_REQUEST);
+        assertTrue(result.isSuccessfulDecode());
+
+        assertArrayEquals(EAP_AKA_REQUEST, result.eapTypeData.encode());
+    }
+
+    @Test
     public void testDecodeWithOptionalAttributes() {
         DecodeResult<EapAkaTypeData> result = mEapAkaTypeDataDecoder.decode(EAP_AKA_REQUEST);
 
diff --git a/tests/iketests/src/java/com/android/internal/net/eap/message/simaka/EapSimTypeDataTest.java b/tests/iketests/src/java/com/android/internal/net/eap/message/simaka/EapSimTypeDataTest.java
index 678a812..bd6352a 100644
--- a/tests/iketests/src/java/com/android/internal/net/eap/message/simaka/EapSimTypeDataTest.java
+++ b/tests/iketests/src/java/com/android/internal/net/eap/message/simaka/EapSimTypeDataTest.java
@@ -16,6 +16,7 @@
 
 package com.android.internal.net.eap.message.simaka;
 
+import static com.android.internal.net.TestUtils.hexStringToByteArray;
 import static com.android.internal.net.eap.message.EapTestMessageDefinitions.EAP_SIM_START_DUPLICATE_ATTRIBUTES;
 import static com.android.internal.net.eap.message.EapTestMessageDefinitions.EAP_SIM_START_SUBTYPE;
 import static com.android.internal.net.eap.message.EapTestMessageDefinitions.INVALID_SUBTYPE;
@@ -24,6 +25,8 @@
 import static com.android.internal.net.eap.message.EapTestMessageDefinitions.TYPE_DATA_INVALID_AT_RAND;
 import static com.android.internal.net.eap.message.simaka.EapSimAkaAttribute.EAP_AT_PERMANENT_ID_REQ;
 import static com.android.internal.net.eap.message.simaka.EapSimAkaAttribute.EAP_AT_VERSION_LIST;
+import static com.android.internal.net.eap.message.simaka.attributes.EapTestAttributeDefinitions.RAND_1;
+import static com.android.internal.net.eap.message.simaka.attributes.EapTestAttributeDefinitions.RAND_2;
 
 import static junit.framework.TestCase.fail;
 
@@ -52,6 +55,12 @@
     private static final int EAP_SIM_START = 10;
     private static final int INVALID_SUBTYPE_INT = -1;
 
+    private static final byte[] EAP_SIM_CHALLENGE_REQUEST =
+            hexStringToByteArray(
+                    "0b0A0B" // Challenge | 2B padding
+                            + "01091A1B" + RAND_1 + RAND_2 // EAP-SIM AT_RAND attribute
+                            + "0B052A2BFFEEDDCCBBAA998877665544332211FF"); // AT_MAC attribute
+
     private EapSimTypeDataDecoder mEapSimTypeDataDecoder;
 
     @Before
@@ -100,6 +109,15 @@
     }
 
     @Test
+    public void testDecodeEncode() {
+        DecodeResult<EapSimTypeData> result =
+                mEapSimTypeDataDecoder.decode(EAP_SIM_CHALLENGE_REQUEST);
+        assertTrue(result.isSuccessfulDecode());
+
+        assertArrayEquals(EAP_SIM_CHALLENGE_REQUEST, result.eapTypeData.encode());
+    }
+
+    @Test
     public void testDecodeNullTypeData() {
         DecodeResult<EapSimTypeData> result = mEapSimTypeDataDecoder.decode(null);
         assertFalse(result.isSuccessfulDecode());
diff --git a/tests/iketests/src/java/com/android/internal/net/eap/message/simaka/attributes/AtMacTest.java b/tests/iketests/src/java/com/android/internal/net/eap/message/simaka/attributes/AtMacTest.java
index 82b066d..042519a 100644
--- a/tests/iketests/src/java/com/android/internal/net/eap/message/simaka/attributes/AtMacTest.java
+++ b/tests/iketests/src/java/com/android/internal/net/eap/message/simaka/attributes/AtMacTest.java
@@ -43,6 +43,7 @@
     private static final int MAC_LENGTH = 16;
     private static final byte[] MAC_BYTES = hexStringToByteArray(MAC);
     private static final byte[] INVALID_MAC = {(byte) 1, (byte) 2, (byte) 3};
+    private static final byte[] RESERVED_BYTES = {(byte) 0x0A, (byte) 0x0B};
 
     private EapSimAkaAttributeFactory mAttributeFactory;
 
@@ -107,4 +108,13 @@
         atMac.encode(result);
         assertArrayEquals(AT_MAC, result.array());
     }
+
+    @Test
+    public void testGetAtMacWithMacCleared() throws Exception {
+        AtMac original = new AtMac(RESERVED_BYTES, MAC_BYTES);
+
+        AtMac clearedMac = original.getAtMacWithMacCleared();
+        assertArrayEquals(RESERVED_BYTES, clearedMac.reservedBytes);
+        assertArrayEquals(new byte[MAC_LENGTH], clearedMac.mac);
+    }
 }
diff --git a/tests/iketests/src/java/com/android/internal/net/eap/statemachine/EapSimAkaMethodStateMachineTest.java b/tests/iketests/src/java/com/android/internal/net/eap/statemachine/EapSimAkaMethodStateMachineTest.java
index bf9c660..686555b 100644
--- a/tests/iketests/src/java/com/android/internal/net/eap/statemachine/EapSimAkaMethodStateMachineTest.java
+++ b/tests/iketests/src/java/com/android/internal/net/eap/statemachine/EapSimAkaMethodStateMachineTest.java
@@ -17,6 +17,7 @@
 package com.android.internal.net.eap.statemachine;
 
 import static com.android.internal.net.TestUtils.hexStringToByteArray;
+import static com.android.internal.net.eap.message.EapData.EAP_TYPE_AKA;
 import static com.android.internal.net.eap.message.EapData.EAP_TYPE_SIM;
 import static com.android.internal.net.eap.message.EapMessage.EAP_CODE_REQUEST;
 import static com.android.internal.net.eap.message.EapTestMessageDefinitions.COMPUTED_MAC;
@@ -46,8 +47,13 @@
 import static com.android.internal.net.eap.message.EapTestMessageDefinitions.ORIGINAL_MAC;
 import static com.android.internal.net.eap.message.EapTestMessageDefinitions.SRES_1;
 import static com.android.internal.net.eap.message.EapTestMessageDefinitions.SRES_BYTES;
+import static com.android.internal.net.eap.message.simaka.EapAkaTypeData.EAP_AKA_CHALLENGE;
+import static com.android.internal.net.eap.message.simaka.EapAkaTypeData.EAP_AKA_CLIENT_ERROR;
 import static com.android.internal.net.eap.message.simaka.EapSimAkaAttribute.AtNotification.GENERAL_FAILURE_POST_CHALLENGE;
 import static com.android.internal.net.eap.message.simaka.EapSimAkaAttribute.AtNotification.GENERAL_FAILURE_PRE_CHALLENGE;
+import static com.android.internal.net.eap.message.simaka.EapSimAkaAttribute.EAP_AT_CHECKCODE;
+import static com.android.internal.net.eap.message.simaka.EapSimAkaAttribute.EAP_AT_ENCR_DATA;
+import static com.android.internal.net.eap.message.simaka.EapSimAkaAttribute.EAP_AT_IV;
 import static com.android.internal.net.eap.message.simaka.EapSimAkaAttribute.EAP_AT_MAC;
 import static com.android.internal.net.eap.message.simaka.EapSimTypeData.EAP_SIM_CHALLENGE;
 import static com.android.internal.net.eap.message.simaka.EapSimTypeData.EAP_SIM_CLIENT_ERROR;
@@ -80,6 +86,7 @@
 import static org.mockito.Mockito.verifyNoMoreInteractions;
 import static org.mockito.Mockito.when;
 
+import android.net.eap.EapSessionConfig.EapAkaConfig;
 import android.net.eap.EapSessionConfig.EapSimConfig;
 import android.telephony.TelephonyManager;
 
@@ -91,13 +98,17 @@
 import com.android.internal.net.eap.exceptions.simaka.EapSimAkaAuthenticationFailureException;
 import com.android.internal.net.eap.message.EapData;
 import com.android.internal.net.eap.message.EapMessage;
+import com.android.internal.net.eap.message.simaka.EapAkaTypeData;
 import com.android.internal.net.eap.message.simaka.EapSimAkaAttribute;
+import com.android.internal.net.eap.message.simaka.EapSimAkaAttribute.AtAutn;
 import com.android.internal.net.eap.message.simaka.EapSimAkaAttribute.AtClientErrorCode;
 import com.android.internal.net.eap.message.simaka.EapSimAkaAttribute.AtIdentity;
 import com.android.internal.net.eap.message.simaka.EapSimAkaAttribute.AtMac;
 import com.android.internal.net.eap.message.simaka.EapSimAkaAttribute.AtNotification;
+import com.android.internal.net.eap.message.simaka.EapSimAkaAttribute.AtRandAka;
 import com.android.internal.net.eap.message.simaka.EapSimAkaAttribute.AtRandSim;
 import com.android.internal.net.eap.message.simaka.EapSimAkaAttribute.AtSelectedVersion;
+import com.android.internal.net.eap.message.simaka.EapSimAkaAttribute.EapSimAkaUnsupportedAttribute;
 import com.android.internal.net.eap.message.simaka.EapSimAkaTypeData;
 import com.android.internal.net.eap.message.simaka.EapSimTypeData;
 import com.android.internal.net.eap.statemachine.EapMethodStateMachine.EapMethodState;
@@ -134,18 +145,40 @@
     // K_encr + K_aut + MSK + EMSK
     private static final int PRF_OUTPUT_BYTES = (2 * KEY_LEN) + (2 * SESSION_KEY_LENGTH);
 
+    private static final String AKA_RAND = "648EAAB01CA1BFEB9E9708852D445DA5";
+    private static final String AUTN = "80CEABF08239000093281F9A178246B8";
+    private static final String IV_DATA = "3232C4A5A2D97B39BCF55FA7BEFCCBF52D26";
+    private static final String ENCR_DATA =
+            "3566EA8CF174FB4A94488E56B6E8DFC25F05B100BEABDA5DDBAC18968D8158FEDF1F";
+    private static final String AKA_MAC_RESERVED_BYTES = "7469";
+    private static final String AKA_MAC = "5198169B1AC51CA0A193FDEEE7981E16";
+    private static final int AT_CHECKCODE_LENGTH = 4;
+    private static final int AT_IV_LENGTH = 20;
+    private static final int AT_ENCR_DATA_LENGTH = 36; // variable length, not IANA specified
+
+    private static final byte[] EAP_AKA_REQUEST_FOR_MAC =
+            hexStringToByteArray(
+                    "01020080" // EAP-Request | ID | length in bytes
+                            + "17010000" // EAP-AKA | Challenge | 2B padding
+                            + "01050000" + AKA_RAND // EAP-AKA AT_RAND attr
+                            + "02050000" + AUTN // AT_AUTN attr
+                            + "86010000" // AT_CHECKCODE attr
+                            + "8105" + IV_DATA // AT_IV attr
+                            + "8209" + ENCR_DATA // AT_ENCR_DATA attr
+                            + "0B05" + AKA_MAC_RESERVED_BYTES + AKA_MAC); // AT_MAC attr
+
     private TelephonyManager mMockTelephonyManager;
-    private EapSimConfig mEapSimConfig;
     private EapSimAkaMethodStateMachine mStateMachine;
 
     @Before
     public void setUp() {
         mMockTelephonyManager = mock(TelephonyManager.class);
-        mEapSimConfig = new EapSimConfig(SUB_ID, TelephonyManager.APPTYPE_USIM);
 
         mStateMachine =
                 new EapSimAkaMethodStateMachine(
-                        mMockTelephonyManager, EAP_IDENTITY_BYTES, mEapSimConfig) {
+                        mMockTelephonyManager,
+                        EAP_IDENTITY_BYTES,
+                        new EapSimConfig(SUB_ID, TelephonyManager.APPTYPE_USIM)) {
                     @Override
                     EapSimAkaTypeData getEapSimAkaTypeData(AtClientErrorCode clientErrorCode) {
                         return new EapSimTypeData(
@@ -474,4 +507,69 @@
         assertEquals(SESSION_KEY_LENGTH, mStateMachine.getMskLength());
         assertEquals(SESSION_KEY_LENGTH, mStateMachine.getEmskLength());
     }
+
+    @Test
+    public void testIsValidMac() throws Exception {
+        // Data expects an EAP-AKA state machine
+        mStateMachine = buildEapAkaStateMachineWithKAut(K_AUT);
+
+        EapMessage message = EapMessage.decode(EAP_AKA_REQUEST_FOR_MAC);
+
+        AtRandAka atRand = new AtRandAka(hexStringToByteArray(AKA_RAND));
+        AtAutn atAutn = new AtAutn(hexStringToByteArray(AUTN));
+
+        // AT_CHECKCODE is formatted: ATTR_TYPE (1B) + Length (1B) + reserved bytes (2B)
+        EapSimAkaUnsupportedAttribute atCheckcode =
+                new EapSimAkaUnsupportedAttribute(
+                        EAP_AT_CHECKCODE, AT_CHECKCODE_LENGTH, new byte[2]);
+        EapSimAkaUnsupportedAttribute atIv =
+                new EapSimAkaUnsupportedAttribute(
+                        EAP_AT_IV, AT_IV_LENGTH, hexStringToByteArray(IV_DATA));
+        EapSimAkaUnsupportedAttribute atEncrData =
+                new EapSimAkaUnsupportedAttribute(
+                        EAP_AT_ENCR_DATA, AT_ENCR_DATA_LENGTH, hexStringToByteArray(ENCR_DATA));
+        AtMac atMac =
+                new AtMac(
+                        hexStringToByteArray(AKA_MAC_RESERVED_BYTES),
+                        hexStringToByteArray(AKA_MAC));
+        EapSimAkaTypeData typeData =
+                new EapAkaTypeData(
+                        EAP_AKA_CHALLENGE,
+                        Arrays.asList(atRand, atAutn, atCheckcode, atIv, atEncrData, atMac));
+
+        // No extra data for EAP-AKA
+        byte[] extraData = new byte[0];
+
+        assertTrue(mStateMachine.isValidMac("testIsValidMac", message, typeData, extraData));
+    }
+
+    private EapSimAkaMethodStateMachine buildEapAkaStateMachineWithKAut(byte[] kAut) {
+        EapSimAkaMethodStateMachine stateMachine =
+                new EapSimAkaMethodStateMachine(
+                        mMockTelephonyManager,
+                        EAP_IDENTITY_BYTES,
+                        new EapAkaConfig(SUB_ID, TelephonyManager.APPTYPE_USIM)) {
+                    @Override
+                    EapSimAkaTypeData getEapSimAkaTypeData(AtClientErrorCode clientErrorCode) {
+                        return new EapAkaTypeData(
+                                EAP_AKA_CLIENT_ERROR, Arrays.asList(clientErrorCode));
+                    }
+
+                    @Override
+                    EapSimAkaTypeData getEapSimAkaTypeData(
+                            int eapSubtype, List<EapSimAkaAttribute> attributes) {
+                        return new EapAkaTypeData(eapSubtype, attributes);
+                    }
+
+                    @Override
+                    int getEapMethod() {
+                        return EAP_TYPE_AKA;
+                    }
+                };
+
+        // set K_AUT for the state machine
+        System.arraycopy(kAut, 0, stateMachine.mKAut, 0, stateMachine.getKAutLength());
+
+        return stateMachine;
+    }
 }
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 e662294..71dc37e 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
@@ -748,6 +748,9 @@
         when(mMockIkeLocalAddressGenerator.generateLocalAddress(
                         eq(mMockDefaultNetwork), eq(true /* isIpv4 */), any(), anyInt()))
                 .thenReturn(LOCAL_ADDRESS);
+        when(mMockIkeLocalAddressGenerator.generateLocalAddress(
+                        eq(mMockDefaultNetwork), eq(false /* isIpv4 */), any(), anyInt()))
+                .thenReturn(LOCAL_ADDRESS_V6);
 
         mMockEapAuthenticatorFactory = mock(IkeEapAuthenticatorFactory.class);
         mMockEapAuthenticator = mock(EapAuthenticator.class);
@@ -2316,7 +2319,10 @@
                         respIdPayload,
                         authRelatedPayloads,
                         hasChildPayloads,
-                        hasConfigPayloadInResp);
+                        hasConfigPayloadInResp,
+                        false /* isMobikeEnabled */,
+                        true /* isIpv4 */,
+                        0 /* ike3gppCallbackInvocations */);
 
         verify(spyAuthPayload)
                 .verifyInboundSignature(
@@ -2358,6 +2364,27 @@
             boolean isMobikeEnabled,
             int ike3gppCallbackInvocations)
             throws Exception {
+        return verifySharedKeyAuthentication(
+                spyAuthPayload,
+                respIdPayload,
+                authRelatedPayloads,
+                hasChildPayloads,
+                hasConfigPayloadInResp,
+                isMobikeEnabled,
+                true /* isIpv4 */,
+                ike3gppCallbackInvocations);
+    }
+
+    private IkeMessage verifySharedKeyAuthentication(
+            IkeAuthPskPayload spyAuthPayload,
+            IkeIdPayload respIdPayload,
+            List<IkePayload> authRelatedPayloads,
+            boolean hasChildPayloads,
+            boolean hasConfigPayloadInResp,
+            boolean isMobikeEnabled,
+            boolean isIpv4,
+            int ike3gppCallbackInvocations)
+            throws Exception {
         IkeMessage ikeAuthReqMessage =
                 verifyAuthenticationCommonAndGetIkeMessage(
                         respIdPayload,
@@ -2365,6 +2392,7 @@
                         hasChildPayloads,
                         hasConfigPayloadInResp,
                         isMobikeEnabled,
+                        isIpv4,
                         ike3gppCallbackInvocations);
 
         // Validate authentication is done. Cannot use matchers because IkeAuthPskPayload is final.
@@ -2388,23 +2416,9 @@
             IkeIdPayload respIdPayload,
             List<IkePayload> authRelatedPayloads,
             boolean hasChildPayloads,
-            boolean hasConfigPayloadInResp)
-            throws Exception {
-        return verifyAuthenticationCommonAndGetIkeMessage(
-                respIdPayload,
-                authRelatedPayloads,
-                hasChildPayloads,
-                hasConfigPayloadInResp,
-                false /* isMobikeEnabled */,
-                0 /* ike3gppCallbackInvocations */);
-    }
-
-    private IkeMessage verifyAuthenticationCommonAndGetIkeMessage(
-            IkeIdPayload respIdPayload,
-            List<IkePayload> authRelatedPayloads,
-            boolean hasChildPayloads,
             boolean hasConfigPayloadInResp,
             boolean isMobikeEnabled,
+            boolean isIpv4,
             int ike3gppCallbackInvocations)
             throws Exception {
         // Send IKE AUTH response to IKE state machine
@@ -2470,8 +2484,12 @@
                 sessionConfig.isIkeExtensionEnabled(IkeSessionConfiguration.EXTENSION_TYPE_MOBIKE));
 
         IkeSessionConnectionInfo ikeConnInfo = sessionConfig.getIkeSessionConnectionInfo();
-        assertEquals(LOCAL_ADDRESS, ikeConnInfo.getLocalAddress());
-        assertEquals(REMOTE_ADDRESS, ikeConnInfo.getRemoteAddress());
+
+        InetAddress expectedLocalAddress = isIpv4 ? LOCAL_ADDRESS : LOCAL_ADDRESS_V6;
+        InetAddress expectedRemoteAddress = isIpv4 ? REMOTE_ADDRESS : REMOTE_ADDRESS_V6;
+
+        assertEquals(expectedLocalAddress, ikeConnInfo.getLocalAddress());
+        assertEquals(expectedRemoteAddress, ikeConnInfo.getRemoteAddress());
         assertEquals(mMockDefaultNetwork, ikeConnInfo.getNetwork());
 
         // Verify payload list pair for first Child negotiation
@@ -2483,8 +2501,8 @@
                 .handleFirstChildExchange(
                         mReqPayloadListCaptor.capture(),
                         mRespPayloadListCaptor.capture(),
-                        eq(LOCAL_ADDRESS),
-                        eq(REMOTE_ADDRESS),
+                        eq(expectedLocalAddress),
+                        eq(expectedRemoteAddress),
                         any(), // udpEncapSocket
                         eq(mIkeSessionStateMachine.mIkePrf),
                         any()); // sk_d
@@ -5311,6 +5329,34 @@
         assertTrue(mIkeSessionStateMachine.mSupportMobike);
     }
 
+    @Test
+    public void testMobikeEnabledNattSupportedIpv4() throws Exception {
+        verifyMobikeEnabled(true /* doesPeerSupportNatt */, true /* isIpv4 */);
+
+        killSessionAndVerifyNetworkCallback(true /* expectCallbackUnregistered */);
+    }
+
+    @Test
+    public void testMobikeEnabledNattUnsupportedIpv4() throws Exception {
+        verifyMobikeEnabled(false /* doesPeerSupportNatt */, true /* isIpv4 */);
+
+        killSessionAndVerifyNetworkCallback(true /* expectCallbackUnregistered */);
+    }
+
+    @Test
+    public void testMobikeEnabledNattSupportedIpv6() throws Exception {
+        verifyMobikeEnabled(true /* doesPeerSupportNatt */, false /* isIpv4 */);
+
+        killSessionAndVerifyNetworkCallback(true /* expectCallbackUnregistered */);
+    }
+
+    @Test
+    public void testMobikeEnabledNattUnsupportedIpv6() throws Exception {
+        verifyMobikeEnabled(false /* doesPeerSupportNatt */, false /* isIpv4 */);
+
+        killSessionAndVerifyNetworkCallback(true /* expectCallbackUnregistered */);
+    }
+
     /**
      * Restarts the IkeSessionStateMachine with MOBIKE enabled. If doesPeerSupportMobike, MOBIKE
      * will be active for the Session.
@@ -5328,9 +5374,75 @@
     @Nullable
     private IkeNetworkCallbackBase verifyMobikeEnabled(
             boolean doesPeerSupportMobike, Network configuredNetwork) throws Exception {
+        return verifyMobikeEnabled(
+                doesPeerSupportMobike,
+                true /* doesPeerSupportNatt */,
+                true /* isIpv4 */,
+                configuredNetwork);
+    }
+
+    @Nullable
+    private IkeDefaultNetworkCallback verifyMobikeEnabled(
+            boolean doesPeerSupportNatt, boolean isIpv4) throws Exception {
+        // Can cast to IkeDefaultNetworkCallback because no Network is specified
+        return (IkeDefaultNetworkCallback)
+                verifyMobikeEnabled(
+                        true /* doesPeerSupportMobike */,
+                        doesPeerSupportNatt,
+                        isIpv4,
+                        null /* configuredNetwork */);
+    }
+
+    /** Returns the expected IkeSocket type when MOBIKE is supported by both sides */
+    private Class<? extends IkeSocket> getExpectedSocketType(
+            boolean doesPeerSupportNatt, boolean isIpv4) {
+        if (doesPeerSupportNatt) {
+            if (isIpv4) {
+                return IkeUdpEncapSocket.class;
+            } else {
+                return IkeUdp6WithEncapPortSocket.class;
+            }
+        } else {
+            if (isIpv4) {
+                return IkeUdp4Socket.class;
+            } else {
+                return IkeUdp6Socket.class;
+            }
+        }
+    }
+
+    @Nullable
+    private IkeNetworkCallbackBase verifyMobikeEnabled(
+            boolean doesPeerSupportMobike,
+            boolean doesPeerSupportNatt,
+            boolean isIpv4,
+            Network configuredNetwork)
+            throws Exception {
         mIkeSessionStateMachine = restartStateMachineWithMobikeConfigured(configuredNetwork);
         mockIkeInitAndTransitionToIkeAuth(mIkeSessionStateMachine.mCreateIkeLocalIkeAuth);
 
+        // 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 (doesPeerSupportNatt && isIpv4) {
+            // Assume NATs are detected on both sides
+            mIkeSessionStateMachine.mLocalNatDetected = true;
+            mIkeSessionStateMachine.mRemoteNatDetected = true;
+
+            mIkeSessionStateMachine.mIkeSocket = mSpyIkeUdpEncapSocket;
+        } else {
+            mIkeSessionStateMachine.mLocalNatDetected = false;
+            mIkeSessionStateMachine.mRemoteNatDetected = false;
+
+            if (isIpv4) {
+                mIkeSessionStateMachine.mIkeSocket = mSpyIkeUdp4Socket;
+            } else {
+                mIkeSessionStateMachine.mIkeSocket = mSpyIkeUdp6Socket;
+            }
+        }
+
         // Build IKE AUTH response. Include MOBIKE_SUPPORTED if doesPeerSupportMobike is true
         List<IkePayload> authRelatedPayloads = new ArrayList<>();
         IkeAuthPskPayload spyAuthPayload = makeSpyRespPskPayload();
@@ -5352,6 +5464,7 @@
                         true /* hasChildPayloads */,
                         true /* hasConfigPayloadInResp */,
                         doesPeerSupportMobike,
+                        isIpv4,
                         0 /* ike3gppCallbackInvocations */);
         verifyRetransmissionStopped();
 
@@ -5389,6 +5502,9 @@
                             ? IkeDefaultNetworkCallback.class
                             : IkeSpecificNetworkCallback.class;
             assertTrue(expectedCallbackType.isInstance(networkCallback));
+            assertTrue(
+                    getExpectedSocketType(doesPeerSupportNatt, isIpv4)
+                            .isInstance(mIkeSessionStateMachine.mIkeSocket));
         }
         return networkCallback;
     }
@@ -5414,7 +5530,7 @@
         // makeAndStartIkeSession() expects no use of ConnectivityManager#getActiveNetwork when
         // there is a configured Network. Use reset() to forget usage in setUp()
         if (configuredNetwork != null) {
-            reset(mMockConnectManager);
+            resetMockConnectManager();
         }
 
         setupChildStateMachineFactory(mMockChildSessionStateMachine);
@@ -5466,7 +5582,7 @@
     public void testMobikeActiveMobilityEvent() throws Exception {
         IkeDefaultNetworkCallback callback = verifyMobikeEnabled(true /* doesPeerSupportMobike */);
 
-        Network newNetwork = mockNewNetworkAndAddress();
+        Network newNetwork = mockNewNetworkAndAddress(true /* isIpv4 */);
 
         callback.onAvailable(newNetwork);
         mLooper.dispatchAll();
@@ -5474,20 +5590,19 @@
         verifyNetworkAndLocalAddress(newNetwork, UPDATED_LOCAL_ADDRESS, REMOTE_ADDRESS, callback);
         verify(mMockIkeLocalAddressGenerator)
                 .generateLocalAddress(
-                        eq(newNetwork),
-                        eq(true /* isIpv4 */),
-                        eq(REMOTE_ADDRESS),
-                        eq(IkeSocket.SERVER_PORT_NON_UDP_ENCAPSULATED));
+                        eq(newNetwork), eq(true /* isIpv4 */), eq(REMOTE_ADDRESS), anyInt());
     }
 
-    private Network mockNewNetworkAndAddress() throws Exception {
+    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;
         when(mMockIkeLocalAddressGenerator.generateLocalAddress(
-                        eq(newNetwork),
-                        eq(true /* isIpv4 */),
-                        eq(REMOTE_ADDRESS),
-                        eq(IkeSocket.SERVER_PORT_NON_UDP_ENCAPSULATED)))
-                .thenReturn(UPDATED_LOCAL_ADDRESS);
+                        eq(newNetwork), eq(isIpv4), eq(expectedRemoteAddress), anyInt()))
+                .thenReturn(injectedLocalAddress);
+
         return newNetwork;
     }
 
@@ -5540,18 +5655,29 @@
     private void verifySetNetwork(
             IkeNetworkCallbackBase callback, IkeSaRecord rekeySaRecord, State expectedState)
             throws Exception {
-        Network newNetwork = mockNewNetworkAndAddress();
+        verifySetNetwork(callback, rekeySaRecord, expectedState, true /* isIpv4 */);
+    }
+
+    private void verifySetNetwork(
+            IkeNetworkCallbackBase callback,
+            IkeSaRecord rekeySaRecord,
+            State expectedState,
+            boolean isIpv4)
+            throws Exception {
+        Network newNetwork = mockNewNetworkAndAddress(isIpv4);
 
         mIkeSessionStateMachine.setNetwork(newNetwork);
         mLooper.dispatchAll();
 
-        verifyNetworkAndLocalAddress(newNetwork, UPDATED_LOCAL_ADDRESS, REMOTE_ADDRESS, callback);
+        InetAddress expectedUpdatedLocalAddress =
+                isIpv4 ? UPDATED_LOCAL_ADDRESS : UPDATED_LOCAL_ADDRESS_V6;
+        InetAddress expectedRemoteAddress = isIpv4 ? REMOTE_ADDRESS : REMOTE_ADDRESS_V6;
+
+        verifyNetworkAndLocalAddress(
+                newNetwork, expectedUpdatedLocalAddress, expectedRemoteAddress, callback);
         verify(mMockIkeLocalAddressGenerator)
                 .generateLocalAddress(
-                        eq(newNetwork),
-                        eq(true /* isIpv4 */),
-                        eq(REMOTE_ADDRESS),
-                        eq(IkeSocket.SERVER_PORT_NON_UDP_ENCAPSULATED));
+                        eq(newNetwork), eq(isIpv4), eq(expectedRemoteAddress), anyInt());
 
         assertEquals(
                 mIkeSessionStateMachine,
@@ -5559,7 +5685,7 @@
                         mIkeSessionStateMachine.mCurrentIkeSaRecord.getLocalSpi()));
 
         if (rekeySaRecord != null) {
-            verifyIkeSaAddresses(rekeySaRecord, UPDATED_LOCAL_ADDRESS, REMOTE_ADDRESS);
+            verifyIkeSaAddresses(rekeySaRecord, expectedUpdatedLocalAddress, expectedRemoteAddress);
             assertEquals(
                     mIkeSessionStateMachine,
                     mIkeSessionStateMachine.mIkeSocket.mSpiToIkeSession.get(
@@ -5570,8 +5696,17 @@
     }
 
     private IkeNetworkCallbackBase setupIdleStateMachineWithMobike() throws Exception {
+        return setupIdleStateMachineWithMobike(true /* doesPeerSupportNatt */, true /* isIpv4 */);
+    }
+
+    private IkeNetworkCallbackBase setupIdleStateMachineWithMobike(
+            boolean doesPeerSupportNatt, boolean isIpv4) throws Exception {
         IkeNetworkCallbackBase callback =
-                verifyMobikeEnabled(true /* doesPeerSupportMobike */, mMockDefaultNetwork);
+                verifyMobikeEnabled(
+                        true /* doesPeerSupportMobike */,
+                        doesPeerSupportNatt,
+                        isIpv4,
+                        mMockDefaultNetwork);
 
         // reset IkeMessageHelper to make verifying outbound req easier
         resetMockIkeMessageHelper();
@@ -5587,12 +5722,39 @@
         return callback;
     }
 
-    @Test
-    public void testSetNetworkIdleState() throws Exception {
-        IkeNetworkCallbackBase callback = setupIdleStateMachineWithMobike();
+    private void verifySetNetworkIdleState(boolean doesPeerSupportNatt, boolean isIpv4)
+            throws Exception {
+        IkeNetworkCallbackBase callback =
+                setupIdleStateMachineWithMobike(doesPeerSupportNatt, isIpv4);
 
         verifySetNetwork(
-                callback, null /* rekeySaRecord */, mIkeSessionStateMachine.mMobikeLocalInfo);
+                callback,
+                null /* rekeySaRecord */,
+                mIkeSessionStateMachine.mMobikeLocalInfo,
+                isIpv4);
+        assertTrue(
+                getExpectedSocketType(doesPeerSupportNatt, isIpv4)
+                        .isInstance(mIkeSessionStateMachine.mIkeSocket));
+    }
+
+    @Test
+    public void testSetNetworkIdleStateNattSupportedIpv4() throws Exception {
+        verifySetNetworkIdleState(true /* doesPeerSupportNatt */, true /* isIpv4 */);
+    }
+
+    @Test
+    public void testSetNetworkIdleStateNattSupportedIpv6() throws Exception {
+        verifySetNetworkIdleState(true /* doesPeerSupportNatt */, false /* isIpv4 */);
+    }
+
+    @Test
+    public void testSetNetworkIdleStateNattUnsupportedIpv4() throws Exception {
+        verifySetNetworkIdleState(false /* doesPeerSupportNatt */, true /* isIpv4 */);
+    }
+
+    @Test
+    public void testSetNetworkIdleStateNattUnsupportedIpv6() throws Exception {
+        verifySetNetworkIdleState(false /* doesPeerSupportNatt */, false /* isIpv4 */);
     }
 
     @Test
@@ -5829,8 +5991,6 @@
                 mIkeSessionStateMachine.getCurrentState()
                         instanceof IkeSessionStateMachine.ChildProcedureOngoing);
         verify(mMockChildSessionStateMachine).rekeyChildSessionForMobike();
-
-        // TODO(b/173237734): check IkeSocket - if includeNatDetection then expect UdpEncap
     }
 
     @Test
@@ -5840,8 +6000,11 @@
         verifySetNetwork(
                 callback, null /* rekeySaRecord */, mIkeSessionStateMachine.mMobikeLocalInfo);
 
+        // Keepalive for the old UDP encap socket stopped
         verify(mMockIkeNattKeepalive).stop();
-        assertNull(mIkeSessionStateMachine.mIkeNattKeepalive);
+
+        // Keepalive for the new UDP encap socket started
+        assertNotNull(mIkeSessionStateMachine.mIkeNattKeepalive);
     }
 
     @Test
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 0daeb9d..5d14a04 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
@@ -27,6 +27,7 @@
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.eq;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.reset;
 import static org.mockito.Mockito.spy;
 
 import android.content.Context;
@@ -48,6 +49,7 @@
 import org.junit.Before;
 
 import java.net.Inet4Address;
+import java.net.Inet6Address;
 import java.util.concurrent.Executor;
 
 public abstract class IkeSessionTestBase {
@@ -57,6 +59,12 @@
             (Inet4Address) InetAddresses.parseNumericAddress("192.0.2.201");
     protected static final Inet4Address REMOTE_ADDRESS =
             (Inet4Address) InetAddresses.parseNumericAddress("127.0.0.1");
+    protected static final Inet6Address LOCAL_ADDRESS_V6 =
+            (Inet6Address) InetAddresses.parseNumericAddress("2001:db8::200");
+    protected static final Inet6Address UPDATED_LOCAL_ADDRESS_V6 =
+            (Inet6Address) InetAddresses.parseNumericAddress("2001:db8::201");
+    protected static final Inet6Address REMOTE_ADDRESS_V6 =
+            (Inet6Address) InetAddresses.parseNumericAddress("::1");
     protected static final String REMOTE_HOSTNAME = "ike.test.android.com";
 
     protected PowerManager.WakeLock mMockBusyWakelock;
@@ -100,15 +108,29 @@
                 .when(mPowerManager)
                 .newWakeLock(anyInt(), argThat(tag -> tag.contains(LOCAL_REQUEST_WAKE_LOCK_TAG)));
 
-        mMockConnectManager = mock(ConnectivityManager.class);
         mMockDefaultNetwork = mock(Network.class);
-        doReturn(mMockDefaultNetwork).when(mMockConnectManager).getActiveNetwork();
         doReturn(REMOTE_ADDRESS).when(mMockDefaultNetwork).getByName(REMOTE_HOSTNAME);
         doReturn(REMOTE_ADDRESS)
                 .when(mMockDefaultNetwork)
                 .getByName(REMOTE_ADDRESS.getHostAddress());
 
         mMockSocketKeepalive = mock(SocketKeepalive.class);
+
+        mMockNetworkCapabilities = mock(NetworkCapabilities.class);
+        doReturn(false)
+                .when(mMockNetworkCapabilities)
+                .hasTransport(RandomnessFactory.TRANSPORT_TEST);
+
+        mMockConnectManager = mock(ConnectivityManager.class);
+        doReturn(mMockConnectManager)
+                .when(mSpyContext)
+                .getSystemService(Context.CONNECTIVITY_SERVICE);
+        resetMockConnectManager();
+    }
+
+    protected void resetMockConnectManager() {
+        reset(mMockConnectManager);
+        doReturn(mMockDefaultNetwork).when(mMockConnectManager).getActiveNetwork();
         doReturn(mMockSocketKeepalive)
                 .when(mMockConnectManager)
                 .createSocketKeepalive(
@@ -118,16 +140,8 @@
                         any(Inet4Address.class),
                         any(Executor.class),
                         any(SocketKeepalive.Callback.class));
-        doReturn(mMockConnectManager)
-                .when(mSpyContext)
-                .getSystemService(Context.CONNECTIVITY_SERVICE);
-
-        mMockNetworkCapabilities = mock(NetworkCapabilities.class);
         doReturn(mMockNetworkCapabilities)
                 .when(mMockConnectManager)
                 .getNetworkCapabilities(any(Network.class));
-        doReturn(false)
-                .when(mMockNetworkCapabilities)
-                .hasTransport(RandomnessFactory.TRANSPORT_TEST);
     }
 }