Snap for 5430508 from 87976f13ea5b6e54d62c96e010658d10d8cf7812 to qt-release

Change-Id: If822aed02176aa465c9959d97fd6ae025df7e308
diff --git a/src/java/com/android/ike/ikev2/ChildSessionOptions.java b/src/java/com/android/ike/ikev2/ChildSessionOptions.java
index cbca9f4..b311c66 100644
--- a/src/java/com/android/ike/ikev2/ChildSessionOptions.java
+++ b/src/java/com/android/ike/ikev2/ChildSessionOptions.java
@@ -20,6 +20,6 @@
  * ChildSessionOptions contains user-provided Child SA proposals and negotiated Child SA
  * information.
  */
-public class ChildSessionOptions {
+public final class ChildSessionOptions {
     // TODO: Implement it.
 }
diff --git a/src/java/com/android/ike/ikev2/IkeSessionOptions.java b/src/java/com/android/ike/ikev2/IkeSessionOptions.java
index fe06b85..9f2094f 100644
--- a/src/java/com/android/ike/ikev2/IkeSessionOptions.java
+++ b/src/java/com/android/ike/ikev2/IkeSessionOptions.java
@@ -29,7 +29,7 @@
  *
  * <p>TODO: Make this doc more user-friendly.
  */
-public class IkeSessionOptions {
+public final class IkeSessionOptions {
     private final InetAddress mServerAddress;
     private final UdpEncapsulationSocket mUdpEncapSocket;
     private final SaProposal[] mSaProposals;
@@ -92,7 +92,7 @@
          * @throws IllegalArgumentException if input proposal is not IKE SA proposal.
          */
         public Builder addSaProposal(SaProposal proposal) {
-            if (proposal.mProtocolId != IkePayload.PROTOCOL_ID_IKE) {
+            if (proposal.getProtocolId() != IkePayload.PROTOCOL_ID_IKE) {
                 throw new IllegalArgumentException(
                         "Expected IKE SA Proposal but received Child SA proposal");
             }
diff --git a/src/java/com/android/ike/ikev2/IkeSessionStateMachine.java b/src/java/com/android/ike/ikev2/IkeSessionStateMachine.java
index c887ecb..b4a7b9c 100644
--- a/src/java/com/android/ike/ikev2/IkeSessionStateMachine.java
+++ b/src/java/com/android/ike/ikev2/IkeSessionStateMachine.java
@@ -17,20 +17,30 @@
 
 import android.os.Looper;
 import android.os.Message;
+import android.system.ErrnoException;
 import android.util.LongSparseArray;
-import android.util.Pair;
 import android.util.SparseArray;
 
 import com.android.ike.ikev2.SaRecord.IkeSaRecord;
 import com.android.ike.ikev2.exceptions.IkeException;
 import com.android.ike.ikev2.message.IkeHeader;
+import com.android.ike.ikev2.message.IkeKePayload;
 import com.android.ike.ikev2.message.IkeMessage;
+import com.android.ike.ikev2.message.IkeNoncePayload;
 import com.android.ike.ikev2.message.IkeNotifyPayload;
+import com.android.ike.ikev2.message.IkePayload;
+import com.android.ike.ikev2.message.IkeSaPayload;
+import com.android.ike.ikev2.message.IkeSaPayload.DhGroupTransform;
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.util.State;
 import com.android.internal.util.StateMachine;
 
 import java.security.GeneralSecurityException;
+import java.security.SecureRandom;
+import java.util.HashSet;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.Set;
 
 /**
  * IkeSessionStateMachine tracks states and manages exchanges of this IKE session.
@@ -74,6 +84,11 @@
     static final int CMD_LOCAL_REQUEST_REKEY_CHILD = CMD_LOCAL_REQUEST_BASE + 7;
     // TODO: Add signals for other procedure types and notificaitons.
 
+    // Remember locally assigned IKE SPIs to avoid SPI collision.
+    private static final Set<Long> ASSIGNED_LOCAL_IKE_SPI_SET = new HashSet<>();
+    private static final int MAX_ASSIGN_IKE_SPI_ATTEMPTS = 100;
+    private static final SecureRandom IKE_SPI_RANDOM = new SecureRandom();
+
     private final IkeSessionOptions mIkeSessionOptions;
     private final ChildSessionOptions mFirstChildSessionOptions;
     /** Map that stores all IkeSaRecords, keyed by remotely generated IKE SPI. */
@@ -85,6 +100,12 @@
      */
     private final SparseArray<ChildSessionStateMachine> mSpiToChildSessionMap;
 
+    /**
+     * Package private socket that sends and receives encoded IKE message. Initialized in Initial
+     * State.
+     */
+    @VisibleForTesting IkeSocket mIkeSocket;
+
     /** Package */
     @VisibleForTesting IkeSaRecord mCurrentIkeSaRecord;
     /** Package */
@@ -146,9 +167,55 @@
         setInitialState(mInitial);
     }
 
+    // Generate IKE SPI. Throw an exception if it failed and handle this exception in current State.
+    private static Long getIkeSpiOrThrow() {
+        for (int i = 0; i < MAX_ASSIGN_IKE_SPI_ATTEMPTS; i++) {
+            long spi = IKE_SPI_RANDOM.nextLong();
+            if (ASSIGNED_LOCAL_IKE_SPI_SET.add(spi)) return spi;
+        }
+        throw new IllegalStateException("Failed to generate IKE SPI.");
+    }
+
     private IkeMessage buildIkeInitReq() {
-        // TODO:Build packet according to mIkeSessionOptions.
-        return null;
+        // TODO: Handle IKE SPI assigning error in CreateIkeLocalIkeInit State.
+
+        List<IkePayload> payloadList = new LinkedList<>();
+
+        // Generate IKE SPI
+        long initSpi = getIkeSpiOrThrow();
+        long respSpi = 0;
+
+        // It is validated in IkeSessionOptions.Builder to ensure IkeSessionOptions has at least one
+        // SaProposal and all SaProposals are valid for IKE SA negotiation.
+        SaProposal[] saProposals = mIkeSessionOptions.getSaProposals();
+
+        // Build SA Payload
+        IkeSaPayload saPayload = new IkeSaPayload(saProposals);
+        payloadList.add(saPayload);
+
+        // Build KE Payload using the first DH group number in the first SaProposal.
+        DhGroupTransform dhGroupTransform = saProposals[0].getDhGroupTransforms()[0];
+        IkeKePayload kePayload = new IkeKePayload(dhGroupTransform.id);
+        payloadList.add(kePayload);
+
+        // Build Nonce Payload
+        IkeNoncePayload noncePayload = new IkeNoncePayload();
+        payloadList.add(noncePayload);
+
+        // TODO: Add Notification Payloads according to user configurations.
+
+        // Build IKE header
+        IkeHeader ikeHeader =
+                new IkeHeader(
+                        initSpi,
+                        respSpi,
+                        IkePayload.PAYLOAD_TYPE_SA,
+                        IkeHeader.EXCHANGE_TYPE_IKE_SA_INIT,
+                        false /*isResponseMsg*/,
+                        true /*fromIkeInitiator*/,
+                        0 /*messageId*/);
+
+        return new IkeMessage(ikeHeader, payloadList);
     }
 
     private IkeMessage buildIkeAuthReq() {
@@ -214,6 +281,19 @@
     }
 
     /**
+     * Receive IKE packet from remote server.
+     *
+     * <p>This method is called synchronously from IkeSocket. It proxies the synchronous call as an
+     * asynchronous job to the IkeSessionStateMachine handler.
+     *
+     * @param ikeHeader the decoded IKE header.
+     * @param ikePacketBytes the byte array of the entire received IKE packet.
+     */
+    public void receiveIkePacket(IkeHeader ikeHeader, byte[] ikePacketBytes) {
+        sendMessage(CMD_RECEIVE_IKE_PACKET, new ReceivedIkePacket(ikeHeader, ikePacketBytes));
+    }
+
+    /**
      * ReceivedIkePacket is a package private data container consists of decoded IkeHeader and
      * encoded IKE packet in a byte array.
      */
@@ -223,9 +303,9 @@
         /** Entire encoded IKE message including IKE header */
         public final byte[] ikePacketBytes;
 
-        ReceivedIkePacket(Pair<IkeHeader, byte[]> ikePacketPair) {
-            ikeHeader = ikePacketPair.first;
-            ikePacketBytes = ikePacketPair.second;
+        ReceivedIkePacket(IkeHeader ikeHeader, byte[] ikePacketBytes) {
+            this.ikeHeader = ikeHeader;
+            this.ikePacketBytes = ikePacketBytes;
         }
     }
 
@@ -260,6 +340,15 @@
     /** Initial state of IkeSessionStateMachine. */
     class Initial extends State {
         @Override
+        public void enter() {
+            try {
+                mIkeSocket = IkeSocket.getIkeSocket(mIkeSessionOptions.getUdpEncapsulationSocket());
+            } catch (ErrnoException e) {
+                // TODO: handle exception and close IkeSession.
+            }
+        }
+
+        @Override
         public boolean processMessage(Message message) {
             switch (message.what) {
                 case CMD_LOCAL_REQUEST_CREATE_IKE:
@@ -417,6 +506,7 @@
         public void enter() {
             mRequestMsg = buildRequest();
             mRequestPacket = encodeRequest();
+            mIkeSocket.sendIkePacket(mRequestPacket, mIkeSessionOptions.getServerAddress());
             // TODO: Send out packet and start retransmission timer.
         }
 
@@ -432,12 +522,20 @@
         // CreateIkeLocalInit should override encodeRequest() to encode unencrypted packet
         protected byte[] encodeRequest() {
             // TODO: encrypt and encode mRequestMsg
-            return null;
-        };
+            return new byte[0];
+        }
     }
 
     /** CreateIkeLocalIkeInit represents state when IKE library initiates IKE_INIT exchange. */
     class CreateIkeLocalIkeInit extends LocalNewExchangeBase {
+
+        @Override
+        public void enter() {
+            super.enter();
+            mIkeSocket.registerIke(
+                    mRequestMsg.ikeHeader.ikeInitiatorSpi, IkeSessionStateMachine.this);
+        }
+
         @Override
         protected IkeMessage buildRequest() {
             return buildIkeInitReq();
@@ -445,8 +543,7 @@
 
         @Override
         protected byte[] encodeRequest() {
-            // TODO: Encode an unencrypted IKE packet.
-            return null;
+            return mRequestMsg.encode();
         }
 
         @Override
@@ -522,8 +619,7 @@
                                         mFirstChildSessionOptions);
                         // TODO: Replace null input params to payload lists in IKE_AUTH request and
                         // IKE_AUTH response for negotiating Child SA.
-                        firstChild.handleFirstChildExchange(
-                                null, null, new ChildSessionCallback());
+                        firstChild.handleFirstChildExchange(null, null, new ChildSessionCallback());
 
                         transitionTo(mIdle);
                     } catch (IkeException e) {
diff --git a/src/java/com/android/ike/ikev2/IkeSocket.java b/src/java/com/android/ike/ikev2/IkeSocket.java
index 71daddd..e235ab8 100644
--- a/src/java/com/android/ike/ikev2/IkeSocket.java
+++ b/src/java/com/android/ike/ikev2/IkeSocket.java
@@ -25,8 +25,11 @@
 import android.os.Handler;
 import android.system.ErrnoException;
 import android.system.Os;
+import android.util.Log;
 import android.util.LongSparseArray;
 
+import com.android.ike.ikev2.exceptions.IkeException;
+import com.android.ike.ikev2.message.IkeHeader;
 import com.android.internal.annotations.VisibleForTesting;
 
 import java.io.FileDescriptor;
@@ -60,24 +63,30 @@
  * thread, there will not be concurrent modification problems.
  */
 public final class IkeSocket extends PacketReader implements AutoCloseable {
+    private static final String TAG = "IkeSocket";
+
     // TODO: b/129358324 Consider supporting IKE exchange without UDP Encapsulation.
     // UDP-encapsulated IKE packets MUST be sent to 4500.
     @VisibleForTesting static final int IKE_SERVER_PORT = 4500;
+
     // A Non-ESP marker helps the recipient to distinguish IKE packets from ESP packets.
     @VisibleForTesting static final int NON_ESP_MARKER_LEN = 4;
+    @VisibleForTesting static final byte[] NON_ESP_MARKER = new byte[NON_ESP_MARKER_LEN];
 
-    // Package private map from UdpEncapsulationSocket to IkeSocket instances.
-    static Map<UdpEncapsulationSocket, IkeSocket> sFdToIkeSocketMap = new HashMap<>();
+    // Map from UdpEncapsulationSocket to IkeSocket instances.
+    private static Map<UdpEncapsulationSocket, IkeSocket> sFdToIkeSocketMap = new HashMap<>();
 
     private static IPacketReceiver sPacketReceiver = new PacketReceiver();
 
-    // Map from locally generated IKE SPI to IkeSessionStateMachine instances.
-    private final LongSparseArray<IkeSessionStateMachine> mSpiToIkeSession =
+    // Package private map from locally generated IKE SPI to IkeSessionStateMachine instances.
+    @VisibleForTesting
+    final LongSparseArray<IkeSessionStateMachine> mSpiToIkeSession =
             new LongSparseArray<>();
     // UdpEncapsulationSocket for sending and receving IKE packet.
     private final UdpEncapsulationSocket mUdpEncapSocket;
 
     /** Package private */
+    @VisibleForTesting
     int mRefCount;
 
     private IkeSocket(UdpEncapsulationSocket udpEncapSocket, Handler handler) {
@@ -138,10 +147,50 @@
     }
 
     /** Package private */
+    @VisibleForTesting
     static final class PacketReceiver implements IPacketReceiver {
         public void handlePacket(
                 byte[] recvbuf, LongSparseArray<IkeSessionStateMachine> spiToIkeSession) {
-            // TODO: Decode IKE header and demultiplex IKE packet
+            // TODO: b/129708574 Consider only logging the error some % of the time it happens, or
+            // only logging the error the first time it happens and then keep a count to prevent
+            // logspam.
+
+            ByteBuffer byteBuffer = ByteBuffer.wrap(recvbuf);
+
+            // Check the existence of the Non-ESP Marker. A received packet can be either an IKE
+            // packet starts with 4 zero-valued bytes Non-ESP Marker or an ESP packet starts with 4
+            // bytes ESP SPI. ESP SPI value can never be zero.
+            byte[] espMarker = new byte[NON_ESP_MARKER_LEN];
+            byteBuffer.get(espMarker);
+            if (!Arrays.equals(NON_ESP_MARKER, espMarker)) {
+                // Drop the received ESP packet.
+                Log.e(TAG, "Receive an ESP packet.");
+                return;
+            }
+
+            try {
+                // Re-direct IKE packet to IkeSessionStateMachine according to the locally generated
+                // IKE SPI.
+                byte[] ikePacketBytes = new byte[byteBuffer.remaining()];
+                byteBuffer.get(ikePacketBytes);
+                IkeHeader ikeHeader = new IkeHeader(ikePacketBytes);
+
+                long localGeneratedSpi =
+                        ikeHeader.fromIkeInitiator
+                                ? ikeHeader.ikeResponderSpi
+                                : ikeHeader.ikeInitiatorSpi;
+
+                IkeSessionStateMachine ikeStateMachine = spiToIkeSession.get(localGeneratedSpi);
+                if (ikeStateMachine == null) {
+                    Log.e(TAG, "Unrecognized IKE SPI.");
+                    // TODO: Handle invalid IKE SPI error
+                } else {
+                    ikeStateMachine.receiveIkePacket(ikeHeader, ikePacketBytes);
+                }
+            } catch (IkeException e) {
+                // Handle invalid IKE header
+                Log.e(TAG, "Can't parse malformed IKE packet header.");
+            }
         }
     }
 
@@ -171,7 +220,7 @@
             ByteBuffer buffer = ByteBuffer.allocate(NON_ESP_MARKER_LEN + ikePacket.length);
 
             // Build outbound UDP Encapsulation packet body for sending IKE message.
-            buffer.put(new byte[NON_ESP_MARKER_LEN]).put(ikePacket);
+            buffer.put(NON_ESP_MARKER).put(ikePacket);
             buffer.rewind();
 
             // Use unconnected UDP socket because one {@UdpEncapsulationSocket} may be shared by
diff --git a/src/java/com/android/ike/ikev2/SaProposal.java b/src/java/com/android/ike/ikev2/SaProposal.java
index 04c06d0..53206cc 100644
--- a/src/java/com/android/ike/ikev2/SaProposal.java
+++ b/src/java/com/android/ike/ikev2/SaProposal.java
@@ -139,17 +139,17 @@
     }
 
     /** Package private */
-    @IkePayload.ProtocolId final int mProtocolId;
+    @IkePayload.ProtocolId private final int mProtocolId;
     /** Package private */
-    final EncryptionTransform[] mEncryptionAlgorithms;
+    private final EncryptionTransform[] mEncryptionAlgorithms;
     /** Package private */
-    final PrfTransform[] mPseudorandomFunctions;
+    private final PrfTransform[] mPseudorandomFunctions;
     /** Package private */
-    final IntegrityTransform[] mIntegrityAlgorithms;
+    private final IntegrityTransform[] mIntegrityAlgorithms;
     /** Package private */
-    final DhGroupTransform[] mDhGroups;
+    private final DhGroupTransform[] mDhGroups;
     /** Package private */
-    final EsnTransform[] mEsns;
+    private final EsnTransform[] mEsns;
 
     private SaProposal(
             @IkePayload.ProtocolId int protocol,
@@ -227,6 +227,37 @@
         return Arrays.asList(selectFrom).contains(selected[0]);
     }
 
+    /*Package private*/
+    @IkePayload.ProtocolId
+    int getProtocolId() {
+        return mProtocolId;
+    }
+
+    /*Package private*/
+    EncryptionTransform[] getEncryptionTransforms() {
+        return mEncryptionAlgorithms;
+    }
+
+    /*Package private*/
+    PrfTransform[] getPrfTransforms() {
+        return mPseudorandomFunctions;
+    }
+
+    /*Package private*/
+    IntegrityTransform[] getIntegrityTransforms() {
+        return mIntegrityAlgorithms;
+    }
+
+    /*Package private*/
+    DhGroupTransform[] getDhGroupTransforms() {
+        return mDhGroups;
+    }
+
+    /*Package private*/
+    EsnTransform[] getEsnTransforms() {
+        return mEsns;
+    }
+
     /**
      * Return all SA Transforms in this SaProposal to be encoded for building an outbound IKE
      * message.
diff --git a/src/java/com/android/ike/ikev2/SaRecord.java b/src/java/com/android/ike/ikev2/SaRecord.java
index 355d41d..4dcddec 100644
--- a/src/java/com/android/ike/ikev2/SaRecord.java
+++ b/src/java/com/android/ike/ikev2/SaRecord.java
@@ -17,9 +17,16 @@
 
 import com.android.ike.ikev2.message.IkeMessage;
 import com.android.ike.ikev2.message.IkePayload;
+import com.android.internal.annotations.VisibleForTesting;
 
+import java.nio.ByteBuffer;
+import java.security.InvalidKeyException;
+import java.security.NoSuchAlgorithmException;
 import java.util.List;
 
+import javax.crypto.Mac;
+import javax.crypto.spec.SecretKeySpec;
+
 /**
  * SaRecord represents common information of an IKE SA and a Child SA.
  *
@@ -199,4 +206,71 @@
         ChildSaRecord makeChildSaRecord(
                 List<IkePayload> reqPayloads, List<IkePayload> respPayloads);
     }
+
+    /** Generate SKEYSEED using negotiated PRF. */
+    @VisibleForTesting
+    static byte[] generateSKeySeed(
+            String prfAlgorithm, byte[] nonceInit, byte[] nonceResp, byte[] sharedDhKey) {
+        try {
+            ByteBuffer keyBuffer = ByteBuffer.allocate(nonceInit.length + nonceResp.length);
+            keyBuffer.put(nonceInit).put(nonceResp);
+            SecretKeySpec prfKeySpec = new SecretKeySpec(keyBuffer.array(), prfAlgorithm);
+
+            Mac prfMac = Mac.getInstance(prfAlgorithm, IkeMessage.getSecurityProvider());
+            prfMac.init(prfKeySpec);
+
+            ByteBuffer sharedKeyBuffer = ByteBuffer.wrap(sharedDhKey);
+            prfMac.update(sharedKeyBuffer);
+
+            return prfMac.doFinal();
+        } catch (InvalidKeyException | NoSuchAlgorithmException e) {
+            throw new IllegalArgumentException("Failed to generate SKEYSEED", e);
+        }
+    }
+
+    /**
+     * Derives key materials using negotiated PRF.
+     *
+     * <p>prf+(K, S) outputs a pseudorandom stream by using negotiated PRF iteratively. In this way
+     * it can generate long enough keying material containing all the keys for this IKE/Child SA.
+     *
+     * @see <a href="https://tools.ietf.org/html/rfc7296#section-2.13">RFC 7296 nternet Key Exchange
+     *     Protocol Version 2 (IKEv2) 2.13. Generating Keying Material </a>
+     */
+    @VisibleForTesting
+    static byte[] generateKeyMat(
+            String prfAlgorithm, byte[] prfKey, byte[] dataToSign, int keyMaterialLen)
+            throws InvalidKeyException {
+        try {
+            SecretKeySpec prfKeySpec = new SecretKeySpec(prfKey, prfAlgorithm);
+            Mac prfMac = Mac.getInstance(prfAlgorithm, IkeMessage.getSecurityProvider());
+
+            ByteBuffer keyMatBuffer = ByteBuffer.allocate(keyMaterialLen);
+
+            byte[] previousMac = new byte[0];
+            final int padLen = 1;
+            byte padValue = 1;
+
+            while (keyMatBuffer.remaining() > 0) {
+                prfMac.init(prfKeySpec);
+
+                ByteBuffer dataToSignBuffer =
+                        ByteBuffer.allocate(previousMac.length + dataToSign.length + padLen);
+                dataToSignBuffer.put(previousMac).put(dataToSign).put(padValue);
+                dataToSignBuffer.rewind();
+
+                prfMac.update(dataToSignBuffer);
+
+                previousMac = prfMac.doFinal();
+                keyMatBuffer.put(
+                        previousMac, 0, Math.min(previousMac.length, keyMatBuffer.remaining()));
+
+                padValue++;
+            }
+
+            return keyMatBuffer.array();
+        } catch (InvalidKeyException | NoSuchAlgorithmException e) {
+            throw new IllegalArgumentException("Failed to generate keying material", e);
+        }
+    }
 }
diff --git a/src/java/com/android/ike/ikev2/message/IkeEncryptedPayloadBody.java b/src/java/com/android/ike/ikev2/message/IkeEncryptedPayloadBody.java
index f47fb88..a41d855 100644
--- a/src/java/com/android/ike/ikev2/message/IkeEncryptedPayloadBody.java
+++ b/src/java/com/android/ike/ikev2/message/IkeEncryptedPayloadBody.java
@@ -139,7 +139,12 @@
         ByteBuffer authenticatedSectionBuffer = ByteBuffer.allocate(dataToAuthenticateLength);
 
         // Encode IKE header
-        ikeHeader.encodeToByteBuffer(authenticatedSectionBuffer);
+        int encryptedPayloadLength =
+                IkePayload.GENERIC_HEADER_LENGTH
+                        + iv.length
+                        + mEncryptedAndPaddedData.length
+                        + checksumLen;
+        ikeHeader.encodeToByteBuffer(authenticatedSectionBuffer, encryptedPayloadLength);
 
         // Encode payload header. The next payload type field indicates the first payload nested in
         // this SkPayload/SkfPayload.
diff --git a/src/java/com/android/ike/ikev2/message/IkeHeader.java b/src/java/com/android/ike/ikev2/message/IkeHeader.java
index ebd3553..c4f215c 100644
--- a/src/java/com/android/ike/ikev2/message/IkeHeader.java
+++ b/src/java/com/android/ike/ikev2/message/IkeHeader.java
@@ -23,6 +23,7 @@
 import com.android.ike.ikev2.exceptions.IkeException;
 import com.android.ike.ikev2.exceptions.InvalidMajorVersionException;
 import com.android.ike.ikev2.exceptions.InvalidSyntaxException;
+import com.android.internal.annotations.VisibleForTesting;
 
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
@@ -36,7 +37,7 @@
  *     Protocol Version 2 (IKEv2)</a>
  */
 public final class IkeHeader {
-    //TODO: b/122838549 Change IkeHeader to static inner class of IkeMessage.
+    // TODO: b/122838549 Change IkeHeader to static inner class of IkeMessage.
     private static final byte IKE_HEADER_VERSION_INFO = (byte) 0x20;
 
     // Indicate whether this message is a response message
@@ -69,7 +70,14 @@
     public final boolean isResponseMsg;
     public final boolean fromIkeInitiator;
     public final int messageId;
-    public final int messageLength;
+
+    // Cannot assign encoded message length value for an outbound IKE message before it's encoded.
+    private static final int ENCODED_MESSAGE_LEN_UNAVAILABLE = -1;
+
+    // mEncodedMessageLength is only set for an inbound IkeMessage. When building an outbound
+    // IkeMessage, message length is not set because message body length is unknown until it gets
+    // encrypted and encoded.
+    private final int mEncodedMessageLength;
 
     /**
      * Construct an instance of IkeHeader. It is only called in the process of building outbound
@@ -82,7 +90,6 @@
      * @param isResp indicates if this message is a response or a request
      * @param fromInit indictaes if this message is sent from the IKE initiator or the IKE responder
      * @param msgId the message identifier
-     * @param length the length of the total message in octets
      */
     public IkeHeader(
             long iSpi,
@@ -91,8 +98,7 @@
             @ExchangeType int eType,
             boolean isResp,
             boolean fromInit,
-            int msgId,
-            int length) {
+            int msgId) {
         ikeInitiatorSpi = iSpi;
         ikeResponderSpi = rSpi;
         nextPayloadType = nextPType;
@@ -100,7 +106,8 @@
         isResponseMsg = isResp;
         fromIkeInitiator = fromInit;
         messageId = msgId;
-        messageLength = length;
+
+        mEncodedMessageLength = ENCODED_MESSAGE_LEN_UNAVAILABLE;
 
         // Major version of IKE protocol in use; it must be set to 2 when building an IKEv2 message.
         majorVersion = 2;
@@ -135,11 +142,21 @@
         fromIkeInitiator = ((flagsByte & 0x08) != 0);
 
         messageId = buffer.getInt();
-        messageLength = buffer.getInt();
+        mEncodedMessageLength = buffer.getInt();
     }
 
-    /** Validate syntax and major version. */
-    public void checkValidOrThrow(int packetLength) throws IkeException {
+    /*Package private*/
+    @VisibleForTesting
+    int getInboundMessageLength() {
+        if (mEncodedMessageLength == ENCODED_MESSAGE_LEN_UNAVAILABLE) {
+            throw new UnsupportedOperationException(
+                    "It is not supported to get encoded message length from an outbound message.");
+        }
+        return mEncodedMessageLength;
+    }
+
+    /** Validate syntax and major version of inbound IKE header. */
+    public void checkInboundValidOrThrow(int packetLength) throws IkeException {
         if (majorVersion > 2) {
             // Receive higher version of protocol. Stop parsing.
             throw new InvalidMajorVersionException(majorVersion);
@@ -155,13 +172,13 @@
                 || exchangeType > EXCHANGE_TYPE_INFORMATIONAL) {
             throw new InvalidSyntaxException("Invalid IKE Exchange Type.");
         }
-        if (messageLength != packetLength) {
+        if (mEncodedMessageLength != packetLength) {
             throw new InvalidSyntaxException("Invalid IKE Message Length.");
         }
     }
 
     /** Encode IKE header to ByteBuffer */
-    public void encodeToByteBuffer(ByteBuffer byteBuffer) {
+    public void encodeToByteBuffer(ByteBuffer byteBuffer, int encodedMessageBodyLen) {
         byteBuffer
                 .putLong(ikeInitiatorSpi)
                 .putLong(ikeResponderSpi)
@@ -177,6 +194,6 @@
             flag |= IKE_HEADER_FLAG_FROM_IKE_INITIATOR;
         }
 
-        byteBuffer.put(flag).putInt(messageId).putInt(messageLength);
+        byteBuffer.put(flag).putInt(messageId).putInt(IKE_HEADER_LENGTH + encodedMessageBodyLen);
     }
 }
diff --git a/src/java/com/android/ike/ikev2/message/IkeKePayload.java b/src/java/com/android/ike/ikev2/message/IkeKePayload.java
index a742d94..92e2238 100644
--- a/src/java/com/android/ike/ikev2/message/IkeKePayload.java
+++ b/src/java/com/android/ike/ikev2/message/IkeKePayload.java
@@ -16,7 +16,7 @@
 
 package com.android.ike.ikev2.message;
 
-import android.util.Pair;
+import android.annotation.Nullable;
 
 import com.android.ike.ikev2.IkeDhParams;
 import com.android.ike.ikev2.SaProposal;
@@ -27,9 +27,12 @@
 import java.math.BigInteger;
 import java.nio.ByteBuffer;
 import java.security.GeneralSecurityException;
+import java.security.InvalidAlgorithmParameterException;
 import java.security.KeyFactory;
 import java.security.KeyPair;
 import java.security.KeyPairGenerator;
+import java.security.NoSuchAlgorithmException;
+import java.security.ProviderException;
 import java.security.SecureRandom;
 
 import javax.crypto.KeyAgreement;
@@ -64,8 +67,22 @@
     /** Supported dhGroup falls into {@link DhGroup} */
     public final int dhGroup;
 
+    /** Public DH key for the recipient to calculate shared key. */
     public final byte[] keyExchangeData;
 
+    /** Flag indicates if this is an outbound payload. */
+    public final boolean isOutbound;
+
+    /**
+     * localPrivateKey caches the locally generated private key when building an outbound KE
+     * payload. It will not be sent out. It is only used to calculate DH shared
+     * key when IKE library receives a public key from the remote server.
+     *
+     * <p>localPrivateKey of a inbound payload will be set to null. Caller MUST ensure its an
+     * outbound payload before using localPrivateKey.
+     */
+    @Nullable public final DHPrivateKeySpec localPrivateKey;
+
     /**
      * Construct an instance of IkeKePayload in the context of IkePayloadFactory
      *
@@ -79,6 +96,9 @@
     IkeKePayload(boolean critical, byte[] payloadBody) throws IkeException {
         super(PAYLOAD_TYPE_KE, critical);
 
+        isOutbound = false;
+        localPrivateKey = null;
+
         ByteBuffer inputBuffer = ByteBuffer.wrap(payloadBody);
 
         dhGroup = Short.toUnsignedInt(inputBuffer.getShort());
@@ -110,17 +130,67 @@
     /**
      * Construct an instance of IkeKePayload for building an outbound packet.
      *
+     * <p>Generate a DH key pair. Cache the private key and and send out the public key as
+     * keyExchangeData.
+     *
      * <p>Critical bit in this payload must not be set as instructed in RFC 7296.
      *
      * @param dh DH group for this KE payload
-     * @param keData the Key Exchange data
      * @see <a href="https://tools.ietf.org/html/rfc7296#page-76">RFC 7296, Internet Key Exchange
      *     Protocol Version 2 (IKEv2), Critical.
      */
-    private IkeKePayload(@SaProposal.DhGroup int dh, byte[] keData) {
+    public IkeKePayload(@SaProposal.DhGroup int dh) {
         super(PAYLOAD_TYPE_KE, false);
+
         dhGroup = dh;
-        keyExchangeData = keData;
+        isOutbound = true;
+
+        BigInteger prime = BigInteger.ZERO;
+        int keySize = 0;
+        switch (dhGroup) {
+            case SaProposal.DH_GROUP_1024_BIT_MODP:
+                prime =
+                        BigIntegerUtils.unsignedHexStringToBigInteger(
+                                IkeDhParams.PRIME_1024_BIT_MODP);
+                keySize = DH_GROUP_1024_BIT_MODP_DATA_LEN;
+                break;
+            case SaProposal.DH_GROUP_2048_BIT_MODP:
+                prime =
+                        BigIntegerUtils.unsignedHexStringToBigInteger(
+                                IkeDhParams.PRIME_2048_BIT_MODP);
+                keySize = DH_GROUP_2048_BIT_MODP_DATA_LEN;
+                break;
+            default:
+                throw new IllegalArgumentException("DH group not supported: " + dh);
+        }
+
+        try {
+            BigInteger baseGen = BigInteger.valueOf(IkeDhParams.BASE_GENERATOR_MODP);
+            DHParameterSpec dhParams = new DHParameterSpec(prime, baseGen);
+
+            KeyPairGenerator dhKeyPairGen =
+                    KeyPairGenerator.getInstance(
+                            KEY_EXCHANGE_ALGORITHM, IkeMessage.getSecurityProvider());
+            // By default SecureRandom uses AndroidOpenSSL provided SHA1PRNG Algorithm, which takes
+            // /dev/urandom as seed source.
+            dhKeyPairGen.initialize(dhParams, new SecureRandom());
+
+            KeyPair keyPair = dhKeyPairGen.generateKeyPair();
+
+            DHPrivateKey privateKey = (DHPrivateKey) keyPair.getPrivate();
+            DHPrivateKeySpec dhPrivateKeyspec =
+                    new DHPrivateKeySpec(privateKey.getX(), prime, baseGen);
+            DHPublicKey publicKey = (DHPublicKey) keyPair.getPublic();
+
+            // Zero-pad the public key without the sign bit
+            keyExchangeData =
+                    BigIntegerUtils.bigIntegerToUnsignedByteArray(publicKey.getY(), keySize);
+            localPrivateKey = dhPrivateKeyspec;
+        } catch (NoSuchAlgorithmException e) {
+            throw new ProviderException("Failed to obtain " + KEY_EXCHANGE_ALGORITHM, e);
+        } catch (InvalidAlgorithmParameterException e) {
+            throw new IllegalArgumentException("Failed to initialize key generator", e);
+        }
     }
 
     /**
@@ -149,56 +219,6 @@
     }
 
     /**
-     * Construct an instance of IkeKePayload according to its {@link DhGroup}.
-     *
-     * @param dh the Dh-Group. It should be in {@link DhGroup}
-     * @return Pair of generated private key and an instance of IkeKePayload with key exchange data.
-     * @throws GeneralSecurityException for security-related exception.
-     */
-    public static Pair<DHPrivateKeySpec, IkeKePayload> getKePayload(@SaProposal.DhGroup int dh)
-            throws GeneralSecurityException {
-        BigInteger baseGen = BigInteger.valueOf(IkeDhParams.BASE_GENERATOR_MODP);
-        BigInteger prime = BigInteger.ZERO;
-        int keySize = 0;
-        switch (dh) {
-            case SaProposal.DH_GROUP_1024_BIT_MODP:
-                prime =
-                        BigIntegerUtils.unsignedHexStringToBigInteger(
-                                IkeDhParams.PRIME_1024_BIT_MODP);
-                keySize = DH_GROUP_1024_BIT_MODP_DATA_LEN;
-                break;
-            case SaProposal.DH_GROUP_2048_BIT_MODP:
-                prime =
-                        BigIntegerUtils.unsignedHexStringToBigInteger(
-                                IkeDhParams.PRIME_2048_BIT_MODP);
-                keySize = DH_GROUP_2048_BIT_MODP_DATA_LEN;
-                break;
-            default:
-                throw new IllegalArgumentException("DH group not supported: " + dh);
-        }
-
-        DHParameterSpec dhParams = new DHParameterSpec(prime, baseGen);
-
-        KeyPairGenerator dhKeyPairGen =
-                KeyPairGenerator.getInstance(
-                        KEY_EXCHANGE_ALGORITHM, IkeMessage.getSecurityProvider());
-        // By default SecureRandom uses AndroidOpenSSL provided SHA1PRNG Algorithm, which takes
-        // /dev/urandom as seed source.
-        dhKeyPairGen.initialize(dhParams, new SecureRandom());
-
-        KeyPair keyPair = dhKeyPairGen.generateKeyPair();
-
-        DHPrivateKey privateKey = (DHPrivateKey) keyPair.getPrivate();
-        DHPrivateKeySpec dhPrivateKeyspec = new DHPrivateKeySpec(privateKey.getX(), prime, baseGen);
-        DHPublicKey publicKey = (DHPublicKey) keyPair.getPublic();
-
-        // Zero-pad the public key without the sign bit
-        byte[] keData = BigIntegerUtils.bigIntegerToUnsignedByteArray(publicKey.getY(), keySize);
-
-        return new Pair(dhPrivateKeyspec, new IkeKePayload(dh, keData));
-    }
-
-    /**
      * Calculate the shared secret.
      *
      * @param privateKeySpec contains the local private key, DH prime and DH base generator.
diff --git a/src/java/com/android/ike/ikev2/message/IkeMessage.java b/src/java/com/android/ike/ikev2/message/IkeMessage.java
index cfdf07b..5b63ee7 100644
--- a/src/java/com/android/ike/ikev2/message/IkeMessage.java
+++ b/src/java/com/android/ike/ikev2/message/IkeMessage.java
@@ -103,7 +103,15 @@
         ikePayloadList = payloadList;
     }
 
-    static Provider getSecurityProvider() {
+    /**
+     * Get security provider for IKE library
+     *
+     * <p>Use BouncyCastleProvider as the default security provider.
+     *
+     * @return the security provider of IKE library.
+     */
+    public static Provider getSecurityProvider() {
+        // TODO: Move this getter out of IKE message package since not only this package uses it.
         return SECURITY_PROVIDER;
     }
 
@@ -231,7 +239,7 @@
     byte[] attachEncodedHeader(byte[] encodedIkeBody) {
         ByteBuffer outputBuffer =
                 ByteBuffer.allocate(IkeHeader.IKE_HEADER_LENGTH + encodedIkeBody.length);
-        ikeHeader.encodeToByteBuffer(outputBuffer);
+        ikeHeader.encodeToByteBuffer(outputBuffer, encodedIkeBody.length);
         outputBuffer.put(encodedIkeBody);
         return outputBuffer.array();
     }
@@ -344,7 +352,7 @@
 
             ByteBuffer outputBuffer =
                     ByteBuffer.allocate(IkeHeader.IKE_HEADER_LENGTH + skPayload.getPayloadLength());
-            ikeHeader.encodeToByteBuffer(outputBuffer);
+            ikeHeader.encodeToByteBuffer(outputBuffer, skPayload.getPayloadLength());
             skPayload.encodeToByteBuffer(firstPayload, outputBuffer);
 
             return outputBuffer.array();
@@ -352,7 +360,7 @@
 
         @Override
         public IkeMessage decode(IkeHeader header, byte[] inputPacket) throws IkeException {
-            header.checkValidOrThrow(inputPacket.length);
+            header.checkInboundValidOrThrow(inputPacket.length);
 
             byte[] unencryptedPayloads =
                     Arrays.copyOfRange(
@@ -389,7 +397,7 @@
                 SecretKey dKey)
                 throws IkeException, GeneralSecurityException {
 
-            header.checkValidOrThrow(inputPacket.length);
+            header.checkInboundValidOrThrow(inputPacket.length);
 
             if (header.nextPayloadType != IkePayload.PAYLOAD_TYPE_SK) {
                 // TODO: b/123372339 Handle message containing unprotected payloads.
diff --git a/tests/iketests/src/java/com/android/ike/ikev2/ChildSessionStateMachineTest.java b/tests/iketests/src/java/com/android/ike/ikev2/ChildSessionStateMachineTest.java
index e81b17a..1d760b7 100644
--- a/tests/iketests/src/java/com/android/ike/ikev2/ChildSessionStateMachineTest.java
+++ b/tests/iketests/src/java/com/android/ike/ikev2/ChildSessionStateMachineTest.java
@@ -63,12 +63,13 @@
 
     private ISaRecordHelper mMockSaRecordHelper;
     private IChildSessionCallback mMockChildSessionCallback;
-    private ChildSessionOptions mMockChildSessionOptions;
+    private ChildSessionOptions mChildSessionOptions;
 
     public ChildSessionStateMachineTest() {
         mMockSaRecordHelper = mock(SaRecord.ISaRecordHelper.class);
         mMockChildSessionCallback = mock(IChildSessionCallback.class);
-        mMockChildSessionOptions = mock(ChildSessionOptions.class);
+
+        mChildSessionOptions = new ChildSessionOptions();
     }
 
     @Before
@@ -77,7 +78,7 @@
         mLooper = new TestLooper();
         mChildSessionStateMachine =
                 new ChildSessionStateMachine(
-                        "ChildSessionStateMachine", mLooper.getLooper(), mMockChildSessionOptions);
+                        "ChildSessionStateMachine", mLooper.getLooper(), mChildSessionOptions);
         mChildSessionStateMachine.setDbg(true);
         SaRecord.setSaRecordHelper(mMockSaRecordHelper);
 
diff --git a/tests/iketests/src/java/com/android/ike/ikev2/IkeSessionStateMachineTest.java b/tests/iketests/src/java/com/android/ike/ikev2/IkeSessionStateMachineTest.java
index 02aaa4c..1dcf5b8 100644
--- a/tests/iketests/src/java/com/android/ike/ikev2/IkeSessionStateMachineTest.java
+++ b/tests/iketests/src/java/com/android/ike/ikev2/IkeSessionStateMachineTest.java
@@ -17,6 +17,9 @@
 package com.android.ike.ikev2;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotEquals;
+import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.Matchers.any;
 import static org.mockito.Matchers.eq;
@@ -26,8 +29,13 @@
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import android.content.Context;
+import android.net.IpSecManager;
+import android.net.IpSecManager.UdpEncapsulationSocket;
+import android.os.Looper;
 import android.os.test.TestLooper;
-import android.util.Pair;
+
+import androidx.test.InstrumentationRegistry;
 
 import com.android.ike.ikev2.ChildSessionStateMachineFactory.ChildSessionFactoryHelper;
 import com.android.ike.ikev2.ChildSessionStateMachineFactory.IChildSessionFactoryHelper;
@@ -44,35 +52,46 @@
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
+import org.mockito.ArgumentCaptor;
 
+import java.net.InetAddress;
 import java.util.LinkedList;
+import java.util.List;
 
 public final class IkeSessionStateMachineTest {
 
+    private static final String SERVER_ADDRESS = "192.0.2.100";
+
+    private UdpEncapsulationSocket mUdpEncapSocket;
+
     private TestLooper mLooper;
     private IkeSessionStateMachine mIkeSessionStateMachine;
 
+    private IkeSessionOptions mIkeSessionOptions;
+    private ChildSessionOptions mChildSessionOptions;
+
     private IIkeMessageHelper mMockIkeMessageHelper;
     private ISaRecordHelper mMockSaRecordHelper;
-    private IkeSessionOptions mMockIkeSessionOptions;
 
     private ChildSessionStateMachine mMockChildSessionStateMachine;
-    private ChildSessionOptions mMockChildSessionOptions;
     private IChildSessionFactoryHelper mMockChildSessionFactoryHelper;
 
     private IkeSaRecord mSpyCurrentIkeSaRecord;
     private IkeSaRecord mSpyLocalInitIkeSaRecord;
     private IkeSaRecord mSpyRemoteInitIkeSaRecord;
 
+    private ArgumentCaptor<IkeMessage> mIkeMessageCaptor =
+            ArgumentCaptor.forClass(IkeMessage.class);
+
     private ReceivedIkePacket makeDummyUnencryptedReceivedIkePacket(int packetType)
             throws Exception {
         IkeMessage dummyIkeMessage = makeDummyIkeMessageForTest(0, 0, false, false);
-        Pair<IkeHeader, byte[]> dummyIkePacketPair =
-                new Pair<>(dummyIkeMessage.ikeHeader, new byte[0]);
-        when(mMockIkeMessageHelper.decode(dummyIkePacketPair.first, dummyIkePacketPair.second))
+        byte[] dummyIkePacketBytes = new byte[0];
+
+        when(mMockIkeMessageHelper.decode(dummyIkeMessage.ikeHeader, dummyIkePacketBytes))
                 .thenReturn(dummyIkeMessage);
         when(mMockIkeMessageHelper.getMessageType(dummyIkeMessage)).thenReturn(packetType);
-        return new ReceivedIkePacket(dummyIkePacketPair);
+        return new ReceivedIkePacket(dummyIkeMessage.ikeHeader, dummyIkePacketBytes);
     }
 
     private ReceivedIkePacket makeDummyEncryptedReceivedIkePacket(
@@ -81,16 +100,16 @@
         IkeMessage dummyIkeMessage =
                 makeDummyIkeMessageForTest(
                         ikeSaRecord.initiatorSpi, ikeSaRecord.responderSpi, fromIkeInit, true);
-        Pair<IkeHeader, byte[]> dummyIkePacketPair =
-                new Pair<>(dummyIkeMessage.ikeHeader, new byte[0]);
+        byte[] dummyIkePacketBytes = new byte[0];
+
         when(mMockIkeMessageHelper.decode(
-                        mMockIkeSessionOptions,
+                        mIkeSessionOptions,
                         ikeSaRecord,
-                        dummyIkePacketPair.first,
-                        dummyIkePacketPair.second))
+                        dummyIkeMessage.ikeHeader,
+                        dummyIkePacketBytes))
                 .thenReturn(dummyIkeMessage);
         when(mMockIkeMessageHelper.getMessageType(dummyIkeMessage)).thenReturn(packetType);
-        return new ReceivedIkePacket(dummyIkePacketPair);
+        return new ReceivedIkePacket(dummyIkeMessage.ikeHeader, dummyIkePacketBytes);
     }
 
     private IkeMessage makeDummyIkeMessageForTest(
@@ -98,27 +117,21 @@
         int firstPayloadType =
                 isEncrypted ? IkePayload.PAYLOAD_TYPE_SK : IkePayload.PAYLOAD_TYPE_NO_NEXT;
         IkeHeader header =
-                new IkeHeader(initSpi, respSpi, firstPayloadType, 0, true, fromikeInit, 0, 0);
+                new IkeHeader(initSpi, respSpi, firstPayloadType, 0, true, fromikeInit, 0);
         return new IkeMessage(header, new LinkedList<IkePayload>());
     }
 
     private void verifyDecodeEncryptedMessage(IkeSaRecord record, ReceivedIkePacket rcvPacket)
             throws Exception {
         verify(mMockIkeMessageHelper)
-                .decode(
-                        mMockIkeSessionOptions,
-                        record,
-                        rcvPacket.ikeHeader,
-                        rcvPacket.ikePacketBytes);
+                .decode(mIkeSessionOptions, record, rcvPacket.ikeHeader, rcvPacket.ikePacketBytes);
     }
 
     public IkeSessionStateMachineTest() {
         mMockIkeMessageHelper = mock(IkeMessage.IIkeMessageHelper.class);
         mMockSaRecordHelper = mock(SaRecord.ISaRecordHelper.class);
-        mMockIkeSessionOptions = mock(IkeSessionOptions.class);
 
         mMockChildSessionStateMachine = mock(ChildSessionStateMachine.class);
-        mMockChildSessionOptions = mock(ChildSessionOptions.class);
         mMockChildSessionFactoryHelper = mock(IChildSessionFactoryHelper.class);
 
         mSpyCurrentIkeSaRecord = spy(new IkeSaRecord(11, 12, true, null, null));
@@ -132,17 +145,25 @@
     }
 
     @Before
-    public void setUp() {
+    public void setUp() throws Exception {
+        Context context = InstrumentationRegistry.getContext();
+        IpSecManager ipSecManager = (IpSecManager) context.getSystemService(Context.IPSEC_SERVICE);
+        mUdpEncapSocket = ipSecManager.openUdpEncapsulationSocket();
+
+        mIkeSessionOptions = buildIkeSessionOptions();
+        mChildSessionOptions = new ChildSessionOptions();
+
         // Setup thread and looper
         mLooper = new TestLooper();
         mIkeSessionStateMachine =
                 new IkeSessionStateMachine(
                         "IkeSessionStateMachine",
                         mLooper.getLooper(),
-                        mMockIkeSessionOptions,
-                        mMockChildSessionOptions);
+                        mIkeSessionOptions,
+                        mChildSessionOptions);
         mIkeSessionStateMachine.setDbg(true);
         mIkeSessionStateMachine.start();
+
         IkeMessage.setIkeMessageHelper(mMockIkeMessageHelper);
         SaRecord.setSaRecordHelper(mMockSaRecordHelper);
         ChildSessionStateMachineFactory.setChildSessionFactoryHelper(
@@ -150,17 +171,46 @@
     }
 
     @After
-    public void tearDown() {
+    public void tearDown() throws Exception {
         mIkeSessionStateMachine.quit();
         mIkeSessionStateMachine.setDbg(false);
+        mUdpEncapSocket.close();
+
         IkeMessage.setIkeMessageHelper(new IkeMessageHelper());
         SaRecord.setSaRecordHelper(new SaRecordHelper());
         ChildSessionStateMachineFactory.setChildSessionFactoryHelper(
                 new ChildSessionFactoryHelper());
     }
 
+    private IkeSessionOptions buildIkeSessionOptions() throws Exception {
+        SaProposal saProposal =
+                SaProposal.Builder.newIkeSaProposalBuilder()
+                        .addEncryptionAlgorithm(
+                                SaProposal.ENCRYPTION_ALGORITHM_AES_CBC, SaProposal.KEY_LEN_AES_128)
+                        .addIntegrityAlgorithm(SaProposal.INTEGRITY_ALGORITHM_HMAC_SHA1_96)
+                        .addPseudorandomFunction(SaProposal.PSEUDORANDOM_FUNCTION_HMAC_SHA1)
+                        .addDhGroup(SaProposal.DH_GROUP_1024_BIT_MODP)
+                        .build();
+
+        InetAddress serveAddress = InetAddress.getByName(SERVER_ADDRESS);
+        IkeSessionOptions sessionOptions =
+                new IkeSessionOptions.Builder(serveAddress, mUdpEncapSocket)
+                        .addSaProposal(saProposal)
+                        .build();
+        return sessionOptions;
+    }
+
+    private static boolean isIkePayloadExist(
+            List<IkePayload> payloadList, @IkePayload.PayloadType int payloadType) {
+        for (IkePayload payload : payloadList) {
+            if (payload.payloadType == payloadType) return true;
+        }
+        return false;
+    }
+
     @Test
     public void testCreateIkeLocalIkeInit() throws Exception {
+        if (Looper.myLooper() == null) Looper.myLooper().prepare();
         // Mock IKE_INIT response.
         ReceivedIkePacket dummyReceivedIkePacket =
                 makeDummyUnencryptedReceivedIkePacket(IkeMessage.MESSAGE_TYPE_IKE_INIT_RESP);
@@ -172,15 +222,37 @@
                 IkeSessionStateMachine.CMD_RECEIVE_IKE_PACKET, dummyReceivedIkePacket);
 
         mLooper.dispatchAll();
+
+        // Validate outbound IKE INIT request
+        verify(mMockIkeMessageHelper).encode(mIkeMessageCaptor.capture());
+        IkeMessage ikeInitReqMessage = mIkeMessageCaptor.getValue();
+
+        IkeHeader ikeHeader = ikeInitReqMessage.ikeHeader;
+        assertEquals(IkeHeader.EXCHANGE_TYPE_IKE_SA_INIT, ikeHeader.exchangeType);
+        assertFalse(ikeHeader.isResponseMsg);
+        assertTrue(ikeHeader.fromIkeInitiator);
+
+        List<IkePayload> payloadList = ikeInitReqMessage.ikePayloadList;
+        assertTrue(isIkePayloadExist(payloadList, IkePayload.PAYLOAD_TYPE_SA));
+        assertTrue(isIkePayloadExist(payloadList, IkePayload.PAYLOAD_TYPE_KE));
+        assertTrue(isIkePayloadExist(payloadList, IkePayload.PAYLOAD_TYPE_NONCE));
+
+        IkeSocket ikeSocket = mIkeSessionStateMachine.mIkeSocket;
+        assertNotNull(ikeSocket);
+        assertNotEquals(
+                -1 /*not found*/, ikeSocket.mSpiToIkeSession.indexOfValue(mIkeSessionStateMachine));
+
         verify(mMockIkeMessageHelper)
                 .decode(dummyReceivedIkePacket.ikeHeader, dummyReceivedIkePacket.ikePacketBytes);
         verify(mMockIkeMessageHelper).getMessageType(any());
+
         assertTrue(
                 mIkeSessionStateMachine.getCurrentState()
                         instanceof IkeSessionStateMachine.CreateIkeLocalIkeAuth);
     }
 
     private void mockIkeSetup() throws Exception {
+        if (Looper.myLooper() == null) Looper.myLooper().prepare();
         // Mock IKE_INIT response
         ReceivedIkePacket dummyIkeInitRespReceivedPacket =
                 makeDummyUnencryptedReceivedIkePacket(IkeMessage.MESSAGE_TYPE_IKE_INIT_RESP);
diff --git a/tests/iketests/src/java/com/android/ike/ikev2/IkeSocketTest.java b/tests/iketests/src/java/com/android/ike/ikev2/IkeSocketTest.java
index 2dd4409..5f15f55 100644
--- a/tests/iketests/src/java/com/android/ike/ikev2/IkeSocketTest.java
+++ b/tests/iketests/src/java/com/android/ike/ikev2/IkeSocketTest.java
@@ -20,6 +20,11 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.fail;
+import static org.mockito.Mockito.any;
+import static org.mockito.Mockito.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
 
 import android.content.Context;
 import android.net.IpSecManager;
@@ -34,9 +39,14 @@
 
 import androidx.test.InstrumentationRegistry;
 
+import com.android.ike.ikev2.IkeSocket.PacketReceiver;
+import com.android.ike.ikev2.message.IkeHeader;
+import com.android.ike.ikev2.message.TestUtils;
+
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
+import org.mockito.ArgumentCaptor;
 
 import java.io.FileDescriptor;
 import java.net.InetAddress;
@@ -50,6 +60,26 @@
     private static final int REMOTE_RECV_BUFF_SIZE = 2048;
     private static final int TIMEOUT = 1000;
 
+    private static final String NON_ESP_MARKER_HEX_STRING = "00000000";
+    private static final String IKE_REQ_MESSAGE_HEX_STRING =
+            "5f54bf6d8b48e6e100000000000000002120220800000000"
+                    + "00000150220000300000002c010100040300000c0100000c"
+                    + "800e00800300000803000002030000080400000200000008"
+                    + "020000022800008800020000b4a2faf4bb54878ae21d6385"
+                    + "12ece55d9236fc5046ab6cef82220f421f3ce6361faf3656"
+                    + "4ecb6d28798a94aad7b2b4b603ddeaaa5630adb9ece8ac37"
+                    + "534036040610ebdd92f46bef84f0be7db860351843858f8a"
+                    + "cf87056e272377f70c9f2d81e29c7b0ce4f291a3a72476bb"
+                    + "0b278fd4b7b0a4c26bbeb08214c707137607958729000024"
+                    + "c39b7f368f4681b89fa9b7be6465abd7c5f68b6ed5d3b4c7"
+                    + "2cb4240eb5c464122900001c00004004e54f73b7d83f6beb"
+                    + "881eab2051d8663f421d10b02b00001c00004005d915368c"
+                    + "a036004cb578ae3e3fb268509aeab1900000002069936922"
+                    + "8741c6d4ca094c93e242c9de19e7b7c60000000500000500";
+
+    private static final String LOCAL_SPI = "0000000000000000";
+    private static final String REMOTE_SPI = "5f54bf6d8b48e6e1";
+
     private static final String DATA_ONE = "one 1";
     private static final String DATA_TWO = "two 2";
 
@@ -58,10 +88,18 @@
     private byte[] mDataOne;
     private byte[] mDataTwo;
 
+    private long mLocalSpi;
+    private long mRemoteSpi;
+
+    private LongSparseArray mSpiToIkeStateMachineMap;
+    private PacketReceiver mPacketReceiver;
+
     private UdpEncapsulationSocket mClientUdpEncapSocket;
     private InetAddress mLocalAddress;
     private FileDescriptor mDummyRemoteServerFd;
 
+    private IkeSessionStateMachine mMockIkeSessionStateMachine;
+
     @Before
     public void setUp() throws Exception {
         Context context = InstrumentationRegistry.getContext();
@@ -73,12 +111,24 @@
 
         mDataOne = DATA_ONE.getBytes("UTF-8");
         mDataTwo = DATA_TWO.getBytes("UTF-8");
+
+        ByteBuffer localSpiBuffer = ByteBuffer.wrap(TestUtils.hexStringToByteArray(LOCAL_SPI));
+        mLocalSpi = localSpiBuffer.getLong();
+        ByteBuffer remoteSpiBuffer = ByteBuffer.wrap(TestUtils.hexStringToByteArray(REMOTE_SPI));
+        mRemoteSpi = remoteSpiBuffer.getLong();
+
+        mMockIkeSessionStateMachine = mock(IkeSessionStateMachine.class);
+
+        mSpiToIkeStateMachineMap = new LongSparseArray<IkeSessionStateMachine>();
+        mSpiToIkeStateMachineMap.put(mLocalSpi, mMockIkeSessionStateMachine);
+
+        mPacketReceiver = new IkeSocket.PacketReceiver();
     }
 
     @After
     public void tearDown() throws Exception {
         mClientUdpEncapSocket.close();
-        IkeSocket.setPacketReceiver(new IkeSocket.PacketReceiver());
+        IkeSocket.setPacketReceiver(mPacketReceiver);
         Os.close(mDummyRemoteServerFd);
     }
 
@@ -120,7 +170,7 @@
         // Verify received data
         ByteBuffer expectedBuffer =
                 ByteBuffer.allocate(IkeSocket.NON_ESP_MARKER_LEN + mDataOne.length);
-        expectedBuffer.put(new byte[IkeSocket.NON_ESP_MARKER_LEN]).put(mDataOne);
+        expectedBuffer.put(IkeSocket.NON_ESP_MARKER).put(mDataOne);
 
         assertArrayEquals(expectedBuffer.array(), receivedData);
 
@@ -187,6 +237,49 @@
         mIkeThread.quitSafely();
     }
 
+    @Test
+    public void testHandlePacket() throws Exception {
+        byte[] recvBuf =
+                TestUtils.hexStringToByteArray(
+                        NON_ESP_MARKER_HEX_STRING + IKE_REQ_MESSAGE_HEX_STRING);
+
+        mPacketReceiver.handlePacket(recvBuf, mSpiToIkeStateMachineMap);
+
+        byte[] expectedIkePacketBytes = TestUtils.hexStringToByteArray(IKE_REQ_MESSAGE_HEX_STRING);
+        ArgumentCaptor<IkeHeader> ikeHeaderCaptor = ArgumentCaptor.forClass(IkeHeader.class);
+        verify(mMockIkeSessionStateMachine)
+                .receiveIkePacket(ikeHeaderCaptor.capture(), eq(expectedIkePacketBytes));
+
+        IkeHeader capturedIkeHeader = ikeHeaderCaptor.getValue();
+        assertEquals(mRemoteSpi, capturedIkeHeader.ikeInitiatorSpi);
+        assertEquals(mLocalSpi, capturedIkeHeader.ikeResponderSpi);
+    }
+
+    @Test
+    public void testHandleEspPacket() throws Exception {
+        byte[] recvBuf =
+                TestUtils.hexStringToByteArray(
+                        NON_ESP_MARKER_HEX_STRING + IKE_REQ_MESSAGE_HEX_STRING);
+        // Modify Non-ESP Marker
+        recvBuf[0] = 1;
+
+        mPacketReceiver.handlePacket(recvBuf, mSpiToIkeStateMachineMap);
+
+        verify(mMockIkeSessionStateMachine, never()).receiveIkePacket(any(), any());
+    }
+
+    @Test
+    public void testHandlePacketWithMalformedHeader() throws Exception {
+        String malformedIkePacketHexString = "5f54bf6d8b48e6e100000000000000002120220800000000";
+        byte[] recvBuf =
+                TestUtils.hexStringToByteArray(
+                        NON_ESP_MARKER_HEX_STRING + malformedIkePacketHexString);
+
+        mPacketReceiver.handlePacket(recvBuf, mSpiToIkeStateMachineMap);
+
+        verify(mMockIkeSessionStateMachine, never()).receiveIkePacket(any(), any());
+    }
+
     private byte[] receive(FileDescriptor mfd) throws Exception {
         byte[] receiveBuffer = new byte[REMOTE_RECV_BUFF_SIZE];
         AtomicInteger bytesRead = new AtomicInteger(-1);
diff --git a/tests/iketests/src/java/com/android/ike/ikev2/SaProposalTest.java b/tests/iketests/src/java/com/android/ike/ikev2/SaProposalTest.java
index 428c028..7f40d72 100644
--- a/tests/iketests/src/java/com/android/ike/ikev2/SaProposalTest.java
+++ b/tests/iketests/src/java/com/android/ike/ikev2/SaProposalTest.java
@@ -66,16 +66,17 @@
                         .addDhGroup(SaProposal.DH_GROUP_1024_BIT_MODP)
                         .build();
 
-        assertEquals(IkePayload.PROTOCOL_ID_IKE, proposal.mProtocolId);
+        assertEquals(IkePayload.PROTOCOL_ID_IKE, proposal.getProtocolId());
         assertArrayEquals(
                 new EncryptionTransform[] {mEncryption3DesTransform},
-                proposal.mEncryptionAlgorithms);
+                proposal.getEncryptionTransforms());
         assertArrayEquals(
                 new IntegrityTransform[] {mIntegrityHmacSha1Transform},
-                proposal.mIntegrityAlgorithms);
+                proposal.getIntegrityTransforms());
         assertArrayEquals(
-                new PrfTransform[] {mPrfAes128XCbcTransform}, proposal.mPseudorandomFunctions);
-        assertArrayEquals(new DhGroupTransform[] {mDhGroup1024Transform}, proposal.mDhGroups);
+                new PrfTransform[] {mPrfAes128XCbcTransform}, proposal.getPrfTransforms());
+        assertArrayEquals(
+                new DhGroupTransform[] {mDhGroup1024Transform}, proposal.getDhGroupTransforms());
     }
 
     @Test
@@ -89,14 +90,15 @@
                         .addDhGroup(SaProposal.DH_GROUP_1024_BIT_MODP)
                         .build();
 
-        assertEquals(IkePayload.PROTOCOL_ID_IKE, proposal.mProtocolId);
+        assertEquals(IkePayload.PROTOCOL_ID_IKE, proposal.getProtocolId());
         assertArrayEquals(
                 new EncryptionTransform[] {mEncryptionAesGcm8Transform},
-                proposal.mEncryptionAlgorithms);
+                proposal.getEncryptionTransforms());
         assertArrayEquals(
-                new PrfTransform[] {mPrfAes128XCbcTransform}, proposal.mPseudorandomFunctions);
-        assertArrayEquals(new DhGroupTransform[] {mDhGroup1024Transform}, proposal.mDhGroups);
-        assertTrue(proposal.mIntegrityAlgorithms.length == 0);
+                new PrfTransform[] {mPrfAes128XCbcTransform}, proposal.getPrfTransforms());
+        assertArrayEquals(
+                new DhGroupTransform[] {mDhGroup1024Transform}, proposal.getDhGroupTransforms());
+        assertTrue(proposal.getIntegrityTransforms().length == 0);
     }
 
     @Test
@@ -109,14 +111,15 @@
                         .addIntegrityAlgorithm(SaProposal.INTEGRITY_ALGORITHM_NONE)
                         .build();
 
-        assertEquals(IkePayload.PROTOCOL_ID_ESP, proposal.mProtocolId);
+        assertEquals(IkePayload.PROTOCOL_ID_ESP, proposal.getProtocolId());
         assertArrayEquals(
                 new EncryptionTransform[] {mEncryptionAesGcm8Transform},
-                proposal.mEncryptionAlgorithms);
+                proposal.getEncryptionTransforms());
         assertArrayEquals(
-                new IntegrityTransform[] {mIntegrityNoneTransform}, proposal.mIntegrityAlgorithms);
-        assertTrue(proposal.mPseudorandomFunctions.length == 0);
-        assertTrue(proposal.mDhGroups.length == 0);
+                new IntegrityTransform[] {mIntegrityNoneTransform},
+                proposal.getIntegrityTransforms());
+        assertTrue(proposal.getPrfTransforms().length == 0);
+        assertTrue(proposal.getDhGroupTransforms().length == 0);
     }
 
     @Test
@@ -129,14 +132,16 @@
                         .addDhGroup(SaProposal.DH_GROUP_1024_BIT_MODP)
                         .build();
 
-        assertEquals(IkePayload.PROTOCOL_ID_ESP, proposal.mProtocolId);
+        assertEquals(IkePayload.PROTOCOL_ID_ESP, proposal.getProtocolId());
         assertArrayEquals(
                 new EncryptionTransform[] {mEncryption3DesTransform},
-                proposal.mEncryptionAlgorithms);
+                proposal.getEncryptionTransforms());
         assertArrayEquals(
-                new IntegrityTransform[] {mIntegrityNoneTransform}, proposal.mIntegrityAlgorithms);
-        assertArrayEquals(new DhGroupTransform[] {mDhGroup1024Transform}, proposal.mDhGroups);
-        assertTrue(proposal.mPseudorandomFunctions.length == 0);
+                new IntegrityTransform[] {mIntegrityNoneTransform},
+                proposal.getIntegrityTransforms());
+        assertArrayEquals(
+                new DhGroupTransform[] {mDhGroup1024Transform}, proposal.getDhGroupTransforms());
+        assertTrue(proposal.getPrfTransforms().length == 0);
     }
 
     @Test
