fix: do not coalesce udp payloads

Author: JeroenRobben

app/src/main/java/tech/httptoolkit/android/vpn/Session.java

@@ -32,6 +32,7 @@
 import java.nio.ByteBuffer;
 import java.nio.channels.SelectionKey;
 import java.nio.channels.spi.AbstractSelectableChannel;
+import java.util.ArrayDeque;
 
 /**
  * store information about a socket connection from a VPN client.
@@ -73,8 +74,12 @@ public class Session {
 	//receiving buffer for storing data from remote host
 	private final ByteArrayOutputStream receivingStream;
 
-	//sending buffer for storing data from vpn client to be send to destination host
+	//sending buffer for storing data from vpn client to be send to destination host (TCP only)
 	private final ByteArrayOutputStream sendingStream;
+
+	//queue of discrete datagrams to be sent to the destination host (UDP only). UDP must
+	//preserve datagram boundaries, so unlike TCP it cannot use a flat byte stream.
+	private final ArrayDeque<byte[]> sendingDatagrams = new ArrayDeque<>();
 
 	private boolean hasReceivedLastSegment = false;
 
@@ -171,13 +176,21 @@ public boolean hasReceivedData(){
 	}
 
 	/**
-	 * set data to be sent to destination server
+	 * set data to be sent to destination server.
+	 * For UDP each call is queued as a discrete datagram
+	 * For TCP the bytes are appended to the send stream.
 	 * @param data Data to be sent
-	 * @return boolean Success or not
+	 * @return int number of bytes accepted
 	 */
 	public synchronized int setSendingData(ByteBuffer data) {
 		final int remaining = data.remaining();
-		sendingStream.write(data.array(), data.position(), data.remaining());
+		if (protocol == SessionProtocol.UDP) {
+			byte[] datagram = new byte[remaining];
+			System.arraycopy(data.array(), data.position(), datagram, 0, remaining);
+			sendingDatagrams.addLast(datagram);
+		} else {
+			sendingStream.write(data.array(), data.position(), remaining);
+		}
 		return remaining;
 	}
 
@@ -186,20 +199,38 @@ int getSendingDataSize(){
 	}
 
 	/**
-	 * dequeue data for sending to server
+	 * dequeue all stream data for sending to the server (TCP).
 	 * @return byte[]
 	 */
 	public synchronized byte[] getSendingData(){
 		byte[] data = sendingStream.toByteArray();
 		sendingStream.reset();
 		return data;
 	}
+
+	/**
+	 * dequeue the next datagram for sending to the server (UDP), or null if none remain.
+	 * @return byte[]
+	 */
+	public synchronized byte[] pollSendingDatagram(){
+		return sendingDatagrams.pollFirst();
+	}
+
+	/**
+	 * return a datagram to the head of the queue when it could not be written yet (UDP).
+	 */
+	public synchronized void requeueSendingDatagram(byte[] datagram){
+		sendingDatagrams.addFirst(datagram);
+	}
+
 	/**
 	 * buffer contains data for sending to destination server
 	 * @return boolean
 	 */
