Merge pull request #224 from pontusmelke/1.0-tls-write-issue

Better handling of BUFFER_OVERFLOW on writes

Author: pontusmelke

driver/src/main/java/org/neo4j/driver/internal/connector/socket/TLSSocketChannel.java

@@ -28,9 +28,9 @@
 import javax.net.ssl.SSLEngineResult.HandshakeStatus;
 import javax.net.ssl.SSLEngineResult.Status;
 
-import org.neo4j.driver.v1.Logger;
 import org.neo4j.driver.internal.util.BytePrinter;
 import org.neo4j.driver.v1.Config.TrustStrategy;
+import org.neo4j.driver.v1.Logger;
 import org.neo4j.driver.v1.exceptions.ClientException;
 
 import static javax.net.ssl.SSLEngineResult.HandshakeStatus.FINISHED;
@@ -163,8 +163,8 @@ private HandshakeStatus runDelegatedTasks()
      * To verify if deciphering is done successfully, we could check if any bytes has been read into {@code buffer},
      * as the deciphered bytes will only be saved into {@code buffer} when deciphering is carried out successfully.
      *
-     * @param buffer
-     * @return
+     * @param buffer to read data into.
+     * @return The status of the current handshake.
      * @throws IOException
      */
     private HandshakeStatus unwrap( ByteBuffer buffer ) throws IOException
@@ -261,7 +261,7 @@ private HandshakeStatus unwrap( ByteBuffer buffer ) throws IOException
      * a loop
      *
      * @param buffer contains the bytes to send to channel
-     * @return
+     * @return The status of the current handshake
      * @throws IOException
      */
     private HandshakeStatus wrap( ByteBuffer buffer ) throws IOException
@@ -277,7 +277,7 @@ private HandshakeStatus wrap( ByteBuffer buffer ) throws IOException
         case OK:
             handshakeStatus = runDelegatedTasks();
             cipherOut.flip();
-            while(cipherOut.hasRemaining())
+            while ( cipherOut.hasRemaining() )
             {
                 channel.write( cipherOut );
             }
@@ -287,19 +287,30 @@ private HandshakeStatus wrap( ByteBuffer buffer ) throws IOException
             // Enlarge the buffer and return the old status
             int curNetSize = cipherOut.capacity();
             int netSize = sslEngine.getSession().getPacketBufferSize();
-            if ( curNetSize >= netSize || buffer.capacity() > netSize )
+            if ( netSize > curNetSize )
             {
-                // TODO
-                throw new ClientException(
-                        String.format( "Failed to enlarge network buffer from %s to %s. This is either because the " +
-                                       "new size is however less than the old size, or because the application " +
-                                       "buffer size %s is so big that the application data still cannot fit into the " +
-                                       "new network buffer.", curNetSize, netSize, buffer.capacity() ) );
+                // enlarge the peer application data buffer
+                cipherOut = ByteBuffer.allocate( netSize );
+                logger.debug( "Enlarged network output buffer from %s to %s. " +
+                              "This operation should be a rare operation.", curNetSize, netSize );
+            }
+            else
+            {
+                // flush as much data as possible
+                cipherOut.flip();
+                int written = channel.write( cipherOut );
+                if (written == 0)
+                {
+                    throw new ClientException(
+                            String.format(
+                                    "Failed to enlarge network buffer from %s to %s. This is either because the " +
+                                    "new size is however less than the old size, or because the application " +
+                                    "buffer size %s is so big that the application data still cannot fit into the " +
+                                    "new network buffer.", curNetSize, netSize, buffer.capacity() ) );
+                }
+                cipherOut.compact();
+                logger.debug( "Network output buffer couldn't be enlarged, flushing data to the channel instead." );
             }
-
-            cipherOut = ByteBuffer.allocate( netSize );
-            logger.debug( "Enlarged network output buffer from %s to %s. " +
-                          "This operation should be a rare operation.", curNetSize, netSize );
             break;
         default:
             throw new ClientException( "Got unexpected status " + status );
