Prevent CTS from hanging if no UDP packet was received

Previously, CTS tests could block forever on attempting to read a packet
from a UDP socket.

This patch modifies the receive methods of the NativeUdpSocket wrapper,
allowing it to gracefully fail if UDP packets are not received.
This patch also updates the IKE over UDP encap socket test to use the
NativeUdpSocket wrapper, preventing the tests from hanging in the same
fashion.

Bug: 70938121
Test: Ran updated tests on walleye + marlin
Merged-In: I390e4585e85647eb8555e706d44b76f95dec931f
Change-Id: I390e4585e85647eb8555e706d44b76f95dec931f
(cherry picked from commit e2fcb9d64ba10f07a69a6b96a7f7e1465856172a)
diff --git a/tests/cts/net/src/android/net/cts/IpSecBaseTest.java b/tests/cts/net/src/android/net/cts/IpSecBaseTest.java
index 57cfe2a..7132ecf 100644
--- a/tests/cts/net/src/android/net/cts/IpSecBaseTest.java
+++ b/tests/cts/net/src/android/net/cts/IpSecBaseTest.java
@@ -25,6 +25,7 @@
 import android.system.Os;
 import android.system.OsConstants;
 import android.test.AndroidTestCase;
+import android.util.Log;
 
 import java.io.FileDescriptor;
 import java.io.IOException;
@@ -38,6 +39,7 @@
 import java.net.Socket;
 import java.net.SocketException;
 import java.util.Arrays;
