Merge changes I21e0b157,I8c90e134

* changes:
  Create IkeSocket using v6 UDP socket and dest port 4500
  Move common test code to IkeSocketTestBase
diff --git a/src/java/com/android/internal/net/ipsec/ike/IkeUdp6Socket.java b/src/java/com/android/internal/net/ipsec/ike/IkeUdp6Socket.java
index 799db7b..647583d 100644
--- a/src/java/com/android/internal/net/ipsec/ike/IkeUdp6Socket.java
+++ b/src/java/com/android/internal/net/ipsec/ike/IkeUdp6Socket.java
@@ -41,7 +41,7 @@
  * {@link Network} will only be bound to by one IkeUdp6Socket instance. When caller requests an
  * IkeUdp6Socket with an already bound {@link Network}, the existing instance will be returned.
  */
-public final class IkeUdp6Socket extends IkeUdpSocket {
+public class IkeUdp6Socket extends IkeUdpSocket {
     private static final String TAG = IkeUdp6Socket.class.getSimpleName();
 
     private static final InetAddress INADDR_ANY = InetAddresses.parseNumericAddress("::");
@@ -49,7 +49,7 @@
     // Map from Network to IkeUdp6Socket instances.
     private static Map<Network, IkeUdp6Socket> sNetworkToUdp6SocketMap = new HashMap<>();
 
-    private IkeUdp6Socket(FileDescriptor socket, Network network, Handler handler) {
+    protected IkeUdp6Socket(FileDescriptor socket, Network network, Handler handler) {
         super(socket, network, handler == null ? new Handler() : handler);
     }
 
@@ -69,14 +69,7 @@
             throws ErrnoException, IOException {
         IkeUdp6Socket ikeSocket = sNetworkToUdp6SocketMap.get(network);
         if (ikeSocket == null) {
-            FileDescriptor sock = Os.socket(AF_INET6, SOCK_DGRAM, IPPROTO_UDP);
-            Os.bind(sock, INADDR_ANY, 0);
-
-            // {@link PacketReader} requires non-blocking I/O access. Set SOCK_NONBLOCK here.
-            Os.fcntlInt(sock, F_SETFL, SOCK_DGRAM | SOCK_NONBLOCK);
-            network.bindSocket(sock);
-
-            ikeSocket = new IkeUdp6Socket(sock, network, handler);
+            ikeSocket = new IkeUdp6Socket(openUdp6SockNonBlock(network), network, handler);
 
             // Create and register FileDescriptor for receiving IKE packet on current thread.
             ikeSocket.start();
@@ -87,6 +80,18 @@
         return ikeSocket;
     }
 
+    protected static FileDescriptor openUdp6SockNonBlock(Network network)
+            throws ErrnoException, IOException {
+        FileDescriptor sock = Os.socket(AF_INET6, SOCK_DGRAM, IPPROTO_UDP);
+        Os.bind(sock, INADDR_ANY, 0);
+
+        // {@link PacketReader} requires non-blocking I/O access. Set SOCK_NONBLOCK here.
+        Os.fcntlInt(sock, F_SETFL, SOCK_DGRAM | SOCK_NONBLOCK);
+        network.bindSocket(sock);
+
+        return sock;
+    }
+
     /** Implement {@link AutoCloseable#close()} */
     @Override
     public void close() {
diff --git a/src/java/com/android/internal/net/ipsec/ike/IkeUdp6WithEncapPortSocket.java b/src/java/com/android/internal/net/ipsec/ike/IkeUdp6WithEncapPortSocket.java
new file mode 100644
index 0000000..a51670f
--- /dev/null
+++ b/src/java/com/android/internal/net/ipsec/ike/IkeUdp6WithEncapPortSocket.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 com.android.internal.net.ipsec.ike;
+
+import android.net.Network;
+import android.os.Handler;
+import android.system.ErrnoException;
+
+import java.io.FileDescriptor;
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * IkeUdp6WithEncapPortSocket uses an IPv6-bound {@link FileDescriptor} to send and receive IKE
+ * packets.
+ *
+ * <p>IkeUdp6WithEncapPortSocket is usually used when IKE Session has IPv6 address and is required
+ * to send message to port 4500, as per MOBIKE spec (RFC 4555).
+ *
+ * <p>Caller MUST provide one {@link Network} when trying to get an instance of
+ * IkeUdp6WithEncapPortSocket. Each {@link Network} will only be bound to by one
+ * IkeUdp6WithEncapPortSocket instance. When caller requests an IkeUdp6WithEncapPortSocket with an
+ * already bound {@link Network}, the existing instance will be returned.
+ */
+public final class IkeUdp6WithEncapPortSocket extends IkeUdp6Socket {
+    private static final String TAG = IkeUdp6WithEncapPortSocket.class.getSimpleName();
+
+    // Map from Network to IkeUdp6WithEncapPortSocket instances.
+    private static Map<Network, IkeUdp6WithEncapPortSocket> sNetworkToUdp6SocketMap =
+            new HashMap<>();
+
+    private IkeUdp6WithEncapPortSocket(FileDescriptor socket, Network network, Handler handler) {
+        super(socket, network, handler);
+    }
+
+    /**
+     * Get an IkeUdp6WithEncapPortSocket instance.
+     *
+     * <p>Return the existing IkeUdp6WithEncapPortSocket instance if it has been created for the
+     * input Network. Otherwise, create and return a new IkeUdp6WithEncapPortSocket instance.
+     *
+     * @param network the Network this socket will be bound to
+     * @param ikeSession the IkeSessionStateMachine that is requesting an
+     *     IkeUdp6WithEncapPortSocket.
+     * @param handler the Handler used to process received packets
+     * @return an IkeUdp6WithEncapPortSocket instance
+     */
+    public static IkeUdp6WithEncapPortSocket getInstance(
+            Network network, IkeSessionStateMachine ikeSession, Handler handler)
+            throws ErrnoException, IOException {
+        IkeUdp6WithEncapPortSocket ikeSocket = sNetworkToUdp6SocketMap.get(network);
+        if (ikeSocket == null) {
+            ikeSocket =
+                    new IkeUdp6WithEncapPortSocket(openUdp6SockNonBlock(network), network, handler);
+
+            // Create and register FileDescriptor for receiving IKE packet on current thread.
+            ikeSocket.start();
+
+            sNetworkToUdp6SocketMap.put(network, ikeSocket);
+        }
+        ikeSocket.mAliveIkeSessions.add(ikeSession);
+        return ikeSocket;
+    }
+
+    @Override
+    public int getIkeServerPort() {
+        return SERVER_PORT_UDP_ENCAPSULATED;
+    }
+
+    /** Implement {@link AutoCloseable#close()} */
+    @Override
+    public void close() {
+        sNetworkToUdp6SocketMap.remove(getNetwork());
+
+        super.close();
+    }
+}
diff --git a/tests/iketests/src/java/com/android/internal/net/ipsec/ike/IkeSocketTestBase.java b/tests/iketests/src/java/com/android/internal/net/ipsec/ike/IkeSocketTestBase.java
index 2363cff..8f13908 100644
--- a/tests/iketests/src/java/com/android/internal/net/ipsec/ike/IkeSocketTestBase.java
+++ b/tests/iketests/src/java/com/android/internal/net/ipsec/ike/IkeSocketTestBase.java
@@ -16,9 +16,16 @@
 
 package com.android.internal.net.ipsec.ike;
 
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
 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.reset;
+import static org.mockito.Mockito.verify;
 
 import android.net.InetAddresses;
 import android.net.Network;
@@ -28,8 +35,13 @@
 import android.util.Log;
 import android.util.LongSparseArray;
 
+import com.android.internal.net.TestUtils;
+import com.android.internal.net.ipsec.ike.message.IkeHeader;
+import com.android.internal.util.HexDump;
+
 import org.junit.After;
 import org.junit.Before;
+import org.mockito.ArgumentCaptor;
 
 import java.io.FileDescriptor;
 import java.io.IOException;
@@ -206,4 +218,113 @@
             reset();
         }
     }
+
+    protected interface IkeSocketFactory {
+        IkeSocket getIkeSocket(Network network, IkeSessionStateMachine ikeSession)
+                throws ErrnoException, IOException;
+    }
+
+    protected void verifyGetAndCloseIkeSocketSameNetwork(
+            IkeSocketFactory ikeUdpSocketFactory, int expectedServerPort) throws Exception {
+        IkeSessionStateMachine mockIkeSessionOne = mock(IkeSessionStateMachine.class);
+        IkeSessionStateMachine mockIkeSessionTwo = mock(IkeSessionStateMachine.class);
+
+        IkeSocket ikeSocketOne = ikeUdpSocketFactory.getIkeSocket(mMockNetwork, mockIkeSessionOne);
+        assertEquals(expectedServerPort, ikeSocketOne.getIkeServerPort());
+        assertEquals(1, ikeSocketOne.mAliveIkeSessions.size());
+
+        IkeSocket ikeSocketTwo = ikeUdpSocketFactory.getIkeSocket(mMockNetwork, mockIkeSessionTwo);
+        assertEquals(expectedServerPort, ikeSocketTwo.getIkeServerPort());
+        assertEquals(2, ikeSocketTwo.mAliveIkeSessions.size());
+        assertEquals(ikeSocketOne, ikeSocketTwo);
+
+        verify(mMockNetwork).bindSocket(eq(ikeSocketOne.getFd()));
+
+        ikeSocketOne.releaseReference(mockIkeSessionOne);
+        assertEquals(1, ikeSocketOne.mAliveIkeSessions.size());
+        assertTrue(isFdOpen(ikeSocketOne.getFd()));
+
+        ikeSocketTwo.releaseReference(mockIkeSessionTwo);
+        assertEquals(0, ikeSocketTwo.mAliveIkeSessions.size());
+        verifyCloseFd(ikeSocketTwo.getFd());
+    }
+
+    protected void verifyGetAndCloseIkeSocketDifferentNetwork(
+            IkeSocketFactory ikeUdpSocketFactory, int expectedServerPort) throws Exception {
+        IkeSessionStateMachine mockIkeSessionOne = mock(IkeSessionStateMachine.class);
+        IkeSessionStateMachine mockIkeSessionTwo = mock(IkeSessionStateMachine.class);
+
+        Network mockNetworkOne = mock(Network.class);
+        Network mockNetworkTwo = mock(Network.class);
+
+        IkeSocket ikeSocketOne =
+                ikeUdpSocketFactory.getIkeSocket(mockNetworkOne, mockIkeSessionOne);
+        assertEquals(expectedServerPort, ikeSocketOne.getIkeServerPort());
+        assertEquals(1, ikeSocketOne.mAliveIkeSessions.size());
+
+        IkeSocket ikeSocketTwo =
+                ikeUdpSocketFactory.getIkeSocket(mockNetworkTwo, mockIkeSessionTwo);
+        assertEquals(expectedServerPort, ikeSocketTwo.getIkeServerPort());
+        assertEquals(1, ikeSocketTwo.mAliveIkeSessions.size());
+
+        assertNotEquals(ikeSocketOne, ikeSocketTwo);
+
+        ArgumentCaptor<FileDescriptor> fdCaptorOne = ArgumentCaptor.forClass(FileDescriptor.class);
+        ArgumentCaptor<FileDescriptor> fdCaptorTwo = ArgumentCaptor.forClass(FileDescriptor.class);
+        verify(mockNetworkOne).bindSocket(fdCaptorOne.capture());
+        verify(mockNetworkTwo).bindSocket(fdCaptorTwo.capture());
+
+        FileDescriptor fdOne = fdCaptorOne.getValue();
+        FileDescriptor fdTwo = fdCaptorTwo.getValue();
+        assertNotNull(fdOne);
+        assertNotNull(fdTwo);
+        assertNotEquals(fdOne, fdTwo);
+
+        ikeSocketOne.releaseReference(mockIkeSessionOne);
+        assertEquals(0, ikeSocketOne.mAliveIkeSessions.size());
+        verifyCloseFd(ikeSocketOne.getFd());
+
+        ikeSocketTwo.releaseReference(mockIkeSessionTwo);
+        assertEquals(0, ikeSocketTwo.mAliveIkeSessions.size());
+        verifyCloseFd(ikeSocketTwo.getFd());
+    }
+
+    protected void verifyHandlePacket(
+            byte[] receivedPacket, IkeSocket.IPacketReceiver packetReceiver) throws Exception {
+        packetReceiver.handlePacket(receivedPacket, 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(REMOTE_SPI, capturedIkeHeader.ikeInitiatorSpi);
+        assertEquals(LOCAL_SPI, capturedIkeHeader.ikeResponderSpi);
+    }
+
+    protected void verifyIkeUdpSocketReceivePacket(
+            IkeSocketFactory ikeUdpSocketFactory, IkeSocket.IPacketReceiver packetReceiver)
+            throws Exception {
+        IkeSessionStateMachine mockIkeSession = mock(IkeSessionStateMachine.class);
+        IkeUdpSocket ikeSocket =
+                (IkeUdpSocket) ikeUdpSocketFactory.getIkeSocket(mMockNetwork, mockIkeSession);
+        assertNotNull(ikeSocket);
+
+        // Set up state
+        ikeSocket.registerIke(LOCAL_SPI, mockIkeSession);
+        IkeSocket.IPacketReceiver mockPacketReceiver = mock(IkeSocket.IPacketReceiver.class);
+        IkeUdpSocket.setPacketReceiver(mockPacketReceiver);
+        try {
+            // Send a packet
+            byte[] pktBytes = HexDump.hexStringToByteArray(IKE_REQ_MESSAGE_HEX_STRING);
+            ikeSocket.handlePacket(pktBytes, pktBytes.length);
+
+            verify(mockPacketReceiver).handlePacket(eq(pktBytes), any());
+
+        } finally {
+            ikeSocket.releaseReference(mockIkeSession);
+            IkeUdpSocket.setPacketReceiver(getPacketReceiver());
+        }
+    }
 }
diff --git a/tests/iketests/src/java/com/android/internal/net/ipsec/ike/IkeUdp4SocketTest.java b/tests/iketests/src/java/com/android/internal/net/ipsec/ike/IkeUdp4SocketTest.java
index ee05885..328533b 100644
--- a/tests/iketests/src/java/com/android/internal/net/ipsec/ike/IkeUdp4SocketTest.java
+++ b/tests/iketests/src/java/com/android/internal/net/ipsec/ike/IkeUdp4SocketTest.java
@@ -16,137 +16,56 @@
 
 package com.android.internal.net.ipsec.ike;
 
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertNotEquals;
-import static org.junit.Assert.assertNotNull;
-import static org.junit.Assert.assertTrue;
-import static org.mockito.Mockito.any;
-import static org.mockito.Mockito.eq;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.verify;
 
 import android.net.Network;
 import android.os.Handler;
 import android.os.test.TestLooper;
+import android.system.ErrnoException;
 
 import com.android.internal.net.TestUtils;
-import com.android.internal.net.ipsec.ike.message.IkeHeader;
-import com.android.internal.util.HexDump;
 
 import org.junit.Test;
-import org.mockito.ArgumentCaptor;
 
-import java.io.FileDescriptor;
+import java.io.IOException;
 
-// TODO: Combine this and IkeUdp6SocketTest, and take a Factory method as an input.
 public final class IkeUdp4SocketTest extends IkeSocketTestBase {
     private final TestLooper mLooper = new TestLooper();
     private final Handler mHandler = new Handler(mLooper.getLooper());
 
+    private final IkeSocketFactory mIkeSocketFactory =
+            new IkeSocketFactory() {
+                @Override
+                public IkeSocket getIkeSocket(Network network, IkeSessionStateMachine ikeSession)
+                        throws ErrnoException, IOException {
+                    return IkeUdp4Socket.getInstance(network, ikeSession, mHandler);
+                }
+            };
+
     @Override
     protected IkeSocket.IPacketReceiver getPacketReceiver() {
         return new IkeUdpSocket.PacketReceiver();
     }
 
     @Test
-    public void testGetAndCloseIkeUdp4SocketSameNetwork() throws Exception {
-        IkeSessionStateMachine mockIkeSessionOne = mock(IkeSessionStateMachine.class);
-        IkeSessionStateMachine mockIkeSessionTwo = mock(IkeSessionStateMachine.class);
-
-        IkeUdp4Socket ikeSocketOne =
-                IkeUdp4Socket.getInstance(mMockNetwork, mockIkeSessionOne, mHandler);
-        assertEquals(1, ikeSocketOne.mAliveIkeSessions.size());
-
-        IkeUdp4Socket ikeSocketTwo =
-                IkeUdp4Socket.getInstance(mMockNetwork, mockIkeSessionTwo, mHandler);
-        assertEquals(2, ikeSocketTwo.mAliveIkeSessions.size());
-        assertEquals(ikeSocketOne, ikeSocketTwo);
-
-        verify(mMockNetwork).bindSocket(eq(ikeSocketOne.getFd()));
-
-        ikeSocketOne.releaseReference(mockIkeSessionOne);
-        assertEquals(1, ikeSocketOne.mAliveIkeSessions.size());
-        assertTrue(isFdOpen(ikeSocketOne.getFd()));
-
-        ikeSocketTwo.releaseReference(mockIkeSessionTwo);
-        assertEquals(0, ikeSocketTwo.mAliveIkeSessions.size());
-        verifyCloseFd(ikeSocketTwo.getFd());
+    public void testGetAndCloseIkeUdp6SocketSameNetwork() throws Exception {
+        verifyGetAndCloseIkeSocketSameNetwork(
+                mIkeSocketFactory, IkeSocket.SERVER_PORT_NON_UDP_ENCAPSULATED);
     }
 
     @Test
-    public void testGetAndCloseIkeUdp4SocketDifferentNetwork() throws Exception {
-        IkeSessionStateMachine mockIkeSessionOne = mock(IkeSessionStateMachine.class);
-        IkeSessionStateMachine mockIkeSessionTwo = mock(IkeSessionStateMachine.class);
-
-        Network mockNetworkOne = mock(Network.class);
-        Network mockNetworkTwo = mock(Network.class);
-
-        IkeUdp4Socket ikeSocketOne =
-                IkeUdp4Socket.getInstance(mockNetworkOne, mockIkeSessionOne, mHandler);
-        assertEquals(1, ikeSocketOne.mAliveIkeSessions.size());
-
-        IkeUdp4Socket ikeSocketTwo =
-                IkeUdp4Socket.getInstance(mockNetworkTwo, mockIkeSessionTwo, mHandler);
-        assertEquals(1, ikeSocketTwo.mAliveIkeSessions.size());
-
-        assertNotEquals(ikeSocketOne, ikeSocketTwo);
-
-        ArgumentCaptor<FileDescriptor> fdCaptorOne = ArgumentCaptor.forClass(FileDescriptor.class);
-        ArgumentCaptor<FileDescriptor> fdCaptorTwo = ArgumentCaptor.forClass(FileDescriptor.class);
-        verify(mockNetworkOne).bindSocket(fdCaptorOne.capture());
-        verify(mockNetworkTwo).bindSocket(fdCaptorTwo.capture());
-
-        FileDescriptor fdOne = fdCaptorOne.getValue();
-        FileDescriptor fdTwo = fdCaptorTwo.getValue();
-        assertNotNull(fdOne);
-        assertNotNull(fdTwo);
-        assertNotEquals(fdOne, fdTwo);
-
-        ikeSocketOne.releaseReference(mockIkeSessionOne);
-        assertEquals(0, ikeSocketOne.mAliveIkeSessions.size());
-        verifyCloseFd(ikeSocketOne.getFd());
-
-        ikeSocketTwo.releaseReference(mockIkeSessionTwo);
-        assertEquals(0, ikeSocketTwo.mAliveIkeSessions.size());
-        verifyCloseFd(ikeSocketTwo.getFd());
+    public void testGetAndCloseIkeUdp6SocketDifferentNetwork() throws Exception {
+        verifyGetAndCloseIkeSocketDifferentNetwork(
+                mIkeSocketFactory, IkeSocket.SERVER_PORT_NON_UDP_ENCAPSULATED);
     }
 
     @Test
     public void testReceiveIkePacket() throws Exception {
-        IkeSessionStateMachine mockIkeSession = mock(IkeSessionStateMachine.class);
-        IkeUdp4Socket ikeSocket = IkeUdp4Socket.getInstance(mMockNetwork, mockIkeSession, mHandler);
-        assertNotNull(ikeSocket);
-
-        // Set up state
-        ikeSocket.registerIke(LOCAL_SPI, mockIkeSession);
-        IkeSocket.IPacketReceiver packetReceiver = mock(IkeSocket.IPacketReceiver.class);
-        IkeUdpSocket.setPacketReceiver(packetReceiver);
-        try {
-            // Send a packet
-            byte[] pktBytes = HexDump.hexStringToByteArray(IKE_REQ_MESSAGE_HEX_STRING);
-            ikeSocket.handlePacket(pktBytes, pktBytes.length);
-
-            verify(packetReceiver).handlePacket(eq(pktBytes), any());
-
-        } finally {
-            ikeSocket.releaseReference(mockIkeSession);
-            IkeUdpSocket.setPacketReceiver(getPacketReceiver());
-        }
+        verifyIkeUdpSocketReceivePacket(mIkeSocketFactory, getPacketReceiver());
     }
 
     @Test
     public void testHandlePacket() throws Exception {
-        byte[] recvBuf = TestUtils.hexStringToByteArray(IKE_REQ_MESSAGE_HEX_STRING);
-
-        getPacketReceiver().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(REMOTE_SPI, capturedIkeHeader.ikeInitiatorSpi);
-        assertEquals(LOCAL_SPI, capturedIkeHeader.ikeResponderSpi);
+        verifyHandlePacket(
+                TestUtils.hexStringToByteArray(IKE_REQ_MESSAGE_HEX_STRING), getPacketReceiver());
     }
 }
diff --git a/tests/iketests/src/java/com/android/internal/net/ipsec/ike/IkeUdp6SocketTest.java b/tests/iketests/src/java/com/android/internal/net/ipsec/ike/IkeUdp6SocketTest.java
index 7c6e30d..24a6e87 100644
--- a/tests/iketests/src/java/com/android/internal/net/ipsec/ike/IkeUdp6SocketTest.java
+++ b/tests/iketests/src/java/com/android/internal/net/ipsec/ike/IkeUdp6SocketTest.java
@@ -16,33 +16,31 @@
 
 package com.android.internal.net.ipsec.ike;
 
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertNotEquals;
-import static org.junit.Assert.assertNotNull;
-import static org.junit.Assert.assertTrue;
-import static org.mockito.Mockito.any;
-import static org.mockito.Mockito.eq;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.verify;
 
 import android.net.Network;
 import android.os.Handler;
 import android.os.test.TestLooper;
+import android.system.ErrnoException;
 
 import com.android.internal.net.TestUtils;
-import com.android.internal.net.ipsec.ike.message.IkeHeader;
-import com.android.internal.util.HexDump;
 
 import org.junit.Test;
-import org.mockito.ArgumentCaptor;
 
-import java.io.FileDescriptor;
+import java.io.IOException;
 
-// TODO: Combine this and IkeUdp4SocketTest, and take a Factory method as an input.
 public final class IkeUdp6SocketTest extends IkeSocketTestBase {
     private final TestLooper mLooper = new TestLooper();
     private final Handler mHandler = new Handler(mLooper.getLooper());
 
+    private final IkeSocketFactory mIkeSocketFactory =
+            new IkeSocketFactory() {
+                @Override
+                public IkeSocket getIkeSocket(Network network, IkeSessionStateMachine ikeSession)
+                        throws ErrnoException, IOException {
+                    return IkeUdp6Socket.getInstance(network, ikeSession, mHandler);
+                }
+            };
+
     @Override
     protected IkeSocket.IPacketReceiver getPacketReceiver() {
         return new IkeUdpSocket.PacketReceiver();
@@ -50,103 +48,24 @@
 
     @Test
     public void testGetAndCloseIkeUdp6SocketSameNetwork() throws Exception {
-        IkeSessionStateMachine mockIkeSessionOne = mock(IkeSessionStateMachine.class);
-        IkeSessionStateMachine mockIkeSessionTwo = mock(IkeSessionStateMachine.class);
-
-        IkeUdp6Socket ikeSocketOne =
-                IkeUdp6Socket.getInstance(mMockNetwork, mockIkeSessionOne, mHandler);
-        assertEquals(1, ikeSocketOne.mAliveIkeSessions.size());
-
-        IkeUdp6Socket ikeSocketTwo =
-                IkeUdp6Socket.getInstance(mMockNetwork, mockIkeSessionTwo, mHandler);
-        assertEquals(2, ikeSocketTwo.mAliveIkeSessions.size());
-        assertEquals(ikeSocketOne, ikeSocketTwo);
-
-        verify(mMockNetwork).bindSocket(eq(ikeSocketOne.getFd()));
-
-        ikeSocketOne.releaseReference(mockIkeSessionOne);
-        assertEquals(1, ikeSocketOne.mAliveIkeSessions.size());
-        assertTrue(isFdOpen(ikeSocketOne.getFd()));
-
-        ikeSocketTwo.releaseReference(mockIkeSessionTwo);
-        assertEquals(0, ikeSocketTwo.mAliveIkeSessions.size());
-        verifyCloseFd(ikeSocketTwo.getFd());
+        verifyGetAndCloseIkeSocketSameNetwork(
+                mIkeSocketFactory, IkeSocket.SERVER_PORT_NON_UDP_ENCAPSULATED);
     }
 
     @Test
     public void testGetAndCloseIkeUdp6SocketDifferentNetwork() throws Exception {
-        IkeSessionStateMachine mockIkeSessionOne = mock(IkeSessionStateMachine.class);
-        IkeSessionStateMachine mockIkeSessionTwo = mock(IkeSessionStateMachine.class);
-
-        Network mockNetworkOne = mock(Network.class);
-        Network mockNetworkTwo = mock(Network.class);
-
-        IkeUdp6Socket ikeSocketOne =
-                IkeUdp6Socket.getInstance(mockNetworkOne, mockIkeSessionOne, mHandler);
-        assertEquals(1, ikeSocketOne.mAliveIkeSessions.size());
-
-        IkeUdp6Socket ikeSocketTwo =
-                IkeUdp6Socket.getInstance(mockNetworkTwo, mockIkeSessionTwo, mHandler);
-        assertEquals(1, ikeSocketTwo.mAliveIkeSessions.size());
-
-        assertNotEquals(ikeSocketOne, ikeSocketTwo);
-
-        ArgumentCaptor<FileDescriptor> fdCaptorOne = ArgumentCaptor.forClass(FileDescriptor.class);
-        ArgumentCaptor<FileDescriptor> fdCaptorTwo = ArgumentCaptor.forClass(FileDescriptor.class);
-        verify(mockNetworkOne).bindSocket(fdCaptorOne.capture());
-        verify(mockNetworkTwo).bindSocket(fdCaptorTwo.capture());
-
-        FileDescriptor fdOne = fdCaptorOne.getValue();
-        FileDescriptor fdTwo = fdCaptorTwo.getValue();
-        assertNotNull(fdOne);
-        assertNotNull(fdTwo);
-        assertNotEquals(fdOne, fdTwo);
-
-        ikeSocketOne.releaseReference(mockIkeSessionOne);
-        assertEquals(0, ikeSocketOne.mAliveIkeSessions.size());
-        verifyCloseFd(ikeSocketOne.getFd());
-
-        ikeSocketTwo.releaseReference(mockIkeSessionTwo);
-        assertEquals(0, ikeSocketTwo.mAliveIkeSessions.size());
-        verifyCloseFd(ikeSocketTwo.getFd());
+        verifyGetAndCloseIkeSocketDifferentNetwork(
+                mIkeSocketFactory, IkeSocket.SERVER_PORT_NON_UDP_ENCAPSULATED);
     }
 
     @Test
     public void testReceiveIkePacket() throws Exception {
-        IkeSessionStateMachine mockIkeSession = mock(IkeSessionStateMachine.class);
-        IkeUdp6Socket ikeSocket = IkeUdp6Socket.getInstance(mMockNetwork, mockIkeSession, mHandler);
-        assertNotNull(ikeSocket);
-
-        // Set up state
-        ikeSocket.registerIke(LOCAL_SPI, mockIkeSession);
-        IkeSocket.IPacketReceiver packetReceiver = mock(IkeSocket.IPacketReceiver.class);
-        IkeUdpSocket.setPacketReceiver(packetReceiver);
-        try {
-            // Send a packet
-            byte[] pktBytes = HexDump.hexStringToByteArray(IKE_REQ_MESSAGE_HEX_STRING);
-            ikeSocket.handlePacket(pktBytes, pktBytes.length);
-
-            verify(packetReceiver).handlePacket(eq(pktBytes), any());
-
-        } finally {
-            ikeSocket.releaseReference(mockIkeSession);
-            IkeUdpSocket.setPacketReceiver(getPacketReceiver());
-        }
+        verifyIkeUdpSocketReceivePacket(mIkeSocketFactory, getPacketReceiver());
     }
 
     @Test
     public void testHandlePacket() throws Exception {
-        byte[] recvBuf = TestUtils.hexStringToByteArray(IKE_REQ_MESSAGE_HEX_STRING);
-
-        getPacketReceiver().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(REMOTE_SPI, capturedIkeHeader.ikeInitiatorSpi);
-        assertEquals(LOCAL_SPI, capturedIkeHeader.ikeResponderSpi);
+        verifyHandlePacket(
+                TestUtils.hexStringToByteArray(IKE_REQ_MESSAGE_HEX_STRING), getPacketReceiver());
     }
 }
diff --git a/tests/iketests/src/java/com/android/internal/net/ipsec/ike/IkeUdp6WithEncapPortSocketTest.java b/tests/iketests/src/java/com/android/internal/net/ipsec/ike/IkeUdp6WithEncapPortSocketTest.java
new file mode 100644
index 0000000..0c79628
--- /dev/null
+++ b/tests/iketests/src/java/com/android/internal/net/ipsec/ike/IkeUdp6WithEncapPortSocketTest.java
@@ -0,0 +1,70 @@
+/*
+ * 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 com.android.internal.net.ipsec.ike;
+
+import android.net.Network;
+import android.os.Handler;
+import android.os.test.TestLooper;
+import android.system.ErrnoException;
+
+import com.android.internal.net.TestUtils;
+
+import org.junit.Test;
+
+import java.io.IOException;
+
+public final class IkeUdp6WithEncapPortSocketTest extends IkeSocketTestBase {
+    private final TestLooper mLooper = new TestLooper();
+    private final Handler mHandler = new Handler(mLooper.getLooper());
+
+    private final IkeSocketFactory mIkeSocketFactory =
+            new IkeSocketFactory() {
+                @Override
+                public IkeSocket getIkeSocket(Network network, IkeSessionStateMachine ikeSession)
+                        throws ErrnoException, IOException {
+                    return IkeUdp6WithEncapPortSocket.getInstance(network, ikeSession, mHandler);
+                }
+            };
+
+    @Override
+    protected IkeSocket.IPacketReceiver getPacketReceiver() {
+        return new IkeUdpSocket.PacketReceiver();
+    }
+
+    @Test
+    public void testGetAndCloseIkeUdp6WithEncapPortSocketTestSameNetwork() throws Exception {
+        verifyGetAndCloseIkeSocketSameNetwork(
+                mIkeSocketFactory, IkeSocket.SERVER_PORT_UDP_ENCAPSULATED);
+    }
+
+    @Test
+    public void testGetAndCloseIkeUdp6WithEncapPortSocketTestDifferentNetwork() throws Exception {
+        verifyGetAndCloseIkeSocketDifferentNetwork(
+                mIkeSocketFactory, IkeSocket.SERVER_PORT_UDP_ENCAPSULATED);
+    }
+
+    @Test
+    public void testReceiveIkePacket() throws Exception {
+        verifyIkeUdpSocketReceivePacket(mIkeSocketFactory, getPacketReceiver());
+    }
+
+    @Test
+    public void testHandlePacket() throws Exception {
+        verifyHandlePacket(
+                TestUtils.hexStringToByteArray(IKE_REQ_MESSAGE_HEX_STRING), getPacketReceiver());
+    }
+}
diff --git a/tests/iketests/src/java/com/android/internal/net/ipsec/ike/IkeUdpEncapSocketTest.java b/tests/iketests/src/java/com/android/internal/net/ipsec/ike/IkeUdpEncapSocketTest.java
index c75c3c6..781d2aa 100644
--- a/tests/iketests/src/java/com/android/internal/net/ipsec/ike/IkeUdpEncapSocketTest.java
+++ b/tests/iketests/src/java/com/android/internal/net/ipsec/ike/IkeUdpEncapSocketTest.java
@@ -17,15 +17,11 @@
 package com.android.internal.net.ipsec.ike;
 
 import static org.junit.Assert.assertArrayEquals;
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertNotEquals;
 import static org.junit.Assert.assertNotNull;
 import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.anyInt;
 import static org.mockito.Mockito.anyObject;
 import static org.mockito.Mockito.doReturn;
-import static org.mockito.Mockito.eq;
-import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.times;
@@ -39,13 +35,13 @@
 import android.net.Network;
 import android.os.HandlerThread;
 import android.os.Looper;
+import android.os.test.TestLooper;
 import android.system.ErrnoException;
 import android.util.Log;
 
 import androidx.test.InstrumentationRegistry;
 
 import com.android.internal.net.TestUtils;
-import com.android.internal.net.ipsec.ike.message.IkeHeader;
 import com.android.internal.net.ipsec.ike.testutils.MockIpSecTestUtils;
 import com.android.server.IpSecService;
 
@@ -53,7 +49,6 @@
 import org.junit.Before;
 import org.junit.Ignore;
 import org.junit.Test;
-import org.mockito.ArgumentCaptor;
 
 import java.io.FileDescriptor;
 import java.io.IOException;
@@ -61,12 +56,28 @@
 import java.nio.ByteBuffer;
 
 public final class IkeUdpEncapSocketTest extends IkeSocketTestBase {
+    private final TestLooper mLooper = new TestLooper();
+
     private UdpEncapsulationSocket mSpyUdpEncapSocket;
 
     private UdpEncapsulationSocket mSpyDummyUdpEncapSocketOne;
     private UdpEncapsulationSocket mSpyDummyUdpEncapSocketTwo;
     private IpSecManager mSpyIpSecManager;
 
+    private final IkeSocketFactory mIkeSocketFactory =
+            new IkeSocketFactory() {
+                @Override
+                public IkeSocket getIkeSocket(Network network, IkeSessionStateMachine ikeSession)
+                        throws ErrnoException, IOException {
+                    try {
+                        return IkeUdpEncapSocket.getIkeUdpEncapSocket(
+                                network, mSpyIpSecManager, ikeSession, mLooper.getLooper());
+                    } catch (ResourceUnavailableException e) {
+                        throw new IllegalStateException(e);
+                    }
+                }
+            };
+
     @Override
     @Before
     public void setUp() throws Exception {
@@ -110,85 +121,26 @@
 
     @Test
     public void testGetAndCloseIkeUdpEncapSocketSameNetwork() throws Exception {
-        // Must be prepared here; AndroidJUnitRunner runs tests on different threads from the
-        // setUp() call. Since the new Handler() call is run in getIkeUdpEncapSocket, the Looper
-        // must be prepared here.
-        if (Looper.myLooper() == null) Looper.prepare();
-
-        IkeSessionStateMachine mockIkeSessionOne = mock(IkeSessionStateMachine.class);
-        IkeSessionStateMachine mockIkeSessionTwo = mock(IkeSessionStateMachine.class);
-
-        IkeUdpEncapSocket ikeSocketOne =
-                IkeUdpEncapSocket.getIkeUdpEncapSocket(
-                        mMockNetwork, mSpyIpSecManager, mockIkeSessionOne, Looper.myLooper());
-        assertEquals(1, ikeSocketOne.mAliveIkeSessions.size());
-
-        IkeUdpEncapSocket ikeSocketTwo =
-                IkeUdpEncapSocket.getIkeUdpEncapSocket(
-                        mMockNetwork, mSpyIpSecManager, mockIkeSessionTwo, Looper.myLooper());
-        assertEquals(2, ikeSocketTwo.mAliveIkeSessions.size());
-        assertEquals(ikeSocketOne, ikeSocketTwo);
-
+        verifyGetAndCloseIkeSocketSameNetwork(
+                mIkeSocketFactory, IkeSocket.SERVER_PORT_UDP_ENCAPSULATED);
         verify(mSpyIpSecManager).openUdpEncapsulationSocket();
-        verify(mMockNetwork).bindSocket(any(FileDescriptor.class));
-
-        ikeSocketOne.releaseReference(mockIkeSessionOne);
-        assertEquals(1, ikeSocketOne.mAliveIkeSessions.size());
-        verify(mSpyDummyUdpEncapSocketOne, never()).close();
-
-        ikeSocketTwo.releaseReference(mockIkeSessionTwo);
-        assertEquals(0, ikeSocketTwo.mAliveIkeSessions.size());
         verify(mSpyDummyUdpEncapSocketOne).close();
     }
 
     @Test
     public void testGetAndCloseIkeUdpEncapSocketDifferentNetwork() throws Exception {
-        // Must be prepared here; AndroidJUnitRunner runs tests on different threads from the
-        // setUp() call. Since the new Handler() call is run in getIkeUdpEncapSocket, the Looper
-        // must be prepared here.
-        if (Looper.myLooper() == null) Looper.prepare();
-
-        IkeSessionStateMachine mockIkeSessionOne = mock(IkeSessionStateMachine.class);
-        IkeSessionStateMachine mockIkeSessionTwo = mock(IkeSessionStateMachine.class);
-
-        Network mockNetworkOne = mock(Network.class);
-        Network mockNetworkTwo = mock(Network.class);
-
-        IkeUdpEncapSocket ikeSocketOne =
-                IkeUdpEncapSocket.getIkeUdpEncapSocket(
-                        mockNetworkOne, mSpyIpSecManager, mockIkeSessionOne, Looper.myLooper());
-        assertEquals(1, ikeSocketOne.mAliveIkeSessions.size());
-
-        IkeUdpEncapSocket ikeSocketTwo =
-                IkeUdpEncapSocket.getIkeUdpEncapSocket(
-                        mockNetworkTwo, mSpyIpSecManager, mockIkeSessionTwo, Looper.myLooper());
-        assertEquals(1, ikeSocketTwo.mAliveIkeSessions.size());
-
-        assertNotEquals(ikeSocketOne, ikeSocketTwo);
+        verifyGetAndCloseIkeSocketDifferentNetwork(
+                mIkeSocketFactory, IkeSocket.SERVER_PORT_UDP_ENCAPSULATED);
         verify(mSpyIpSecManager, times(2)).openUdpEncapsulationSocket();
-
-        ArgumentCaptor<FileDescriptor> fdCaptorOne = ArgumentCaptor.forClass(FileDescriptor.class);
-        ArgumentCaptor<FileDescriptor> fdCaptorTwo = ArgumentCaptor.forClass(FileDescriptor.class);
-        verify(mockNetworkOne).bindSocket(fdCaptorOne.capture());
-        verify(mockNetworkTwo).bindSocket(fdCaptorTwo.capture());
-
-        FileDescriptor fdOne = fdCaptorOne.getValue();
-        FileDescriptor fdTwo = fdCaptorTwo.getValue();
-        assertNotNull(fdOne);
-        assertNotNull(fdTwo);
-        assertNotEquals(fdOne, fdTwo);
-
-        ikeSocketOne.releaseReference(mockIkeSessionOne);
-        assertEquals(0, ikeSocketOne.mAliveIkeSessions.size());
         verify(mSpyDummyUdpEncapSocketOne).close();
-
-        ikeSocketTwo.releaseReference(mockIkeSessionTwo);
-        assertEquals(0, ikeSocketTwo.mAliveIkeSessions.size());
         verify(mSpyDummyUdpEncapSocketTwo).close();
     }
 
     @Ignore
     public void disableTestSendIkePacket() throws Exception {
+        // Must be prepared here; AndroidJUnitRunner runs tests on different threads
+        // from the setUp() call. Since the new Handler() call is run in
+        // getIkeUdpEncapSocket, the Looper must be prepared here.
         if (Looper.myLooper() == null) Looper.prepare();
 
         // Send IKE packet
@@ -293,16 +245,7 @@
                 TestUtils.hexStringToByteArray(
                         NON_ESP_MARKER_HEX_STRING + IKE_REQ_MESSAGE_HEX_STRING);
 
-        getPacketReceiver().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(REMOTE_SPI, capturedIkeHeader.ikeInitiatorSpi);
-        assertEquals(LOCAL_SPI, capturedIkeHeader.ikeResponderSpi);
+        verifyHandlePacket(recvBuf, getPacketReceiver());
     }
 
     @Test