package net.openhft.chronicle.network.ssl;

import org.junit.Test;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.ServerSocketChannel;
import java.nio.channels.SocketChannel;
import java.nio.charset.StandardCharsets;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;

import static org.hamcrest.CoreMatchers.is;
import static org.junit.Assert.assertThat;
import static org.junit.Assert.assertTrue;

public final class NioSslIntegrationTest {
    private static final boolean SEND_DATA_BEFORE_SSL_HANDSHAKE = Boolean.getBoolean("ssl.test.payload");

    @Test
    public void shouldEncryptAndDecryptTraffic() throws Exception {
        final ExecutorService threadPool = Executors.newFixedThreadPool(2);

        final ServerSocketChannel serverChannel = ServerSocketChannel.open();
        serverChannel.bind(new InetSocketAddress("0.0.0.0", 13337));
        serverChannel.configureBlocking(true);

        final SocketChannel channel = SocketChannel.open();
        channel.configureBlocking(false);
        channel.connect(new InetSocketAddress("127.0.0.1", serverChannel.socket().getLocalPort()));

        final Client client = new Client(channel);

        final StateMachineProcessor clientProcessor = new StateMachineProcessor(channel, false,
                SSLContextLoader.getInitialisedContext(), client);

        final SocketChannel serverConnection = serverChannel.accept();
        serverConnection.configureBlocking(false);

        final Server server = new Server(serverConnection);
        final StateMachineProcessor serverProcessor = new StateMachineProcessor(serverConnection, true,
                SSLContextLoader.getInitialisedContext(), server);

        while (!(channel.finishConnect() && serverConnection.finishConnect())) {
            Thread.yield();
        }

        if (SEND_DATA_BEFORE_SSL_HANDSHAKE) {
            testDataConnection(channel, serverConnection);
        }

        threadPool.submit(clientProcessor);
        threadPool.submit(serverProcessor);

        client.waitForResponse(10, TimeUnit.SECONDS);

        serverProcessor.stop();
        clientProcessor.stop();
        threadPool.shutdown();
        assertTrue(threadPool.awaitTermination(10, TimeUnit.SECONDS));
    }

    private void testDataConnection(final SocketChannel channel, final SocketChannel serverConnection) throws IOException {
        final ByteBuffer message = ByteBuffer.wrap("test message".getBytes(StandardCharsets.US_ASCII));

        while (message.hasRemaining()) {
            channel.write(message);
        }

        message.clear();
        while (message.hasRemaining()) {
            serverConnection.read(message);
        }

        message.flip();
        assertThat(new String(message.array()), is("test message"));
    }

    private static final class Client extends AbstractSocketBufferHandler {

        private final CountDownLatch latch = new CountDownLatch(1);
        private int counter = 0;
        private int responseCount = 0;

        Client(final SocketChannel socketChannel) {
            super(socketChannel);
        }

        void waitForResponse(final long duration, final TimeUnit timeUnit) throws InterruptedException {
            assertTrue(latch.await(duration, timeUnit));
        }

        @Override
        public void handleDecryptedData(final ByteBuffer input, final ByteBuffer output) {
            if (input.hasRemaining() && input.remaining() != input.capacity()) {
                final byte[] tmp = new byte[input.remaining()];
                input.get(tmp);
                final String response = new String(tmp, StandardCharsets.UTF_8);

                if (responseCount++ == 5) {
                    latch.countDown();
                }
            }
            output.put(("hello " + (counter++)).getBytes(StandardCharsets.US_ASCII));
        }
    }

    private static final class Server extends AbstractSocketBufferHandler {
        private final ByteBuffer lastReceivedMessage = ByteBuffer.allocateDirect(64);

        Server(final SocketChannel channel) {
            super(channel);
        }

        @Override
        public void handleDecryptedData(final ByteBuffer input, final ByteBuffer output) {
            if (input.hasRemaining() && input.remaining() != input.capacity()) {
                lastReceivedMessage.clear();
                lastReceivedMessage.put(input);
                lastReceivedMessage.flip();
            }

            if (lastReceivedMessage.remaining() != lastReceivedMessage.capacity() && lastReceivedMessage.hasRemaining()) {
                output.put("echo: ".getBytes(StandardCharsets.US_ASCII));
                output.put(lastReceivedMessage);
            }
        }
    }
}
@@ -37,16 +37,64 @@
 k.ssl;%0A%0A
+import net.openhft.chronicle.core.io.Closeable;%0A
 import o
@@ -172,16 +172,55 @@
 ddress;%0A
+import java.net.StandardSocketOptions;%0A
 import j
@@ -1113,24 +1113,99 @@
 %22, 13337));%0A
+        serverChannel.setOption(StandardSocketOptions.SO_REUSEADDR, true);%0A
         serv
@@ -1439,24 +1439,43 @@
 lPort()));%0A%0A
+        try %7B%0A%0A    
         fina
@@ -1510,24 +1510,28 @@
 (channel);%0A%0A
+    
         fina
@@ -1622,32 +1622,36 @@
                 
+    
 SSLContextLoader
@@ -1682,24 +1682,28 @@
 , client);%0A%0A
+    
         fina
@@ -1757,24 +1757,28 @@
 l.accept();%0A
+    
         serv
@@ -1817,32 +1817,36 @@
 alse);%0A%0A        
+    
 final Server ser
@@ -1873,32 +1873,36 @@
 verConnection);%0A
+    
         final St
@@ -1997,32 +1997,36 @@
                 
+   
+ 
 SSLContextLoader
@@ -2057,24 +2057,28 @@
 , server);%0A%0A
+    
         whil
@@ -2158,16 +2158,20 @@
         
+    
 Thread.y
@@ -2178,35 +2178,43 @@
 ield();%0A        
+   
+ 
 %7D%0A%0A
+    
         if (SEND
@@ -2255,16 +2255,20 @@
         
+ 
+   
 testData
@@ -2306,35 +2306,43 @@
 ction);%0A        
+  
+  
 %7D%0A%0A
+    
         threadPo
@@ -2369,32 +2369,36 @@
 essor);%0A        
+    
 threadPool.submi
@@ -2422,24 +2422,28 @@
 );%0A%0A        
+  
+  
 client.waitF
@@ -2476,17 +2476,20 @@
 CONDS);%0A
-%0A
+    
         
@@ -2504,32 +2504,36 @@
 ocessor.stop();%0A
+    
         clientPr
@@ -2540,32 +2540,123 @@
 ocessor.stop();%0A
+        %7D finally %7B%0A            Closeable.closeQuietly(channel, serverChannel);%0A        %7D%0A%0A
         threadPo

Improve test