Big VPN redesign for thread-safety, performance & reliability

Previously, we had one thread to read from the VPN, one reader & writer thread
per VPN session when required, from a pool, and one NIO thread that orchestrated
the threadpool. This had a lot of races: data arrived in the session during
writing, writes couldn't safely set the NIO subscriptions to handle write retries,
and it was easy to fall into useless hot NIO loops or miss events entirely.

In practice, this was easiest demoed with any speed test app: under load, the
connection would lock up and lose data, dramatically dropping connection perf,
and it'd also probably burn a bunch of CPU en route.

Now, we have one thread that handles VPN reads, one that handles VPN writes
(async, with reliable retries), and one that handles all upstream reads &
writes with non-blocking NIO.

Author: pimterry

app/src/main/java/tech/httptoolkit/android/ProxyVpnRunnable.kt

@@ -3,7 +3,7 @@ package tech.httptoolkit.android
 import android.os.ParcelFileDescriptor
 import android.util.Log
 import android.util.SparseArray
-import tech.httptoolkit.android.vpn.ClientPacketWriterImpl
+import tech.httptoolkit.android.vpn.ClientPacketWriter
 import tech.httptoolkit.android.vpn.SessionHandler
 import tech.httptoolkit.android.vpn.SessionManager
 import tech.httptoolkit.android.vpn.socket.SocketNIODataService