-	public boolean hasDataToSend(){
-		return sendingStream.size() > 0;
+	public synchronized boolean hasDataToSend(){
+		return protocol == SessionProtocol.UDP
+				? !sendingDatagrams.isEmpty()
+				: sendingStream.size() > 0;
 	}
 
 	public SessionProtocol getProtocol() {

app/src/main/java/tech/httptoolkit/android/vpn/socket/SocketChannelWriter.java

@@ -80,7 +80,7 @@ public long write(@NonNull Session session) {
 	private long writeUDP(Session session) {
 		long amountBytes = 0;
 		try {
-			amountBytes = writePendingData(session);
+			amountBytes = writePendingUDPData(session);
 			Date dt = new Date();
 			session.connectionStartTime = dt.getTime();
 		}catch(NotYetConnectedException ex2){
@@ -117,21 +117,20 @@ private long writeTCP(Session session) {
 		return amountBytes;
 	}
 
+	/** TCP: a byte stream, so buffered bytes are concatenated and written as-is. */
 	private long writePendingData(Session session) throws IOException {
 		if (!session.hasDataToSend()) return 0;
 
 		long totalBytesWritten = 0;
-		AbstractSelectableChannel channel = session.getChannel();
+		SocketChannel channel = (SocketChannel) session.getChannel();
 
 		byte[] data = session.getSendingData();
 		ByteBuffer buffer = ByteBuffer.allocate(data.length);
 		buffer.put(data);
 		buffer.flip();
 
 		while (buffer.hasRemaining()) {
-			int bytesWritten = channel instanceof SocketChannel
-				? ((SocketChannel) channel).write(buffer)
-				: ((DatagramChannel) channel).write(buffer);
+			int bytesWritten = channel.write(buffer);
 
 			if (bytesWritten == 0) {
 				break;
@@ -149,7 +148,7 @@ private long writePendingData(Session session) throws IOException {
 			// Subscribe to WRITE events, so we know when this is ready to resume.
 			session.subscribeKey(SelectionKey.OP_WRITE);
 		} else {
-			// All done, all good -> wait until the next TCP PSH / UDP packet
+			// All done, all good -> wait until the next TCP PSH packet
 			session.setDataForSendingReady(false);
 
 			// We don't need to know about WRITE events any more, we've written all our data.
@@ -158,4 +157,39 @@ private long writePendingData(Session session) throws IOException {
 		}
 		return totalBytesWritten;
 	}
+
+	/**
+	 * UDP: a datagram protocol, so each queued datagram must be written with its own
+	 * channel.write() to preserve message boundaries. We send one datagram per write cycle
+	 * and resubscribe to OP_WRITE while more remain, mirroring the TCP backpressure pattern.
+	 */
+	private long writePendingUDPData(Session session) throws IOException {
+		byte[] datagram = session.pollSendingDatagram();
+		if (datagram == null) {
+			session.setDataForSendingReady(false);
+			session.unsubscribeKey(SelectionKey.OP_WRITE);
+			return 0;
+		}
+
+		DatagramChannel channel = (DatagramChannel) session.getChannel();
+		// A connected non-blocking DatagramChannel writes the whole datagram or nothing
+		// (0 when the send buffer is full); it never sends a partial datagram.
+		int bytesWritten = channel.write(ByteBuffer.wrap(datagram));
+
+		if (bytesWritten == 0) {
+			// Not ready yet: put the datagram back and resume on the next OP_WRITE.
+			session.requeueSendingDatagram(datagram);
+			session.subscribeKey(SelectionKey.OP_WRITE);
+			return 0;
+		}
+
+		if (session.hasDataToSend()) {
+			// More datagrams queued -> come back for the next one.
+			session.subscribeKey(SelectionKey.OP_WRITE);
+		} else {
+			session.setDataForSendingReady(false);
+			session.unsubscribeKey(SelectionKey.OP_WRITE);
+		}
+		return bytesWritten;
+	}
 }

app/src/test/java/tech/httptoolkit/android/vpn/ConnectionTrackingForwardingTest.kt

@@ -0,0 +1,219 @@
+package tech.httptoolkit.android.vpn
+
+import android.app.Application
+import com.google.common.truth.Truth.assertThat
+import org.junit.After
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.robolectric.RobolectricTestRunner
+import org.robolectric.annotation.Config
+import tech.httptoolkit.android.vpn.transport.ip.IPAddress
+import tech.httptoolkit.android.vpn.transport.ip.IPHeader
+import tech.httptoolkit.android.vpn.transport.tcp.TCPHeader
+import java.net.DatagramPacket
+import java.net.DatagramSocket
+import java.net.InetAddress
+import java.net.ServerSocket
+import java.net.Socket
+import java.util.concurrent.Executors
+import java.util.concurrent.TimeUnit
+
+/**
+ * End-to-end correctness checks on the engine's connection tracking, through the real
+ * forwarding pipeline against loopback peers. Exercises both directions:
+ *
+ *  - egress (packet intercepted from a local app): the engine creates/reuses the right
+ *    session per 5-tuple, and demultiplexes concurrent connections without crossing streams;
+ *  - ingress (data coming back from the LAN peer): each reply is routed back to the exact
+ *    client connection that originated it, and traffic for an unknown connection is rejected.
+ */
+@RunWith(RobolectricTestRunner::class)
+@Config(sdk = [34], application = Application::class)
+class ConnectionTrackingForwardingTest {
+
+    private lateinit var harness: ForwardingTestHarness
+
+    private val clientIp = "10.0.0.2"
+    private val peerIp = "127.0.0.1"
+
+    @Before
+    fun setUp() {
+        harness = ForwardingTestHarness()
+    }
+
+    @After
+    fun tearDown() {
+        harness.close()
+    }
+
+    // --- UDP -----------------------------------------------------------------
+
+    @Test
+    fun `concurrent udp connections demultiplex replies back to the originating client`() {
+        val peer = DatagramSocket(0, InetAddress.getByName(peerIp)).apply { soTimeout = 3000 }
+        val peerPort = peer.localPort
+        try {
+            // Two client connections to the same peer, distinguished only by source port.
+            harness.feed(TestPackets.udpPacket(clientIp, 40001, peerIp, peerPort, "one".toByteArray()))
+            harness.feed(TestPackets.udpPacket(clientIp, 40002, peerIp, peerPort, "two".toByteArray()))
+
+            // The peer sees two distinct source sockets; echo each payload straight back.
+            repeat(2) {
+                val rx = DatagramPacket(ByteArray(64), 64)
+                peer.receive(rx)
+                peer.send(DatagramPacket(rx.data, rx.length, rx.socketAddress))
+            }
+
+            // Collect both TUN replies, then assert each landed on the correct client port:
+            // "one" must return to 40001 and "two" to 40002 (no cross-talk).
+            val byClientPort = buildMap {
+                repeat(2) {
+                    val (_, udp, payload) = harness.parseUdp(harness.awaitTunPacket())
+                    put(udp.destinationPort, String(payload))
+                }
+            }
+            assertThat(byClientPort).containsExactly(40001, "one", 40002, "two")
+        } finally {
+            peer.close()
+        }
+    }
+
+    @Test
+    fun `repeated udp datagrams reuse one connection and accumulate egress`() {
+        val peer = DatagramSocket(0, InetAddress.getByName(peerIp)).apply { soTimeout = 3000 }
+        val peerPort = peer.localPort
+        try {
+            harness.feed(TestPackets.udpPacket(clientIp, 40003, peerIp, peerPort, "p1".toByteArray()))
+            harness.feed(TestPackets.udpPacket(clientIp, 40003, peerIp, peerPort, "p2".toByteArray()))
+
+            // Both datagrams reach the peer over the same upstream socket: the second reused
+            // the connection rather than opening a new one.
+            val first = DatagramPacket(ByteArray(64), 64).also { peer.receive(it) }
+            val second = DatagramPacket(ByteArray(64), 64).also { peer.receive(it) }
+            assertThat(
+                setOf(String(first.data, 0, first.length), String(second.data, 0, second.length))
+            ).containsExactly("p1", "p2")
+            assertThat(second.socketAddress).isEqualTo(first.socketAddress)
+
+            // Exactly one session/flow was created; egress counters accumulate across both.
+            val session = harness.await {
+                harness.sessionByKey(udpKey(40003, peerPort))
+            }
+            assertThat(harness.flowDao.countNotSyncedFlows()).isEqualTo(1)
+            harness.await { session.flow.takeIf { it.packetCountEgress >= 2 } }
+        } finally {
+            peer.close()
+        }
+    }
+
+    @Test
+    fun `back-to-back udp datagrams preserve message boundaries`() {
+        // Regression test for datagram coalescing: two datagrams sent on the same connection
+        // before the writer drains must NOT be merged into one upstream datagram.
+        val peer = DatagramSocket(0, InetAddress.getByName(peerIp)).apply { soTimeout = 3000 }
+        val peerPort = peer.localPort
+        try {
+            harness.feed(TestPackets.udpPacket(clientIp, 40004, peerIp, peerPort, "AAAA".toByteArray()))
+            harness.feed(TestPackets.udpPacket(clientIp, 40004, peerIp, peerPort, "BBBB".toByteArray()))
+
+            // The peer must receive two distinct 4-byte datagrams, not one merged "AAAABBBB".
+            val first = DatagramPacket(ByteArray(64), 64).also { peer.receive(it) }
+            val second = DatagramPacket(ByteArray(64), 64).also { peer.receive(it) }
+            assertThat(first.length).isEqualTo(4)
+            assertThat(second.length).isEqualTo(4)
+            assertThat(listOf(String(first.data, 0, 4), String(second.data, 0, 4)))
+                .containsExactly("AAAA", "BBBB").inOrder()
+        } finally {
+            peer.close()
+        }
+    }
+
+    // --- TCP -----------------------------------------------------------------
+
+    @Test
+    fun `concurrent tcp connections are tracked independently`() {
+        val server = ServerSocket(0, 50, InetAddress.getByName(peerIp))
+        val peerPort = server.localPort
+        val executor = Executors.newFixedThreadPool(2)
+        val isn1 = 1000L
+        val isn2 = 5000L
+        try {
+            val accept1 = executor.submit<Socket> { server.accept() }
+            val accept2 = executor.submit<Socket> { server.accept() }
+
+            harness.feed(syn(40001, peerPort, isn1))
+            harness.feed(syn(40002, peerPort, isn2))
+
+            // Each SYN-ACK must be demultiplexed to its own client port and acknowledge that
+            // connection's ISN — proving the two handshakes are not conflated.
+            val synAcks = buildMap<Int, Pair<IPHeader, TCPHeader>> {
+                repeat(2) {
+                    val pkt = harness.awaitTunPacketMatching {
+                        val (_, tcp) = harness.parseTcp(it); tcp.isSYN && tcp.isACK
+                    }
+                    val parsed = harness.parseTcp(pkt)
+                    put(parsed.second.destinationPort, parsed)
+                }
+            }
+
+            assertThat(synAcks.keys).containsExactly(40001, 40002)
+            assertThat(synAcks.getValue(40001).second.ackNumber).isEqualTo(isn1 + 1)
+            assertThat(synAcks.getValue(40002).second.ackNumber).isEqualTo(isn2 + 1)
+            assertThat(synAcks.getValue(40001).first.destinationIP.toString()).isEqualTo(clientIp)
+            assertThat(synAcks.getValue(40002).first.destinationIP.toString()).isEqualTo(clientIp)
+
+            // Two distinct sessions and two recorded flows.
+            assertThat(harness.sessionByKey(tcpKey(40001, peerPort))).isNotNull()
+            assertThat(harness.sessionByKey(tcpKey(40002, peerPort))).isNotNull()
+            assertThat(harness.flowDao.countNotSyncedFlows()).isEqualTo(2)
+
+            accept1.get(3, TimeUnit.SECONDS)
+            accept2.get(3, TimeUnit.SECONDS)
+        } finally {
+            executor.shutdownNow()
+            server.close()
+        }
+    }
+
+    @Test
+    fun `tcp data for an unknown connection is rejected with RST and creates no session`() {
+        // An ACK with payload but no preceding SYN has no tracked connection: the engine must
+        // reject it with a RST rather than silently adopting it or crashing.
+        harness.feed(
+            TestPackets.tcpPacket(
+                clientIp, 40009, peerIp, 9,
+                seq = 42, ack = 99, flags = TestPackets.ACK, payload = "junk".toByteArray(),
+            )
+        )
+
+        val rst = harness.awaitTunPacketMatching { harness.parseTcp(it).second.isRST }
+        val (ip, tcp) = harness.parseTcp(rst)
+        assertThat(tcp.isRST).isTrue()
+        assertThat(ip.destinationIP.toString()).isEqualTo(clientIp)
+        assertThat(tcp.destinationPort).isEqualTo(40009)
+
+        assertThat(harness.sessionByKey(tcpKey(40009, 9))).isNull()
+        assertThat(harness.flowDao.countNotSyncedFlows()).isEqualTo(0)
+    }
+
+    // --- helpers -------------------------------------------------------------
+
+    private fun syn(clientPort: Int, peerPort: Int, isn: Long): ByteArray =
+        TestPackets.tcpPacket(
+            clientIp, clientPort, peerIp, peerPort,
+            seq = isn, ack = 0, flags = TestPackets.SYN, mss = 1460,
+        )
+
+    private fun udpKey(clientPort: Int, peerPort: Int): String = Session.getSessionKey(
+        SessionProtocol.UDP,
+        IPAddress(TestPackets.ip(peerIp)), peerPort,
+        IPAddress(TestPackets.ip(clientIp)), clientPort,
+    )
+
+    private fun tcpKey(clientPort: Int, peerPort: Int): String = Session.getSessionKey(
+        SessionProtocol.TCP,
+        IPAddress(TestPackets.ip(peerIp)), peerPort,
+        IPAddress(TestPackets.ip(clientIp)), clientPort,
+    )
+}