@@ -320,9 +331,9 @@ private HandshakeStatus wrap( ByteBuffer buffer ) throws IOException
      * After the method call, the new position of {@code from.pos} will be {@code from.pos + p}, and similarly,
      * the new position of {@code to.pos} will be {@code to.pos + p}
      *
-     * @param from
-     * @param to
-     * @return
+     * @param from buffer to copy from
+     * @param to buffer to copy to
+     * @return the number of transferred bytes
      */
     static int bufferCopy( ByteBuffer from, ByteBuffer to )
     {
@@ -341,9 +352,9 @@ static int bufferCopy( ByteBuffer from, ByteBuffer to )
 
     /**
      * Create SSLEngine with the SSLContext just created.
-     * @param host
-     * @param port
-     * @param sslContext
+     * @param host the host to connect to
+     * @param port the port to connect to
+     * @param sslContext the current ssl context
      */
     private static SSLEngine createSSLEngine( String host, int port, SSLContext sslContext )
     {
@@ -431,10 +442,10 @@ public void close() throws IOException
             channel.close();
             logger.debug( "TLS connection closed" );
         }
-        catch(IOException e)
+        catch ( IOException e )
         {
             // Treat this as ok - the connection is closed, even if the TLS session did not exit cleanly.
-            logger.warn( "TLS socket could not be closed cleanly: '"+e.getMessage()+"'", e );
+            logger.warn( "TLS socket could not be closed cleanly: '" + e.getMessage() + "'", e );
         }
     }
 

driver/src/test/java/org/neo4j/driver/v1/integration/TLSSocketChannelFragmentationIT.java renamed to driver/src/test/java/org/neo4j/driver/v1/integration/TLSSocketChannelFragmentation.java

@@ -22,13 +22,8 @@
 import org.junit.Test;
 
 import java.io.IOException;
-import java.io.OutputStream;
-import java.net.InetSocketAddress;
-import java.net.ServerSocket;
-import java.net.Socket;
 import java.nio.ByteBuffer;
 import java.nio.channels.ByteChannel;
-import java.nio.channels.SocketChannel;
 import java.security.GeneralSecurityException;
 import java.security.KeyManagementException;
 import java.security.KeyStore;
@@ -39,35 +34,12 @@
 import java.security.cert.X509Certificate;
 import javax.net.ssl.KeyManagerFactory;
 import javax.net.ssl.SSLContext;
-import javax.net.ssl.SSLEngine;
-import javax.net.ssl.SSLServerSocketFactory;
 import javax.net.ssl.TrustManager;
 import javax.net.ssl.X509TrustManager;
 
-import org.neo4j.driver.internal.connector.socket.TLSSocketChannel;
-import org.neo4j.driver.internal.logging.DevNullLogger;
-
-import static org.hamcrest.core.IsEqual.equalTo;
-import static org.junit.Assert.assertThat;
-
-/**
- * This tests that the TLSSocketChannel handles every combination of network buffer sizes that we
- * can reasonably expect to see in the wild. It exhaustively tests power-of-two sizes up to 2^16
- * for the following variables:
- *
- * - Network frame size
- * - Bolt message size
- * - Read buffer size
- *
- * It tests every possible combination, and it does this currently only for the read path, expanding
- * to the write path as well would be useful. For each size, it sets up a TLS server and tests the
- * handshake, transferring the data, and verifying the data is correct after decryption.
- */
-public class TLSSocketChannelFragmentationIT
+public abstract class TLSSocketChannelFragmentation
 {
-    private SSLContext sslCtx;
-    private byte[] blobOfData;
-    private ServerSocket server;
+    protected SSLContext sslCtx;
 
     @Before
     public void setup() throws Throwable
@@ -76,17 +48,6 @@ public void setup() throws Throwable
         createServer();
     }
 