diff --git a/tests/iketests/src/java/com/android/ike/ikev2/SaRecordTest.java b/tests/iketests/src/java/com/android/ike/ikev2/SaRecordTest.java
new file mode 100644
index 0000000..5c61105
--- /dev/null
+++ b/tests/iketests/src/java/com/android/ike/ikev2/SaRecordTest.java
@@ -0,0 +1,138 @@
+/*
+ * Copyright (C) 2019 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 com.android.ike.ikev2;
+
+import static org.junit.Assert.assertArrayEquals;
+
+import com.android.ike.ikev2.message.TestUtils;
+
+import org.junit.Test;
+
+public final class SaRecordTest {
+    private static final String IKE_INIT_SPI = "5F54BF6D8B48E6E1";
+    private static final String IKE_RESP_SPI = "909232B3D1EDCB5C";
+
+    private static final String IKE_NONCE_INIT_HEX_STRING =
+            "C39B7F368F4681B89FA9B7BE6465ABD7C5F68B6ED5D3B4C72CB4240EB5C46412";
+    private static final String IKE_NONCE_RESP_HEX_STRING =
+            "9756112CA539F5C25ABACC7EE92B73091942A9C06950F98848F1AF1694C4DDFF";
+
+    private static final String IKE_SHARED_DH_KEY_HEX_STRING =
+            "C14155DEA40056BD9C76FB4819687B7A397582F4CD5AFF4B"
+                    + "8F441C56E0C08C84234147A0BA249A555835A048E3CA2980"
+                    + "7D057A61DD26EEFAD9AF9C01497005E52858E29FB42EB849"
+                    + "6731DF96A11CCE1F51137A9A1B900FA81AEE7898E373D4E4"
+                    + "8B899BBECA091314ECD4B6E412EF4B0FEF798F54735F3180"
+                    + "7424A318287F20E8";
+
+    private static final String IKE_SKEYSEED_HEX_STRING =
+            "8C42F3B1F5F81C7BAAC5F33E9A4F01987B2F9657";
+    private static final String IKE_SK_D_HEX_STRING = "C86B56EFCF684DCC2877578AEF3137167FE0EBF6";
+    private static final String IKE_SK_AUTH_INIT_HEX_STRING =
+            "554FBF5A05B7F511E05A30CE23D874DB9EF55E51";
+    private static final String IKE_SK_AUTH_RESP_HEX_STRING =
+            "36D83420788337CA32ECAA46892C48808DCD58B1";
+    private static final String IKE_SK_ENCR_INIT_HEX_STRING = "5CBFD33F75796C0188C4A3A546AEC4A1";
+    private static final String IKE_SK_ENCR_RESP_HEX_STRING = "C33B35FCF29514CD9D8B4A695E1A816E";
+    private static final String IKE_SK_PRF_INIT_HEX_STRING =
+            "094787780EE466E2CB049FA327B43908BC57E485";
+    private static final String IKE_SK_PRF_RESP_HEX_STRING =
+            "A30E6B08BE56C0E6BFF4744143C75219299E1BEB";
+    private static final String IKE_KEY_MAT =
+            IKE_SK_D_HEX_STRING
+                    + IKE_SK_AUTH_INIT_HEX_STRING
+                    + IKE_SK_AUTH_RESP_HEX_STRING
+                    + IKE_SK_ENCR_INIT_HEX_STRING
+                    + IKE_SK_ENCR_RESP_HEX_STRING
+                    + IKE_SK_PRF_INIT_HEX_STRING
+                    + IKE_SK_PRF_RESP_HEX_STRING;
+
+    private static final int IKE_AUTH_ALGO_KEY_LEN = 20;
+    private static final int IKE_ENCR_ALGO_KEY_LEN = 16;
+    private static final int IKE_PRF_KEY_LEN = 20;
+    private static final int IKE_SK_D_KEY_LEN = IKE_PRF_KEY_LEN;
+
+    private static final String FIRST_CHILD_ENCR_INIT_HEX_STRING =
+            "1B865CEA6E2C23973E8C5452ADC5CD7D";
+    private static final String FIRST_CHILD_ENCR_RESP_HEX_STRING =
+            "5E82FEDACC6DCB0756DDD7553907EBD1";
+    private static final String FIRST_CHILD_AUTH_INIT_HEX_STRING =
+            "A7A5A44F7EF4409657206C7DC52B7E692593B51E";
+    private static final String FIRST_CHILD_AUTH_RESP_HEX_STRING =
+            "CDE612189FD46DE870FAEC04F92B40B0BFDBD9E1";
+    private static final String FIRST_CHILD_KEY_MAT =
+            FIRST_CHILD_ENCR_INIT_HEX_STRING
+                    + FIRST_CHILD_AUTH_INIT_HEX_STRING
+                    + FIRST_CHILD_ENCR_RESP_HEX_STRING
+                    + FIRST_CHILD_AUTH_RESP_HEX_STRING;
+
+    private static final int FIRST_CHILD_AUTH_ALGO_KEY_LEN = 20;
+    private static final int FIRST_CHILD_ENCR_ALGO_KEY_LEN = 16;
+
+    private static final String PRF_HMAC_SHA1_ALGO_NAME = "HmacSHA1";
+
+    @Test
+    public void testCalculateSKeySeed() throws Exception {
+        byte[] nonceInit = TestUtils.hexStringToByteArray(IKE_NONCE_INIT_HEX_STRING);
+        byte[] nonceResp = TestUtils.hexStringToByteArray(IKE_NONCE_RESP_HEX_STRING);
+        byte[] sharedDhKey = TestUtils.hexStringToByteArray(IKE_SHARED_DH_KEY_HEX_STRING);
+
+        byte[] calculatedSKeySeed =
+                SaRecord.generateSKeySeed(
+                        PRF_HMAC_SHA1_ALGO_NAME, nonceInit, nonceResp, sharedDhKey);
+
+        byte[] expectedSKeySeed = TestUtils.hexStringToByteArray(IKE_SKEYSEED_HEX_STRING);
+        assertArrayEquals(expectedSKeySeed, calculatedSKeySeed);
+    }
+
+    @Test
+    public void testSignWithPrfPlusForIke() throws Exception {
+        byte[] prfKey = TestUtils.hexStringToByteArray(IKE_SKEYSEED_HEX_STRING);
+        byte[] prfData =
+                TestUtils.hexStringToByteArray(
+                        IKE_NONCE_INIT_HEX_STRING
+                                + IKE_NONCE_RESP_HEX_STRING
+                                + IKE_INIT_SPI
+                                + IKE_RESP_SPI);
+        int keyMaterialLen =
+                IKE_SK_D_KEY_LEN
+                        + IKE_AUTH_ALGO_KEY_LEN * 2
+                        + IKE_ENCR_ALGO_KEY_LEN * 2
+                        + IKE_PRF_KEY_LEN * 2;
+
+        byte[] calculatedKeyMat =
+                SaRecord.generateKeyMat(PRF_HMAC_SHA1_ALGO_NAME, prfKey, prfData, keyMaterialLen);
+
+        byte[] expectedKeyMat = TestUtils.hexStringToByteArray(IKE_KEY_MAT);
+        assertArrayEquals(expectedKeyMat, calculatedKeyMat);
+    }
+
+    @Test
+    public void testSignWithPrfPlusForFirstChild() throws Exception {
+        byte[] prfKey = TestUtils.hexStringToByteArray(IKE_SK_D_HEX_STRING);
+        byte[] prfData =
+                TestUtils.hexStringToByteArray(
+                        IKE_NONCE_INIT_HEX_STRING + IKE_NONCE_RESP_HEX_STRING);
+        int keyMaterialLen = FIRST_CHILD_AUTH_ALGO_KEY_LEN * 2 + FIRST_CHILD_ENCR_ALGO_KEY_LEN * 2;
+
+        byte[] calculatedKeyMat =
+                SaRecord.generateKeyMat(PRF_HMAC_SHA1_ALGO_NAME, prfKey, prfData, keyMaterialLen);
+
+        byte[] expectedKeyMat = TestUtils.hexStringToByteArray(FIRST_CHILD_KEY_MAT);
+        assertArrayEquals(expectedKeyMat, calculatedKeyMat);
+    }
+}
diff --git a/tests/iketests/src/java/com/android/ike/ikev2/message/IkeHeaderTest.java b/tests/iketests/src/java/com/android/ike/ikev2/message/IkeHeaderTest.java
index 1bc52c3..08b1612 100644
--- a/tests/iketests/src/java/com/android/ike/ikev2/message/IkeHeaderTest.java
+++ b/tests/iketests/src/java/com/android/ike/ikev2/message/IkeHeaderTest.java
@@ -62,6 +62,7 @@
 
     private static final int IKE_MSG_ID = 0;
     private static final int IKE_MSG_LENGTH = 336;
+    private static final int IKE_MSG_BODY_LENGTH = IKE_MSG_LENGTH - IkeHeader.IKE_HEADER_LENGTH;
 
     // Byte offsets of version field in IKE message header.
     private static final int VERSION_OFFSET = 17;
@@ -89,7 +90,7 @@
         assertFalse(header.isResponseMsg);
         assertTrue(header.fromIkeInitiator);
         assertEquals(IKE_MSG_ID, header.messageId);
-        assertEquals(IKE_MSG_LENGTH, header.messageLength);
+        assertEquals(IKE_MSG_LENGTH, header.getInboundMessageLength());
     }
 
     @Test
@@ -142,7 +143,7 @@
         IkeHeader header = new IkeHeader(inputPacket);
 
         ByteBuffer byteBuffer = ByteBuffer.allocate(IkeHeader.IKE_HEADER_LENGTH);
-        header.encodeToByteBuffer(byteBuffer);
+        header.encodeToByteBuffer(byteBuffer, IKE_MSG_BODY_LENGTH);
 
         byte[] expectedPacket = TestUtils.hexStringToByteArray(IKE_HEADER_HEX_STRING);
         assertArrayEquals(expectedPacket, byteBuffer.array());
diff --git a/tests/iketests/src/java/com/android/ike/ikev2/message/IkeKePayloadTest.java b/tests/iketests/src/java/com/android/ike/ikev2/message/IkeKePayloadTest.java
index 4f45f26..1bb0b70 100644
--- a/tests/iketests/src/java/com/android/ike/ikev2/message/IkeKePayloadTest.java
+++ b/tests/iketests/src/java/com/android/ike/ikev2/message/IkeKePayloadTest.java
@@ -18,11 +18,10 @@
 
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
-import android.util.Pair;
-
 import com.android.ike.ikev2.IkeDhParams;
 import com.android.ike.ikev2.SaProposal;
 import com.android.ike.ikev2.exceptions.InvalidSyntaxException;
@@ -100,6 +99,7 @@
 
         IkeKePayload payload = new IkeKePayload(CRITICAL_BIT, inputPacket);
 
+        assertFalse(payload.isOutbound);
         assertEquals(EXPECTED_DH_GROUP, payload.dhGroup);
 
         byte[] keyExchangeData = TestUtils.hexStringToByteArray(KEY_EXCHANGE_DATA_RAW_PACKET);
@@ -138,11 +138,11 @@
 
     @Test
     public void testGetIkeKePayload() throws Exception {
-        Pair<DHPrivateKeySpec, IkeKePayload> pair =
-                IkeKePayload.getKePayload(SaProposal.DH_GROUP_1024_BIT_MODP);
+        IkeKePayload payload = new IkeKePayload(SaProposal.DH_GROUP_1024_BIT_MODP);
 
         // Test DHPrivateKeySpec
-        DHPrivateKeySpec privateKeySpec = pair.first;
+        assertTrue(payload.isOutbound);
+        DHPrivateKeySpec privateKeySpec = payload.localPrivateKey;
 
         BigInteger primeValue = privateKeySpec.getP();
         BigInteger expectedPrimeValue = new BigInteger(IkeDhParams.PRIME_1024_BIT_MODP, 16);
@@ -153,8 +153,6 @@
         assertEquals(0, expectedGenValue.compareTo(genValue));
 
         // Test IkeKePayload
-        IkeKePayload payload = pair.second;
-
         assertEquals(EXPECTED_DH_GROUP, payload.dhGroup);
         assertEquals(EXPECTED_KE_DATA_LEN, payload.keyExchangeData.length);
     }