+import java.util.concurrent.atomic.AtomicInteger;
 
 public class IpSecBaseTest extends AndroidTestCase {
 
@@ -51,6 +53,7 @@
 
     protected static final byte[] TEST_DATA = "Best test data ever!".getBytes();
     protected static final int DATA_BUFFER_LEN = 4096;
+    protected static final int SOCK_TIMEOUT = 500;
 
     private static final byte[] KEY_DATA = {
         0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
@@ -124,9 +127,27 @@
         @Override
         public byte[] receive() throws Exception {
             byte[] in = new byte[DATA_BUFFER_LEN];
-            int bytesRead = Os.read(mFd, in, 0, DATA_BUFFER_LEN);
+            AtomicInteger bytesRead = new AtomicInteger(-1);
 
-            return Arrays.copyOfRange(in, 0, bytesRead);
+            Thread readSockThread = new Thread(() -> {
+                long startTime = System.currentTimeMillis();
+                while (bytesRead.get() < 0 && System.currentTimeMillis() < startTime + SOCK_TIMEOUT) {
+                    try {
+                        bytesRead.set(Os.recvfrom(mFd, in, 0, DATA_BUFFER_LEN, 0, null));
+                    } catch (Exception e) {
+                        Log.e(TAG, "Error encountered reading from socket", e);
+                    }
+                }
+            });
+
+            readSockThread.start();
+            readSockThread.join(SOCK_TIMEOUT);
+
+            if (bytesRead.get() < 0) {
+                throw new IOException("No data received from socket");
+            }
+
+            return Arrays.copyOfRange(in, 0, bytesRead.get());
         }
 
         @Override
@@ -174,7 +195,7 @@
         public JavaUdpSocket(InetAddress localAddr) {
             try {
                 mSocket = new DatagramSocket(0, localAddr);
-                mSocket.setSoTimeout(500);
+                mSocket.setSoTimeout(SOCK_TIMEOUT);
             } catch (SocketException e) {
                 // Fail loudly if we can't set up sockets properly. And without the timeout, we
                 // could easily end up in an endless wait.
@@ -227,7 +248,7 @@
         public JavaTcpSocket(Socket socket) {
             mSocket = socket;
             try {
-                mSocket.setSoTimeout(500);
+                mSocket.setSoTimeout(SOCK_TIMEOUT);
             } catch (SocketException e) {
                 // Fail loudly if we can't set up sockets properly. And without the timeout, we
                 // could easily end up in an endless wait.
@@ -279,7 +300,7 @@
         }
     }
 
-    private static void applyTransformBidirectionally(
+    protected static void applyTransformBidirectionally(
             IpSecManager ism, IpSecTransform transform, GenericSocket socket) throws Exception {
         for (int direction : DIRECTIONS) {
             socket.applyTransportModeTransform(ism, direction, transform);
diff --git a/tests/cts/net/src/android/net/cts/IpSecManagerTest.java b/tests/cts/net/src/android/net/cts/IpSecManagerTest.java
index 9be201b..95d91a2 100644
--- a/tests/cts/net/src/android/net/cts/IpSecManagerTest.java
+++ b/tests/cts/net/src/android/net/cts/IpSecManagerTest.java
@@ -576,34 +576,68 @@
         }
     }
 
-    public void testIkeOverUdpEncapSocket() throws Exception {
-        // IPv6 not supported for UDP-encap-ESP
-        InetAddress local = InetAddress.getByName(IPV4_LOOPBACK);
+    private void checkIkePacket(
+            NativeUdpSocket wrappedEncapSocket, InetAddress localAddr) throws Exception {
         StatsChecker.initStatsChecker();
 
-        try (IpSecManager.UdpEncapsulationSocket encapSocket = mISM.openUdpEncapsulationSocket()) {
-            int localPort = getPort(encapSocket.getFileDescriptor());
+        try (NativeUdpSocket remoteSocket = new NativeUdpSocket(getBoundUdpSocket(localAddr))) {
 
-            // Append ESP header - 4 bytes of SPI, 4 bytes of seq number
+            // Append IKE/ESP header - 4 bytes of SPI, 4 bytes of seq number, all zeroed out
+            // If the first four bytes are zero, assume non-ESP (IKE traffic)
             byte[] dataWithEspHeader = new byte[TEST_DATA.length + 8];
             System.arraycopy(TEST_DATA, 0, dataWithEspHeader, 8, TEST_DATA.length);
 
-            byte[] in = new byte[dataWithEspHeader.length];
-            Os.sendto(
-                    encapSocket.getFileDescriptor(),
-                    dataWithEspHeader,
-                    0,
-                    dataWithEspHeader.length,
-                    0,
-                    local,
-                    localPort);
-            Os.read(encapSocket.getFileDescriptor(), in, 0, in.length);
+            // Send the IKE packet from remoteSocket to wrappedEncapSocket. Since IKE packets
+            // are multiplexed over the socket, we expect them to appear on the encap socket
+            // (as opposed to being decrypted and received on the non-encap socket)
+            remoteSocket.sendTo(dataWithEspHeader, localAddr, wrappedEncapSocket.getPort());
+            byte[] in = wrappedEncapSocket.receive();
             assertArrayEquals("Encapsulated data did not match.", dataWithEspHeader, in);
 
-            int ipHdrLen = local instanceof Inet6Address ? IP6_HDRLEN : IP4_HDRLEN;
-            int expectedPacketSize = dataWithEspHeader.length + UDP_HDRLEN + ipHdrLen;
-            StatsChecker.assertUidStatsDelta(expectedPacketSize, 1, expectedPacketSize, 1);
-            StatsChecker.assertIfaceStatsDelta(expectedPacketSize, 1, expectedPacketSize, 1);
+            // Also test that the IKE socket can send data out.
+            wrappedEncapSocket.sendTo(dataWithEspHeader, localAddr, remoteSocket.getPort());
+            in = remoteSocket.receive();
+            assertArrayEquals("Encapsulated data did not match.", dataWithEspHeader, in);
+
+            // Calculate expected packet sizes. Always use IPv4 header, since our kernels only
+            // guarantee support of UDP encap on IPv4.
+            int expectedNumPkts = 2;
+            int expectedPacketSize =
+                    expectedNumPkts * (dataWithEspHeader.length + UDP_HDRLEN + IP4_HDRLEN);
+
+            StatsChecker.waitForNumPackets(expectedNumPkts);
+            StatsChecker.assertUidStatsDelta(
+                    expectedPacketSize, expectedNumPkts, expectedPacketSize, expectedNumPkts);
+            StatsChecker.assertIfaceStatsDelta(
+                    expectedPacketSize, expectedNumPkts, expectedPacketSize, expectedNumPkts);
+        }
+    }
+
+    public void testIkeOverUdpEncapSocket() throws Exception {
+        // IPv6 not supported for UDP-encap-ESP
+        InetAddress local = InetAddress.getByName(IPV4_LOOPBACK);
+        try (IpSecManager.UdpEncapsulationSocket encapSocket = mISM.openUdpEncapsulationSocket()) {
+            NativeUdpSocket wrappedEncapSocket =
+                    new NativeUdpSocket(encapSocket.getFileDescriptor());
+            checkIkePacket(wrappedEncapSocket, local);
+
+            // Now try with a transform applied to a socket using this Encap socket
+            IpSecAlgorithm crypt = new IpSecAlgorithm(IpSecAlgorithm.CRYPT_AES_CBC, CRYPT_KEY);
+            IpSecAlgorithm auth = new IpSecAlgorithm(IpSecAlgorithm.AUTH_HMAC_MD5, getKey(128), 96);
+
+            try (IpSecManager.SecurityParameterIndex spi =
+                            mISM.allocateSecurityParameterIndex(local);
+                    IpSecTransform transform =
+                            new IpSecTransform.Builder(mContext)
+                                    .setEncryption(crypt)
+                                    .setAuthentication(auth)
+                                    .setIpv4Encapsulation(encapSocket, encapSocket.getPort())
+                                    .buildTransportModeTransform(local, spi);
+                    JavaUdpSocket localSocket = new JavaUdpSocket(local)) {
+                applyTransformBidirectionally(mISM, transform, localSocket);
+
+                checkIkePacket(wrappedEncapSocket, local);
+            }
         }
     }
 
@@ -1096,131 +1130,4 @@
             assertTrue("Returned invalid port", encapSocket.getPort() != 0);
         }
     }
-
-    public void testUdpEncapsulation() throws Exception {
-        InetAddress local = InetAddress.getByName(IPV4_LOOPBACK);
-
-        // TODO: Refactor to make this more representative of a normal application use case. (use
-        // separate sockets for inbound and outbound)
-        // Create SPIs, UDP encap socket
-        try (IpSecManager.UdpEncapsulationSocket encapSocket = mISM.openUdpEncapsulationSocket();
-                IpSecManager.SecurityParameterIndex spi =
-                        mISM.allocateSecurityParameterIndex(local);
-                IpSecTransform transform =
-                        buildIpSecTransform(mContext, spi, encapSocket, local)) {
-
-            // Create user socket, apply transform to it
-            FileDescriptor udpSocket = null;
-            try {
-                udpSocket = getBoundUdpSocket(local);
-                int port = getPort(udpSocket);
-
-                mISM.applyTransportModeTransform(
-                        udpSocket, IpSecManager.DIRECTION_IN, transform);
-                mISM.applyTransportModeTransform(
-                        udpSocket, IpSecManager.DIRECTION_OUT, transform);
-
-                // Send an ESP packet from this socket to itself. Since the inbound and
-                // outbound transforms match, we should receive the data we sent.
-                byte[] data = new String("IPSec UDP-encap-ESP test data").getBytes("UTF-8");
-                Os.sendto(udpSocket, data, 0, data.length, 0, local, port);
-                byte[] in = new byte[data.length];
-                Os.read(udpSocket, in, 0, in.length);
-                assertTrue("Encapsulated data did not match.", Arrays.equals(data, in));
-
-                // Send an IKE packet from this socket to itself. IKE packets (SPI of 0)
-                // are not transformed in any way, and should be sent in the clear
-                // We expect this to work too (no inbound transforms)
-                final byte[] header = new byte[] {0, 0, 0, 0};
-                final String message = "Sample IKE Packet";
-                data = (new String(header) + message).getBytes("UTF-8");
-                Os.sendto(
-                        encapSocket.getFileDescriptor(),
-                        data,
-                        0,
-                        data.length,
-                        0,
-                        local,
-                        encapSocket.getPort());
-                in = new byte[data.length];
-                Os.read(encapSocket.getFileDescriptor(), in, 0, in.length);
-                assertTrue(
-                        "Encap socket was unable to send/receive IKE data",
-                        Arrays.equals(data, in));
-
-                mISM.removeTransportModeTransforms(udpSocket);
-            } finally {
-                if (udpSocket != null) {
-                    Os.close(udpSocket);
-                }
-            }
-        }
-    }
-
-    public void testIke() throws Exception {
-        InetAddress localAddr = InetAddress.getByName(IPV4_LOOPBACK);
-
-        // TODO: Refactor to make this more representative of a normal application use case. (use
-        // separate sockets for inbound and outbound)
-        try (IpSecManager.UdpEncapsulationSocket encapSocket = mISM.openUdpEncapsulationSocket();
-                IpSecManager.SecurityParameterIndex spi =
-                        mISM.allocateSecurityParameterIndex(localAddr);
-                IpSecTransform transform =
-                        buildIpSecTransform(mContext, spi, encapSocket, localAddr)) {
-
-            // Create user socket, apply transform to it
-            FileDescriptor sock = null;
-
-            try {
-                sock = getBoundUdpSocket(localAddr);
-                int port = getPort(sock);
-
-                mISM.applyTransportModeTransform(sock, IpSecManager.DIRECTION_IN, transform);
-                mISM.applyTransportModeTransform(sock, IpSecManager.DIRECTION_OUT, transform);
-
-                // TODO: Find a way to set a timeout on the socket, and assert the ESP packet
-                // doesn't make it through. Setting sockopts currently throws EPERM (possibly
-                // because it is owned by a different UID).
-
-                // Send ESP packet from our socket to the encap socket. The SPIs do not
-                // match, and we should expect this packet to be dropped.
-                byte[] header = new byte[] {1, 1, 1, 1};
-                String message = "Sample ESP Packet";
-                byte[] data = (new String(header) + message).getBytes("UTF-8");
-                Os.sendto(sock, data, 0, data.length, 0, localAddr, encapSocket.getPort());
-
-                // Send IKE packet from the encap socket to itself. Since IKE is not
-                // transformed in any way, this should succeed.
-                header = new byte[] {0, 0, 0, 0};
-                message = "Sample IKE Packet";
-                data = (new String(header) + message).getBytes("UTF-8");
-                Os.sendto(
-                        encapSocket.getFileDescriptor(),
-                        data,
-                        0,
-                        data.length,
-                        0,
-                        localAddr,
-                        encapSocket.getPort());
-
-                // ESP data should be dropped, due to different input SPI (as opposed to being
-                // readable from the encapSocket)
-                // Thus, only IKE data should be received from the socket.
-                // If the first four bytes are zero, assume non-ESP (IKE) traffic.
-                // Expect an nulled out SPI just as we sent out, without being modified.
-                byte[] in = new byte[4];
-                in[0] = 1; // Make sure the array has to be overwritten to pass
-                Os.read(encapSocket.getFileDescriptor(), in, 0, in.length);
-                assertTrue(
-                        "Encap socket received UDP-encap-ESP data despite invalid SPIs",
-                        Arrays.equals(header, in));
-
-                mISM.removeTransportModeTransforms(sock);
-            } finally {
-                if (sock != null) {
-                    Os.close(sock);
-                }
-            }
-        }
-    }
 }