-    private void blobOfDataSize( int dataBlobSize )
-    {
-        blobOfData = new byte[dataBlobSize];
-        // If the blob is all zeros, we'd miss data corruption problems in assertions, so
-        // fill the data blob with different values.
-        for ( int i = 0; i < blobOfData.length; i++ )
-        {
-            blobOfData[i] = (byte) (i % 128);
-        }
-    }
-
     @Test
     public void shouldHandleFuzziness() throws Throwable
     {
@@ -109,34 +70,9 @@ public void shouldHandleFuzziness() throws Throwable
         }
     }
 
-    private void testForBufferSizes( int blobOfDataSize, int networkFrameSize, int userBufferSize ) throws IOException, GeneralSecurityException
-    {
-        blobOfDataSize(blobOfDataSize);
-        SSLEngine engine = sslCtx.createSSLEngine();
-        engine.setUseClientMode( true );
-        ByteChannel ch = SocketChannel.open( new InetSocketAddress( server.getInetAddress(), server.getLocalPort() ) );
-        ch = new LittleAtATimeChannel( ch, networkFrameSize );
-
-        TLSSocketChannel channel = new TLSSocketChannel(ch, new DevNullLogger(), engine);
-        try
-        {
-            ByteBuffer readBuffer = ByteBuffer.allocate( blobOfData.length );
-            while ( readBuffer.position() < readBuffer.capacity() )
-            {
-                readBuffer.limit(Math.min( readBuffer.capacity(), readBuffer.position() + userBufferSize ));
-                channel.read( readBuffer );
-            }
-
-            assertThat(readBuffer.array(), equalTo(blobOfData));
-        }
-        finally
-        {
-            channel.close();
-        }
-    }
-
-    private void createSSLContext()
-            throws KeyStoreException, IOException, NoSuchAlgorithmException, CertificateException, UnrecoverableKeyException, KeyManagementException
+    protected void createSSLContext()
+            throws KeyStoreException, IOException, NoSuchAlgorithmException, CertificateException,
+            UnrecoverableKeyException, KeyManagementException
     {
         KeyStore ks = KeyStore.getInstance("JKS");
         char[] password = "password".toCharArray();
@@ -159,40 +95,16 @@ public X509Certificate[] getAcceptedIssuers() {
         }}, null );
     }
 
-    private void createServer() throws IOException
-    {
-        SSLServerSocketFactory ssf = sslCtx.getServerSocketFactory();
-        server = ssf.createServerSocket(0);
+    protected abstract void testForBufferSizes( int blobOfDataSize, int networkFrameSize, int userBufferSize ) throws IOException,
+            GeneralSecurityException;
 
-        new Thread(new Runnable()
-        {
-            @Override
-            public void run()
-            {
-                try
-                {
-                    while(true)
-                    {
-                        Socket client = server.accept();
-                        OutputStream outputStream = client.getOutputStream();
-                        outputStream.write( blobOfData );
-                        outputStream.flush();
-                        // client.close(); // TODO: Uncomment this, fix resulting error handling CLOSED event
-                    }
-                }
-                catch ( IOException e )
-                {
-                    e.printStackTrace();
-                }
-            }
-        }).start();
-    }
+    protected abstract void createServer() throws IOException;
 
     /**
      * Delegates to underlying channel, but only reads up to the set amount at a time, used to emulate
      * different network frame sizes in this test.
      */
