From 8fc8d4587b5ec4eee031d0e41e091740ca99d9de Mon Sep 17 00:00:00 2001 From: Jason Ernst Date: Wed, 16 Dec 2020 08:30:13 -0800 Subject: [PATCH 1/3] refactor to remove statics --- .../grape/service/GrapeVpnService.java | 12 ++----- .../network/grape/service/SessionHandler.java | 33 +++++++---------- .../network/grape/service/SessionManager.java | 12 +++---- .../grape/service/SocketDataReaderWorker.java | 11 +++--- .../grape/service/SocketDataWriterWorker.java | 12 ++++--- .../grape/service/SocketProtector.java | 35 +++---------------- .../java/network/grape/service/VpnWriter.java | 28 +++++++++------ 7 files changed, 60 insertions(+), 83 deletions(-) diff --git a/app/src/main/java/network/grape/service/GrapeVpnService.java b/app/src/main/java/network/grape/service/GrapeVpnService.java index ed705249..9c45da30 100644 --- a/app/src/main/java/network/grape/service/GrapeVpnService.java +++ b/app/src/main/java/network/grape/service/GrapeVpnService.java @@ -69,12 +69,6 @@ public int onStartCommand(Intent intent, int flags, int startId) { public void run() { logger.info("running vpn service"); - // map this class into the socket protector implementation so that other classes like the - // SessionHandler can protect sockets later on (protect is provided by the VpnService which - // this class inherits). - SocketProtector protector = SocketProtector.getInstance(); - protector.setProtector(this); - try { if (startVpnService()) { logger.info("VPN Service started"); @@ -134,11 +128,11 @@ void startTrafficHandler() throws IOException { // Allocate the buffer for a single packet. ByteBuffer packet = ByteBuffer.allocate(MAX_PACKET_LEN); - SessionHandler handler = SessionHandler.getInstance(); - handler.setOutputStream(clientWriter); + SessionManager sessionManager = new SessionManager(); + SessionHandler handler = new SessionHandler(sessionManager, new SocketProtector(this)); // background thread for writing output to the vpn outputstream - vpnWriter = new VpnWriter(clientWriter); + vpnWriter = new VpnWriter(clientWriter, sessionManager); vpnWriterThread = new Thread(vpnWriter); vpnWriterThread.start(); diff --git a/app/src/main/java/network/grape/service/SessionHandler.java b/app/src/main/java/network/grape/service/SessionHandler.java index 1bc09f7b..4b006259 100644 --- a/app/src/main/java/network/grape/service/SessionHandler.java +++ b/app/src/main/java/network/grape/service/SessionHandler.java @@ -3,7 +3,6 @@ import static network.grape.lib.network.ip.IpHeader.IP4_VERSION; import static network.grape.lib.network.ip.IpHeader.IP6_VERSION; -import java.io.FileOutputStream; import java.io.IOException; import java.net.Inet4Address; import java.net.InetSocketAddress; @@ -30,20 +29,14 @@ */ public class SessionHandler { private final Logger logger = LoggerFactory.getLogger(SessionHandler.class); - private static final SessionHandler handler = new SessionHandler(); - private SocketProtector protector = SocketProtector.getInstance(); - private Selector selector = SessionManager.INSTANCE.getSelector(); - private FileOutputStream outputStream; - - public static SessionHandler getInstance() { - return handler; - } - - private SessionHandler() { - } - - void setOutputStream(FileOutputStream outputStream) { - this.outputStream = outputStream; + private final SocketProtector protector; + private final Selector selector; + private final SessionManager sessionManager; + + public SessionHandler(SessionManager sessionManager, SocketProtector protector) { + this.sessionManager = sessionManager; + this.selector = sessionManager.getSelector(); + this.protector = protector; } /** @@ -92,10 +85,10 @@ private void handleTcpPacket(ByteBuffer payload, IpHeader ipHeader, TcpHeader tc private void handleUdpPacket(ByteBuffer payload, IpHeader ipHeader, UdpHeader udpHeader) { // try to find an existing session - Session session = SessionManager.INSTANCE - .getSession(ipHeader.getSourceAddress(), udpHeader.getSourcePort(), - ipHeader.getDestinationAddress(), - udpHeader.getDestinationPort(), TransportHeader.UDP_PROTOCOL); + Session session = sessionManager.getSession(ipHeader.getSourceAddress(), + udpHeader.getSourcePort(), + ipHeader.getDestinationAddress(), + udpHeader.getDestinationPort(), TransportHeader.UDP_PROTOCOL); // otherwise create a new one if (session == null) { @@ -151,7 +144,7 @@ private void handleUdpPacket(ByteBuffer payload, IpHeader ipHeader, UdpHeader ud } session.setChannel(channel); - if (!SessionManager.INSTANCE.putSession(session)) { + if (!sessionManager.putSession(session)) { // just in case we fail to add it (we should hopefully never get here) logger.error("Unable to create a new session in the session manager for " + session); return; diff --git a/app/src/main/java/network/grape/service/SessionManager.java b/app/src/main/java/network/grape/service/SessionManager.java index 1f23badb..c90d53b7 100644 --- a/app/src/main/java/network/grape/service/SessionManager.java +++ b/app/src/main/java/network/grape/service/SessionManager.java @@ -19,15 +19,15 @@ * You can think of this thing as a sort of a local NAT within the phone because it has to keep * track of all of the outbound <-> inbound mappings of ports and IPs like a NAT table does. */ -public enum SessionManager { - INSTANCE; +public class SessionManager { - private final Logger logger = LoggerFactory.getLogger(SessionManager.class); - private final Map table = new ConcurrentHashMap<>(); - @Getter - private Selector selector; + private final Logger logger; + private final Map table; + @Getter private Selector selector; SessionManager() { + logger = LoggerFactory.getLogger(SessionManager.class); + table = new ConcurrentHashMap<>(); try { selector = Selector.open(); } catch (IOException ex) { diff --git a/app/src/main/java/network/grape/service/SocketDataReaderWorker.java b/app/src/main/java/network/grape/service/SocketDataReaderWorker.java index f4b273e0..9408ce22 100644 --- a/app/src/main/java/network/grape/service/SocketDataReaderWorker.java +++ b/app/src/main/java/network/grape/service/SocketDataReaderWorker.java @@ -23,18 +23,21 @@ * terminates. */ public class SocketDataReaderWorker implements Runnable { - private final Logger logger = LoggerFactory.getLogger(SocketDataReaderWorker.class); + private final Logger logger; + private final SessionManager sessionManager; private FileOutputStream outputStream; private String sessionKey; - SocketDataReaderWorker(FileOutputStream outputStream, String sessionKey) { + SocketDataReaderWorker(FileOutputStream outputStream, String sessionKey, SessionManager sessionManager) { + this.logger = LoggerFactory.getLogger(SocketDataReaderWorker.class); this.outputStream = outputStream; this.sessionKey = sessionKey; + this.sessionManager = sessionManager; } @Override public void run() { - Session session = SessionManager.INSTANCE.getSessionByKey(sessionKey); + Session session = sessionManager.getSessionByKey(sessionKey); if (session == null) { logger.error("Session NOT FOUND: " + sessionKey); return; @@ -82,7 +85,7 @@ public void run() { return; } - SessionManager.INSTANCE.closeSession(session); + sessionManager.closeSession(session); } else { session.setBusyRead(false); } diff --git a/app/src/main/java/network/grape/service/SocketDataWriterWorker.java b/app/src/main/java/network/grape/service/SocketDataWriterWorker.java index d3a6e53c..49fe1fdc 100644 --- a/app/src/main/java/network/grape/service/SocketDataWriterWorker.java +++ b/app/src/main/java/network/grape/service/SocketDataWriterWorker.java @@ -16,18 +16,22 @@ * TCP connection, writes an RST to the VPN stream to reset the connection. */ public class SocketDataWriterWorker implements Runnable { - private final Logger logger = LoggerFactory.getLogger(SocketDataWriterWorker.class); + private final Logger logger; private FileOutputStream outputStream; // really only used for sending RST packet on TCP private String sessionKey; + private SessionManager sessionManager; - SocketDataWriterWorker(FileOutputStream outputStream, String sessionKey) { + SocketDataWriterWorker(FileOutputStream outputStream, String sessionKey, + SessionManager sessionManager) { + this.logger = LoggerFactory.getLogger(SocketDataWriterWorker.class); this.outputStream = outputStream; this.sessionKey = sessionKey; + this.sessionManager = sessionManager; } @Override public void run() { - final Session session = SessionManager.INSTANCE.getSessionByKey(sessionKey); + final Session session = sessionManager.getSessionByKey(sessionKey); if (session == null) { logger.error("No session related to " + sessionKey + "for write"); return; @@ -73,7 +77,7 @@ public void run() { return; } - SessionManager.INSTANCE.closeSession(session); + sessionManager.closeSession(session); } } diff --git a/app/src/main/java/network/grape/service/SocketProtector.java b/app/src/main/java/network/grape/service/SocketProtector.java index bda61e6f..58d428d9 100644 --- a/app/src/main/java/network/grape/service/SocketProtector.java +++ b/app/src/main/java/network/grape/service/SocketProtector.java @@ -4,39 +4,14 @@ import java.net.Socket; /** - * Singleton class which gives access to socket protection from the GrapeVpnService more widely. The - * GrapeVpnService sets itself as the socket protector, and then every class using this singleton - * can protect sockets. + * Gives access to socket protection from the GrapeVpnService more widely. The GrapeVpnService sets + * itself as the socket protector, and then every class using this singleton can protect sockets. */ public class SocketProtector { - private static final Object synObject = new Object(); - private static volatile SocketProtector instance = null; - private ProtectSocket protector = null; + private final ProtectSocket protector; - /** - * Provides a socket protector to any class which requires one. - * @return a singleton instance of the socket protector - */ - public static SocketProtector getInstance() { - if (instance == null) { - synchronized (synObject) { - if (instance == null) { - instance = new SocketProtector(); - } - } - } - return instance; - } - - /** - * set class that implement IProtectSocket if only if it was never set before. - * - * @param protector ProtectSocket - */ - public void setProtector(ProtectSocket protector) { - if (this.protector == null) { - this.protector = protector; - } + public SocketProtector(ProtectSocket protector) { + this.protector = protector; } public void protect(Socket socket) { diff --git a/app/src/main/java/network/grape/service/VpnWriter.java b/app/src/main/java/network/grape/service/VpnWriter.java index 4fd2459c..845db944 100644 --- a/app/src/main/java/network/grape/service/VpnWriter.java +++ b/app/src/main/java/network/grape/service/VpnWriter.java @@ -26,10 +26,11 @@ */ public class VpnWriter implements Runnable { - private final Logger logger = LoggerFactory.getLogger(VpnWriter.class); + private final Logger logger; public static final Object syncSelector = new Object(); public static final Object syncSelector2 = new Object(); - private FileOutputStream outputStream; + private final FileOutputStream outputStream; + private final SessionManager sessionManager; private Selector selector; //create thread pool for reading/writing data to socket private ThreadPoolExecutor workerPool; @@ -37,21 +38,28 @@ public class VpnWriter implements Runnable { /** * Construct a new VpnWriter. + * * @param outputStream the stream to write back into the VPN interface. */ - public VpnWriter(FileOutputStream outputStream) { + public VpnWriter(FileOutputStream outputStream, SessionManager sessionManager) { + this.logger = LoggerFactory.getLogger(VpnWriter.class); this.outputStream = outputStream; + this.sessionManager = sessionManager; final BlockingQueue taskQueue = new LinkedBlockingQueue<>(); workerPool = new ThreadPoolExecutor(8, 100, 10, TimeUnit.SECONDS, taskQueue); } /** * Construct a new VpnWriter with the workerpool provided. + * * @param outputStream the stream to write back into the VPN interface. - * @param workerPool the worker pool to execute reader and writer threads in. + * @param workerPool the worker pool to execute reader and writer threads in. */ - public VpnWriter(FileOutputStream outputStream, ThreadPoolExecutor workerPool) { + public VpnWriter(FileOutputStream outputStream, SessionManager sessionManager, + ThreadPoolExecutor workerPool) { + this.logger = LoggerFactory.getLogger(VpnWriter.class); this.outputStream = outputStream; + this.sessionManager = sessionManager; this.workerPool = workerPool; } @@ -60,7 +68,7 @@ public VpnWriter(FileOutputStream outputStream, ThreadPoolExecutor workerPool) { */ public void run() { logger.info("VpnWriter starting in the background"); - selector = SessionManager.INSTANCE.getSelector(); + selector = sessionManager.getSelector(); running = true; while (running) { @@ -109,7 +117,7 @@ protected void processUdpSelectionKey(SelectionKey key) { return; } DatagramChannel channel = (DatagramChannel) key.channel(); - Session session = SessionManager.INSTANCE.getSessionByChannel(channel); + Session session = sessionManager.getSessionByChannel(channel); String keyString = channel.socket().getLocalAddress().toString() + ":" + channel.socket().getLocalPort() + "," + channel.socket().getInetAddress().toString() + ":" @@ -153,7 +161,7 @@ protected void processUdpSelectionKey(SelectionKey key) { * Generic selector handling for both TCP and UDP sessions. * * @param selectionKey the key in the selection set which is marked for reading or writing. - * @param session the session associated with the selection key. + * @param session the session associated with the selection key. */ protected void processSelector(SelectionKey selectionKey, Session session) { // tcp has PSH flag when data is ready for sending, UDP does not have this @@ -161,13 +169,13 @@ protected void processSelector(SelectionKey selectionKey, Session session) { && session.hasDataToSend() && session.isDataForSendingReady()) { session.setBusyWrite(true); final SocketDataWriterWorker worker = - new SocketDataWriterWorker(outputStream, session.getKey()); + new SocketDataWriterWorker(outputStream, session.getKey(), sessionManager); workerPool.execute(worker); } if (selectionKey.isValid() && selectionKey.isReadable() && !session.isBusyRead()) { session.setBusyRead(true); final SocketDataReaderWorker worker = - new SocketDataReaderWorker(outputStream, session.getKey()); + new SocketDataReaderWorker(outputStream, session.getKey(), sessionManager); workerPool.execute(worker); } } From b35df25773c06cccaa7e3cf67e9c0593222e656f Mon Sep 17 00:00:00 2001 From: Jason Ernst Date: Wed, 16 Dec 2020 11:02:11 -0800 Subject: [PATCH 2/3] Tests for udp selectionkey --- .../network/grape/service/VpnWriterTest.java | 113 +++++++++++++++++- 1 file changed, 107 insertions(+), 6 deletions(-) diff --git a/app/src/test/java/network/grape/service/VpnWriterTest.java b/app/src/test/java/network/grape/service/VpnWriterTest.java index 4e86e996..1a0010ae 100644 --- a/app/src/test/java/network/grape/service/VpnWriterTest.java +++ b/app/src/test/java/network/grape/service/VpnWriterTest.java @@ -4,34 +4,135 @@ import static org.mockito.Mockito.doNothing; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; +import static org.mockito.Mockito.reset; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; + import java.io.FileOutputStream; +import java.io.IOException; +import java.net.DatagramSocket; +import java.net.InetAddress; +import java.nio.channels.DatagramChannel; import java.nio.channels.SelectionKey; import java.util.concurrent.ThreadPoolExecutor; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Disabled; import org.junit.jupiter.api.Test; import org.mockito.Mockito; public class VpnWriterTest { - @Test - public void testProcessUdpKey() { - FileOutputStream fileOutputStream = mock(FileOutputStream.class); - ThreadPoolExecutor workerPool = mock(ThreadPoolExecutor.class); - VpnWriter vpnWriter = new VpnWriter(fileOutputStream, workerPool); + FileOutputStream fileOutputStream; + ThreadPoolExecutor workerPool; + SessionManager sessionManager; + + @BeforeEach + public void before() { + fileOutputStream = mock(FileOutputStream.class); + workerPool = mock(ThreadPoolExecutor.class); + sessionManager = mock(SessionManager.class); + } + + SelectionKey prepSelectionKey(boolean exception) throws IOException { SelectionKey selectionKey = mock(SelectionKey.class); + when(selectionKey.isValid()).thenReturn(true); + DatagramChannel channel = mock(DatagramChannel.class); + when(selectionKey.channel()).thenReturn(channel); + DatagramSocket socket = mock(DatagramSocket.class); + when(channel.socket()).thenReturn(socket); + if (exception) { + when(channel.connect(any())).thenThrow(IOException.class); + } else { + when(channel.connect(any())).thenReturn(channel); + when(channel.isConnected()).thenReturn(true); + } + InetAddress inetAddress = mock(InetAddress.class); + when(socket.getLocalAddress()).thenReturn(inetAddress); + when(socket.getInetAddress()).thenReturn(inetAddress); + return selectionKey; + } + @Test + public void invalidUdpSelector() { + VpnWriter vpnWriter = spy(new VpnWriter(fileOutputStream, sessionManager, workerPool)); + SelectionKey selectionKey = mock(SelectionKey.class); when(selectionKey.isValid()).thenReturn(false); vpnWriter.processUdpSelectionKey(selectionKey); } + @Test + public void udpSessionNotFound() throws IOException { + SelectionKey selectionKey = prepSelectionKey(false); + VpnWriter vpnWriter = spy(new VpnWriter(fileOutputStream, sessionManager, workerPool)); + vpnWriter.processUdpSelectionKey(selectionKey); + } + + @Test + public void sessionNotConnectedKeyConnectable() throws IOException { + VpnWriter vpnWriter = spy(new VpnWriter(fileOutputStream, sessionManager, workerPool)); + Session session = mock(Session.class); + + //io exception on connect + SelectionKey selectionKey = prepSelectionKey(true); + InetAddress inetAddress = mock(InetAddress.class); + when(session.getDestinationIp()).thenReturn(inetAddress); + when(session.isConnected()).thenReturn(false); + when(selectionKey.isConnectable()).thenReturn(true); + when(sessionManager.getSessionByChannel(any())).thenReturn(session); + doNothing().when(vpnWriter).processSelector(any(), any()); + vpnWriter.processUdpSelectionKey(selectionKey); + verify(vpnWriter, never()).processSelector(any(), any()); + verify(session, times(1)).setAbortingConnection(true); + verify(vpnWriter, never()).processSelector(any(), any()); + + // good path + selectionKey = prepSelectionKey(false); + inetAddress = mock(InetAddress.class); + when(session.getDestinationIp()).thenReturn(inetAddress); + when(session.isConnected()).thenReturn(false); + when(selectionKey.isConnectable()).thenReturn(true); + when(sessionManager.getSessionByChannel(any())).thenReturn(session); + doNothing().when(vpnWriter).processSelector(any(), any()); + vpnWriter.processUdpSelectionKey(selectionKey); + verify(vpnWriter, times(1)).processSelector(any(), any()); + } + + @Test + public void sessionNotConnectedKeyNotConnectable() throws IOException { + VpnWriter vpnWriter = spy(new VpnWriter(fileOutputStream, sessionManager, workerPool)); + Session session = mock(Session.class); + when(sessionManager.getSessionByChannel(any())).thenReturn(session); + SelectionKey selectionKey = prepSelectionKey(true); + when(session.isConnected()).thenReturn(false); + when(selectionKey.isConnectable()).thenReturn(false); + doNothing().when(vpnWriter).processSelector(any(), any()); + vpnWriter.processUdpSelectionKey(selectionKey); + verify(vpnWriter, never()).processSelector(any(), any()); + } + + @Test + public void sessionConnected() throws IOException { + VpnWriter vpnWriter = spy(new VpnWriter(fileOutputStream, sessionManager, workerPool)); + Session session = mock(Session.class); + when(sessionManager.getSessionByChannel(any())).thenReturn(session); + SelectionKey selectionKey = prepSelectionKey(false); + when(session.isConnected()).thenReturn(true); + doNothing().when(vpnWriter).processSelector(any(), any()); + vpnWriter.processUdpSelectionKey(selectionKey); + verify(vpnWriter, times(1)).processSelector(any(), any()); + verify(session, never()).setChannel(any()); + } + @Test public void testProcessSelector() { FileOutputStream fileOutputStream = mock(FileOutputStream.class); Session session = mock(Session.class); SelectionKey selectionKey = mock(SelectionKey.class); ThreadPoolExecutor workerPool = mock(ThreadPoolExecutor.class); - VpnWriter vpnWriter = new VpnWriter(fileOutputStream, workerPool); + SessionManager sessionManager = mock(SessionManager.class); + VpnWriter vpnWriter = new VpnWriter(fileOutputStream, sessionManager, workerPool); when(selectionKey.isValid()).thenReturn(false); vpnWriter.processSelector(selectionKey, session); From 6254a48d2ac87531f49470f7df769570806f62e4 Mon Sep 17 00:00:00 2001 From: Jason Ernst Date: Wed, 16 Dec 2020 12:59:13 -0800 Subject: [PATCH 3/3] increased coverage of VpnWriter --- .../grape/service/GrapeVpnService.java | 8 +- .../java/network/grape/service/VpnWriter.java | 27 ++-- .../network/grape/service/VpnWriterTest.java | 116 +++++++++++++++++- 3 files changed, 130 insertions(+), 21 deletions(-) diff --git a/app/src/main/java/network/grape/service/GrapeVpnService.java b/app/src/main/java/network/grape/service/GrapeVpnService.java index 9c45da30..c95b6a2c 100644 --- a/app/src/main/java/network/grape/service/GrapeVpnService.java +++ b/app/src/main/java/network/grape/service/GrapeVpnService.java @@ -11,6 +11,10 @@ import java.net.Socket; import java.net.UnknownHostException; import java.nio.ByteBuffer; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.TimeUnit; import network.grape.lib.PacketHeaderException; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -132,7 +136,9 @@ void startTrafficHandler() throws IOException { SessionHandler handler = new SessionHandler(sessionManager, new SocketProtector(this)); // background thread for writing output to the vpn outputstream - vpnWriter = new VpnWriter(clientWriter, sessionManager); + final BlockingQueue taskQueue = new LinkedBlockingQueue<>(); + ThreadPoolExecutor executor = new ThreadPoolExecutor(8, 100, 10, TimeUnit.SECONDS, taskQueue); + vpnWriter = new VpnWriter(clientWriter, sessionManager, executor); vpnWriterThread = new Thread(vpnWriter); vpnWriterThread.start(); diff --git a/app/src/main/java/network/grape/service/VpnWriter.java b/app/src/main/java/network/grape/service/VpnWriter.java index 845db944..882ce469 100644 --- a/app/src/main/java/network/grape/service/VpnWriter.java +++ b/app/src/main/java/network/grape/service/VpnWriter.java @@ -36,19 +36,6 @@ public class VpnWriter implements Runnable { private ThreadPoolExecutor workerPool; private volatile boolean running; - /** - * Construct a new VpnWriter. - * - * @param outputStream the stream to write back into the VPN interface. - */ - public VpnWriter(FileOutputStream outputStream, SessionManager sessionManager) { - this.logger = LoggerFactory.getLogger(VpnWriter.class); - this.outputStream = outputStream; - this.sessionManager = sessionManager; - final BlockingQueue taskQueue = new LinkedBlockingQueue<>(); - workerPool = new ThreadPoolExecutor(8, 100, 10, TimeUnit.SECONDS, taskQueue); - } - /** * Construct a new VpnWriter with the workerpool provided. * @@ -63,6 +50,14 @@ public VpnWriter(FileOutputStream outputStream, SessionManager sessionManager, this.workerPool = workerPool; } + boolean isRunning() { + return running; + } + + boolean notRunning() { + return !running; + } + /** * Main thread for the VpnWriter. */ @@ -70,7 +65,7 @@ public void run() { logger.info("VpnWriter starting in the background"); selector = sessionManager.getSelector(); running = true; - while (running) { + while (isRunning()) { // first just try to wait for a socket to be ready for a connect, read, etc try { @@ -87,7 +82,7 @@ public void run() { continue; } - if (!running) { + if (notRunning()) { break; } @@ -103,7 +98,7 @@ public void run() { processUdpSelectionKey(key); } iterator.remove(); - if (!running) { + if (notRunning()) { break; } } diff --git a/app/src/test/java/network/grape/service/VpnWriterTest.java b/app/src/test/java/network/grape/service/VpnWriterTest.java index 1a0010ae..d68fba6d 100644 --- a/app/src/test/java/network/grape/service/VpnWriterTest.java +++ b/app/src/test/java/network/grape/service/VpnWriterTest.java @@ -10,12 +10,20 @@ import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; + import java.io.FileOutputStream; import java.io.IOException; import java.net.DatagramSocket; import java.net.InetAddress; import java.nio.channels.DatagramChannel; +import java.nio.channels.FileChannel; import java.nio.channels.SelectionKey; +import java.nio.channels.Selector; +import java.nio.channels.ServerSocketChannel; +import java.nio.channels.SocketChannel; +import java.util.HashSet; +import java.util.Iterator; +import java.util.Set; import java.util.concurrent.ThreadPoolExecutor; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Disabled; @@ -127,12 +135,9 @@ public void sessionConnected() throws IOException { @Test public void testProcessSelector() { - FileOutputStream fileOutputStream = mock(FileOutputStream.class); Session session = mock(Session.class); SelectionKey selectionKey = mock(SelectionKey.class); - ThreadPoolExecutor workerPool = mock(ThreadPoolExecutor.class); - SessionManager sessionManager = mock(SessionManager.class); - VpnWriter vpnWriter = new VpnWriter(fileOutputStream, sessionManager, workerPool); + VpnWriter vpnWriter = spy(new VpnWriter(fileOutputStream, sessionManager, workerPool)); when(selectionKey.isValid()).thenReturn(false); vpnWriter.processSelector(selectionKey, session); @@ -185,4 +190,107 @@ public void testProcessSelector() { vpnWriter.processSelector(selectionKey, session); verify(session, Mockito.times(1)).setBusyWrite(true); } + + @Test + public void runTest() throws InterruptedException, IOException { + + // base case, nothing back from the selector + VpnWriter vpnWriter = spy(new VpnWriter(fileOutputStream, sessionManager, workerPool)); + Selector selector = mock(Selector.class); + when(sessionManager.getSelector()).thenReturn(selector); + when(vpnWriter.isRunning()).thenReturn(true).thenReturn(false); + Thread t = new Thread(vpnWriter); + t.start(); + t.join(); + + // exception on select + vpnWriter = spy(new VpnWriter(fileOutputStream, sessionManager, workerPool)); + selector = mock(Selector.class); + when(selector.select()).thenThrow(IOException.class); + when(sessionManager.getSelector()).thenReturn(selector); + when(vpnWriter.isRunning()).thenReturn(true).thenReturn(false); + t = new Thread(vpnWriter); + t.start(); + t.join(); + + // exception on select + interrupt in handler + vpnWriter = spy(new VpnWriter(fileOutputStream, sessionManager, workerPool)); + selector = mock(Selector.class); + when(selector.select()).thenThrow(IOException.class); + when(sessionManager.getSelector()).thenReturn(selector); + when(vpnWriter.isRunning()).thenReturn(true).thenReturn(false); + t = new Thread(vpnWriter); + t.start(); + // there is a chance thread scheduling will be bad and the interrupted exception be thrown + // in time here, but its okay. + Thread.sleep(100); + t.interrupt(); + t.join(); + } + + @Test public void runTestSelectionSet() throws InterruptedException { + // non-empty iterator + Set selectionKeySet = new HashSet<>(); + SelectionKey udpKey = mock(SelectionKey.class); + DatagramChannel udpChannel = mock(DatagramChannel.class); + when(udpKey.channel()).thenReturn(udpChannel); + selectionKeySet.add(udpKey); + + SelectionKey tcpKey = mock(SelectionKey.class); + SocketChannel tcpChannel = mock(SocketChannel.class); + when(tcpKey.channel()).thenReturn(tcpChannel); + selectionKeySet.add(tcpKey); + + SelectionKey serverSocketKey = mock(SelectionKey.class); + ServerSocketChannel serverSocketChannel = mock(ServerSocketChannel.class); + when(serverSocketKey.channel()).thenReturn(serverSocketChannel); + selectionKeySet.add(serverSocketKey); + + VpnWriter vpnWriter = spy(new VpnWriter(fileOutputStream, sessionManager, workerPool)); + doNothing().when(vpnWriter).processUdpSelectionKey(any()); + + Selector selector = mock(Selector.class); + when(sessionManager.getSelector()).thenReturn(selector); + when(selector.selectedKeys()).thenReturn(selectionKeySet); + Thread t = new Thread(vpnWriter); + t.start(); + when(vpnWriter.isRunning()).thenReturn(true).thenReturn(false); + when(vpnWriter.notRunning()).thenReturn(false); + vpnWriter.shutdown(); + t.join(); + } + + @Test public void runTestNotRunning() throws InterruptedException { + // base case, nothing back from the selector + VpnWriter vpnWriter = spy(new VpnWriter(fileOutputStream, sessionManager, workerPool)); + Selector selector = mock(Selector.class); + when(sessionManager.getSelector()).thenReturn(selector); + when(vpnWriter.isRunning()).thenReturn(true).thenReturn(false); + when(vpnWriter.notRunning()).thenReturn(true); + Thread t = new Thread(vpnWriter); + t.start(); + t.join(); + } + + @Test public void runTestNotRunningNonEmptyIterator() throws InterruptedException { + // non-empty iterator + Set selectionKeySet = new HashSet<>(); + SelectionKey udpKey = mock(SelectionKey.class); + DatagramChannel udpChannel = mock(DatagramChannel.class); + when(udpKey.channel()).thenReturn(udpChannel); + selectionKeySet.add(udpKey); + + VpnWriter vpnWriter = spy(new VpnWriter(fileOutputStream, sessionManager, workerPool)); + doNothing().when(vpnWriter).processUdpSelectionKey(any()); + + Selector selector = mock(Selector.class); + when(sessionManager.getSelector()).thenReturn(selector); + when(selector.selectedKeys()).thenReturn(selectionKeySet); + Thread t = new Thread(vpnWriter); + t.start(); + when(vpnWriter.isRunning()).thenReturn(true).thenReturn(false); + when(vpnWriter.notRunning()).thenReturn(false).thenReturn(true); + vpnWriter.shutdown(); + t.join(); + } }