Add tests for IKE UDP{4, 6} Sockets

This tests adds the IkeUdp4Socket and IkeUdp6Socket tests.

Bug: 152460200
Test: FrameworksIkeTests passing, new tests added
Change-Id: Ie43c54a22de6645ed672aa447c08ec4fb292b2ed

6 files changed, 370 insertions, 35 deletions

diff --git a/src/java/com/android/internal/net/ipsec/ike/IkeUdp4Socket.java b/src/java/com/android/internal/net/ipsec/ike/IkeUdp4Socket.java
index fc73626f..c65c27ea 100644
--- a/src/java/com/android/internal/net/ipsec/ike/IkeUdp4Socket.java
+++ b/src/java/com/android/internal/net/ipsec/ike/IkeUdp4Socket.java
@@ -28,6 +28,8 @@ import android.os.Handler;
 import android.system.ErrnoException;
 import android.system.Os;
 
+import com.android.internal.annotations.VisibleForTesting;
+
 import java.io.FileDescriptor;
 import java.io.IOException;
 import java.net.InetAddress;
@@ -50,7 +52,7 @@ public final class IkeUdp4Socket extends IkeUdpSocket {
     private static Map<Network, IkeUdp4Socket> sNetworkToUdp4SocketMap = new HashMap<>();
 
     private IkeUdp4Socket(FileDescriptor socket, Network network, Handler handler) {
-        super(socket, network, handler);
+        super(socket, network, handler == null ? new Handler() : handler);
     }
 
     /**
@@ -65,6 +67,14 @@ public final class IkeUdp4Socket extends IkeUdpSocket {
      */
     public static IkeUdp4Socket getInstance(Network network, IkeSessionStateMachine ikeSession)
             throws ErrnoException, IOException {
+        return getInstance(network, ikeSession, null);
+    }
+
+    // package protected; for testing purposes.
+    @VisibleForTesting
+    static IkeUdp4Socket getInstance(
+            Network network, IkeSessionStateMachine ikeSession, Handler handler)
+            throws ErrnoException, IOException {
         IkeUdp4Socket ikeSocket = sNetworkToUdp4SocketMap.get(network);
         if (ikeSocket == null) {
             FileDescriptor sock = Os.socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
@@ -74,7 +84,7 @@ public final class IkeUdp4Socket extends IkeUdpSocket {
             Os.fcntlInt(sock, F_SETFL, SOCK_DGRAM | SOCK_NONBLOCK);
             network.bindSocket(sock);
 
-            ikeSocket = new IkeUdp4Socket(sock, network, new Handler());
+            ikeSocket = new IkeUdp4Socket(sock, network, handler);
 
             // Create and register FileDescriptor for receiving IKE packet on current thread.
             ikeSocket.start();
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 9be5355d..f17757e0 100644
--- a/src/java/com/android/internal/net/ipsec/ike/IkeUdp6Socket.java
+++ b/src/java/com/android/internal/net/ipsec/ike/IkeUdp6Socket.java
@@ -28,6 +28,8 @@ import android.os.Handler;
 import android.system.ErrnoException;
 import android.system.Os;
 
+import com.android.internal.annotations.VisibleForTesting;
+
 import java.io.FileDescriptor;
 import java.io.IOException;
 import java.net.InetAddress;
@@ -50,7 +52,7 @@ public final class IkeUdp6Socket extends IkeUdpSocket {
     private static Map<Network, IkeUdp6Socket> sNetworkToUdp6SocketMap = new HashMap<>();
 
     private IkeUdp6Socket(FileDescriptor socket, Network network, Handler handler) {
-        super(socket, network, handler);
+        super(socket, network, handler == null ? new Handler() : handler);
     }
 
     /**
@@ -65,6 +67,14 @@ public final class IkeUdp6Socket extends IkeUdpSocket {
      */
     public static IkeUdp6Socket getInstance(Network network, IkeSessionStateMachine ikeSession)
             throws ErrnoException, IOException {
+        return getInstance(network, ikeSession, null);
+    }
+
+    // package protected; for testing purposes.
+    @VisibleForTesting
+    static IkeUdp6Socket getInstance(
+            Network network, IkeSessionStateMachine ikeSession, Handler handler)
+            throws ErrnoException, IOException {
         IkeUdp6Socket ikeSocket = sNetworkToUdp6SocketMap.get(network);
         if (ikeSocket == null) {
             FileDescriptor sock = Os.socket(AF_INET6, SOCK_DGRAM, IPPROTO_UDP);
@@ -74,7 +84,7 @@ public final class IkeUdp6Socket extends IkeUdpSocket {
             Os.fcntlInt(sock, F_SETFL, SOCK_DGRAM | SOCK_NONBLOCK);
             network.bindSocket(sock);
 
-            ikeSocket = new IkeUdp6Socket(sock, network, new Handler());
+            ikeSocket = new IkeUdp6Socket(sock, network, handler);
 
             // Create and register FileDescriptor for receiving IKE packet on current thread.
             ikeSocket.start();
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 5bece076..d9846ffa 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
@@ -94,7 +94,6 @@ public abstract class IkeSocketTestBase {
 
     @After
     public void tearDown() throws Exception {
-        IkeUdpEncapSocket.setPacketReceiver(getPacketReceiver());
         Os.close(mDummyRemoteServerFd);
     }
 
@@ -107,6 +106,15 @@ public abstract class IkeSocketTestBase {
         return sock;
     }
 
+    protected boolean isFdOpen(FileDescriptor fd) {
+        try {
+            Os.getsockname(fd);
+            return true;
+        } catch (ErrnoException ignored) {
+            return false;
+        }
+    }
+
     protected void verifyCloseFd(FileDescriptor fd) {
         try {
             Os.sendto(
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
new file mode 100644
index 00000000..ee058854
--- /dev/null
+++ b/tests/iketests/src/java/com/android/internal/net/ipsec/ike/IkeUdp4SocketTest.java
@@ -0,0 +1,152 @@
+/*
+ * 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 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 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;
+
+// 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());
+
+    @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());
+    }
+
+    @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());
+    }
+
+    @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());
+        }
+    }
+
+    @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);
+    }
+}
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
new file mode 100644
index 00000000..7c6e30d1
--- /dev/null
+++ b/tests/iketests/src/java/com/android/internal/net/ipsec/ike/IkeUdp6SocketTest.java
@@ -0,0 +1,152 @@
+/*
+ * 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 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 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;
+
+// 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());
+
+    @Override
+    protected IkeSocket.IPacketReceiver getPacketReceiver() {
+        return new IkeUdpSocket.PacketReceiver();
+    }
+
+    @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());
+    }
+
+    @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());
+    }
+
+    @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());
+        }
+    }
+
+    @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);
+    }
+}
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 fab2b475..8c0bf4d5 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
@@ -248,36 +248,39 @@ public final class IkeUdpEncapSocketTest extends IkeSocketTestBase {
         TestCountDownLatch receiveLatch = new TestCountDownLatch();
         DummyPacketReceiver packetReceiver = new DummyPacketReceiver(receiveLatch);
         IkeUdpEncapSocket.setPacketReceiver(packetReceiver);
-
-        // Send first packet.
-        sendToIkeUdpEncapSocket(mDummyRemoteServerFd, mDataOne, IPV4_LOOPBACK);
-        receiveLatch.await();
-
-        assertEquals(1, ikeSocket.numPacketsReceived());
-        assertArrayEquals(mDataOne, packetReceiver.mReceivedData);
-
-        // Send second packet.
-        sendToIkeUdpEncapSocket(mDummyRemoteServerFd, mDataTwo, IPV4_LOOPBACK);
-        receiveLatch.await();
-
-        assertEquals(2, ikeSocket.numPacketsReceived());
-        assertArrayEquals(mDataTwo, packetReceiver.mReceivedData);
-
-        // Close IkeUdpEncapSocket.
-        TestCountDownLatch closeLatch = new TestCountDownLatch();
-        ikeSocket
-                .getHandler()
-                .post(
-                        () -> {
-                            ikeSocket.releaseReference(mMockIkeSessionStateMachine);
-                            closeLatch.countDown();
-                        });
-        closeLatch.await();
-
-        verify(mSpyUdpEncapSocket).close();
-        verifyCloseFd(mSpyUdpEncapSocket.getFileDescriptor());
-
-        mIkeThread.quitSafely();
+        try {
+            // Send first packet.
+            sendToIkeUdpEncapSocket(mDummyRemoteServerFd, mDataOne, IPV4_LOOPBACK);
+            receiveLatch.await();
+
+            assertEquals(1, ikeSocket.numPacketsReceived());
+            assertArrayEquals(mDataOne, packetReceiver.mReceivedData);
+
+            // Send second packet.
+            sendToIkeUdpEncapSocket(mDummyRemoteServerFd, mDataTwo, IPV4_LOOPBACK);
+            receiveLatch.await();
+
+            assertEquals(2, ikeSocket.numPacketsReceived());
+            assertArrayEquals(mDataTwo, packetReceiver.mReceivedData);
+
+            // Close IkeUdpEncapSocket.
+            TestCountDownLatch closeLatch = new TestCountDownLatch();
+            ikeSocket
+                    .getHandler()
+                    .post(
+                            () -> {
+                                ikeSocket.releaseReference(mMockIkeSessionStateMachine);
+                                closeLatch.countDown();
+                            });
+            closeLatch.await();
+
+            verify(mSpyUdpEncapSocket).close();
+            verifyCloseFd(mSpyUdpEncapSocket.getFileDescriptor());
+
+            mIkeThread.quitSafely();
+        } finally {
+            IkeUdpEncapSocket.setPacketReceiver(getPacketReceiver());
+        }
     }
 
     @Test