-    private static class LittleAtATimeChannel implements ByteChannel
+    protected static class LittleAtATimeChannel implements ByteChannel
     {
         private final ByteChannel delegate;
         private final int maxFrameSize;

driver/src/test/java/org/neo4j/driver/v1/integration/TLSSocketChannelReadFragmentationIT.java

@@ -0,0 +1,121 @@
+/**
+ * Copyright (c) 2002-2016 "Neo Technology,"
+ * Network Engine for Objects in Lund AB [http://neotechnology.com]
+ *
+ * This file is part of Neo4j.
+ *
+ * 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 org.neo4j.driver.v1.integration;
+
+import java.io.IOException;
+import java.io.OutputStream;
+import java.net.InetSocketAddress;
+import java.net.ServerSocket;
+import java.net.Socket;
+import java.nio.ByteBuffer;
+import java.nio.channels.ByteChannel;
+import java.nio.channels.SocketChannel;
+import java.security.GeneralSecurityException;
+import javax.net.ssl.SSLEngine;
+import javax.net.ssl.SSLServerSocketFactory;
+
+import org.neo4j.driver.internal.connector.socket.TLSSocketChannel;
+import org.neo4j.driver.internal.logging.DevNullLogger;
+
+import static org.hamcrest.core.IsEqual.equalTo;
+import static org.junit.Assert.assertThat;
+
+/**
+ * This tests that the TLSSocketChannel handles every combination of network buffer sizes that we
+ * can reasonably expect to see in the wild. It exhaustively tests power-of-two sizes up to 2^16
+ * for the following variables:
+ *
+ * - Network frame size
+ * - Bolt message size
+ * - Read buffer size
+ *
+ * It tests every possible combination, and it does this currently only for the read path, expanding
+ * to the write path as well would be useful. For each size, it sets up a TLS server and tests the
+ * handshake, transferring the data, and verifying the data is correct after decryption.
+ */
+public class TLSSocketChannelReadFragmentationIT extends TLSSocketChannelFragmentation
+{
+    private byte[] blobOfData;
+    private ServerSocket server;
+
+
+
+    private void blobOfDataSize( int dataBlobSize )
+    {
+        blobOfData = new byte[dataBlobSize];
+        // If the blob is all zeros, we'd miss data corruption problems in assertions, so
+        // fill the data blob with different values.
+        for ( int i = 0; i < blobOfData.length; i++ )
+        {
+            blobOfData[i] = (byte) (i % 128);
+        }
+    }
+
+    protected void testForBufferSizes( int blobOfDataSize, int networkFrameSize, int userBufferSize ) throws IOException, GeneralSecurityException
+    {
+        blobOfDataSize(blobOfDataSize);
+        SSLEngine engine = sslCtx.createSSLEngine();
+        engine.setUseClientMode( true );
+        ByteChannel ch = SocketChannel.open( new InetSocketAddress( server.getInetAddress(), server.getLocalPort() ) );
+        ch = new LittleAtATimeChannel( ch, networkFrameSize );
+
+        try ( TLSSocketChannel channel = new TLSSocketChannel( ch, new DevNullLogger(), engine ) )
+        {
+            ByteBuffer readBuffer = ByteBuffer.allocate( blobOfData.length );
+            while ( readBuffer.position() < readBuffer.capacity() )
+            {
+                readBuffer.limit( Math.min( readBuffer.capacity(), readBuffer.position() + userBufferSize ) );
+                channel.read( readBuffer );
+            }
+
+            assertThat( readBuffer.array(), equalTo( blobOfData ) );
+        }
+    }
+
+    protected void createServer() throws IOException
+    {
+        SSLServerSocketFactory ssf = sslCtx.getServerSocketFactory();
+        server = ssf.createServerSocket(0);
+
+        new Thread(new Runnable()
+        {
+            @Override
+            public void run()
+            {
+                try
+                {
+                    //noinspection InfiniteLoopStatement
+                    while(true)
+                    {
+                        Socket client = server.accept();
+                        OutputStream outputStream = client.getOutputStream();
+                        outputStream.write( blobOfData );
+                        outputStream.flush();
+                        // client.close(); // TODO: Uncomment this, fix resulting error handling CLOSED event
+                    }
+                }
+                catch ( IOException e )
+                {
+                    e.printStackTrace();
+                }
+            }
+        }).start();
+    }
+
+}