@@ -27,23 +27,22 @@ class ProxyVpnRunnable(
     @Volatile private var running = false
 
     // Packets from device apps downstream, heading upstream via this VPN
-    private val clientReader = FileInputStream(vpnInterface.fileDescriptor)
+    private val vpnReadStream = FileInputStream(vpnInterface.fileDescriptor)
 
     // Packets from upstream servers, received by this VPN
-    private val clientWriter = FileOutputStream(vpnInterface.fileDescriptor)
-    private val clientPacketWriter = ClientPacketWriterImpl(clientWriter)
+    private val vpnWriteStream = FileOutputStream(vpnInterface.fileDescriptor)
+    private val vpnPacketWriter = ClientPacketWriter(vpnWriteStream)
+    private val vpnPacketWriterThread = Thread(vpnPacketWriter)
 
-    // Allocate the buffer for a single packet.
-    private val packet = ByteBuffer.allocate(MAX_PACKET_LEN)!!
+    // Background service & task for non-blocking socket
+    private val nioService = SocketNIODataService(vpnPacketWriter)
+    private val dataServiceThread = Thread(nioService, "Socket NIO thread")
 
-    // SessionHandler, which handles sessions whilst writing packets
-    private val handler = SessionHandler.getInstance().apply {
-        setWriter(clientPacketWriter)
-    }
+    private val manager = SessionManager(nioService)
+    private val handler = SessionHandler(manager, nioService, vpnPacketWriter)
 
-    // Background service & task for non-blocking socket
-    private val dataService = SocketNIODataService(clientPacketWriter)
-    private val dataServiceThread = Thread(dataService, "Socket NIO thread")
+    // Allocate the buffer for a single packet.
+    private val packet = ByteBuffer.allocate(MAX_PACKET_LEN)!!
 
     // Our redirect rules, defining which traffic should be forwarded to what proxy address
     private val portRedirections = SparseArray<InetSocketAddress>().apply {
@@ -61,8 +60,9 @@ class ProxyVpnRunnable(
 
         Log.i(TAG, "Vpn thread starting")
 
-        SessionManager.INSTANCE.setTcpPortRedirections(portRedirections)
+        manager.setTcpPortRedirections(portRedirections)
         dataServiceThread.start()
+        vpnPacketWriterThread.start()
 
         var data: ByteArray
         var length: Int
@@ -71,7 +71,7 @@ class ProxyVpnRunnable(
         while (running) {
             data = packet.array()
 
-            length = clientReader.read(data)
+            length = vpnReadStream.read(data)
             if (length > 0) {
                 try {
                     packet.limit(length)
@@ -93,14 +93,16 @@ class ProxyVpnRunnable(
         }
 
         Log.i(TAG, "Vpn thread shutting down")
-
     }
 
     fun stop() {
         if (running) {
             running = false
-            dataService.setShutdown(true)
+            nioService.shutdown()
             dataServiceThread.interrupt()
+
+            vpnPacketWriter.shutdown()
+            vpnPacketWriterThread.interrupt()
         } else {
             Log.w(TAG, "Vpn runnable stopped, but it's not running")
         }

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

@@ -0,0 +1,79 @@
+/*
+ *  Copyright 2014 AT&T
+ *
+ * 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 tech.httptoolkit.android.vpn;
+
+import android.util.Log;
+
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.nio.channels.AsynchronousFileChannel;
+import java.nio.channels.Selector;
+import java.util.Queue;
+import java.util.concurrent.BlockingDeque;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.LinkedBlockingDeque;
+import java.util.concurrent.LinkedBlockingQueue;
+
+import tech.httptoolkit.android.TagKt;
+
+/**
+ * write packet data back to VPN client stream. This class is thread safe.
+ * @author Borey Sao
+ * Date: May 22, 2014
+ */
+public class ClientPacketWriter implements Runnable {
+
+	private final String TAG = TagKt.getTAG(this);
+
+	private final FileOutputStream clientWriter;
+
+	private volatile boolean shutdown = false;
+	private final BlockingDeque<byte[]> packetQueue = new LinkedBlockingDeque<>();
+
+	public ClientPacketWriter(FileOutputStream clientWriter) {
+		this.clientWriter = clientWriter;
+	}
+
+	public void write(byte[] data) {
+		if (data.length > 30000) throw new Error("Packet too large");
+		Log.i(TAG, "Putting " + data.length + " bytes on the write queue");
+		packetQueue.addLast(data);
+	}
+
+	public void shutdown() {
+		this.shutdown = true;
+	}
+
+	@Override
+	public void run() {
+		while (!this.shutdown) {
+			try {
+				byte[] data = this.packetQueue.take();
+				try {
+					this.clientWriter.write(data);
+				} catch (IOException e) {
+					Log.e(TAG, "Error writing " + data.length + " bytes to the VPN");
+					e.printStackTrace();
+
+					this.packetQueue.addFirst(data); // Put the data back, so it's resent
+					Thread.sleep(10); // Add an arbitrary tiny pause, in case that helps
+				}
+			} catch (InterruptedException e) { }
+		}
+	}
+}

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

@@ -19,6 +19,7 @@
 import android.util.Log;
 
 import tech.httptoolkit.android.vpn.network.ip.IPv4Header;
+import tech.httptoolkit.android.vpn.socket.ICloseSession;
 import tech.httptoolkit.android.vpn.transport.tcp.TCPHeader;
 import tech.httptoolkit.android.vpn.transport.udp.UDPHeader;
 import tech.httptoolkit.android.vpn.util.PacketUtil;
@@ -36,15 +37,15 @@
  * Date: May 19, 2014
  */
 public class Session {
-	private static  final String TAG = "Session";
+	private final String TAG = "Session";
 
 	private AbstractSelectableChannel channel;
 
-	private int destIp = 0;
-	private int destPort = 0;
+	private final int destIp;
+	private final int destPort;
 
-	private int sourceIp = 0;
-	private int sourcePort = 0;
+	private final int sourceIp;
+	private final int sourcePort;
 
 	//sequence received from client
 	private long recSequence = 0;
@@ -69,10 +70,10 @@ public class Session {
 	private boolean isConnected = false;
 
 	//receiving buffer for storing data from remote host
-	private ByteArrayOutputStream receivingStream;
+	private final ByteArrayOutputStream receivingStream;
 
 	//sending buffer for storing data from vpn client to be send to destination host
-	private ByteArrayOutputStream sendingStream;
+	private final ByteArrayOutputStream sendingStream;
 
 	private boolean hasReceivedLastSegment = false;
 
@@ -85,7 +86,7 @@ public class Session {
 	private boolean closingConnection = false;
 
 	//indicate data from client is ready for sending to destination
-	private boolean isDataForSendingReady = false;
+	private volatile boolean isDataForSendingReady = false;
 
 	//store data for retransmission
 	private byte[] unackData = null;
@@ -102,24 +103,31 @@ public class Session {
 	//indicate that vpn client has sent FIN flag and it has been acked
 	private boolean ackedToFin = false;
 
-	//indicate that this session is currently being worked on by some SocketDataWorker already
-	private volatile boolean isBusyRead = false;
-	private volatile boolean isBusyWrite = false;
-	
 	//closing session and aborting connection, will be done by background task
 	private volatile boolean abortingConnection = false;
 
 	private SelectionKey selectionkey = null;
 
 	public long connectionStartTime = 0;
+
+	private final ICloseSession sessionCloser;
 
-	Session(int sourceIp, int sourcePort, int destinationIp, int destinationPort){
+	Session(
+		int sourceIp,
+		int sourcePort,
+		int destinationIp,
+		int destinationPort,
+		ICloseSession sessionCloser
+	) {
 		receivingStream = new ByteArrayOutputStream();
 		sendingStream = new ByteArrayOutputStream();
+
 		this.sourceIp = sourceIp;
 		this.sourcePort = sourcePort;
 		this.destIp = destinationIp;
 		this.destPort = destinationPort;
+
+		this.sessionCloser = sessionCloser;
 	}
 
 	/**
@@ -357,18 +365,6 @@ boolean isAckedToFin() {
 		return ackedToFin;
 	}
 
-	public boolean isBusyRead() {
-		return isBusyRead;
-	}
-	public void setBusyread(boolean isbusyread) {
-		this.isBusyRead = isbusyread;
-	}
-	public boolean isBusywrite() {
-		return isBusyWrite;
-	}
-	public void setBusywrite(boolean isbusywrite) {
-		this.isBusyWrite = isbusywrite;
-	}
 	public boolean isAbortingConnection() {
 		return abortingConnection;
 	}
@@ -378,7 +374,7 @@ public void setAbortingConnection(boolean abortingConnection) {
 	public SelectionKey getSelectionKey() {
 		return selectionkey;
 	}
-	void setSelectionKey(SelectionKey selectionkey) {
+	public void setSelectionKey(SelectionKey selectionkey) {
 		this.selectionkey = selectionkey;
 	}
 
@@ -403,6 +399,10 @@ public void unsubscribeKey(int OP) {
 		}
 	}
 
+	public void closeSession() {
+		this.sessionCloser.closeSession(this);
+	}
+
 	public String getSessionKey() {
 		return Session.getSessionKey(this.destIp, this.destPort, this.sourceIp, this.sourcePort);
 	}