Merge pull request #442 from lutovich/1.4-handshake-with-timeout

Use connect timeout in Bolt and TLS handshake

Author: ali-ince

driver/src/main/java/org/neo4j/driver/internal/net/ChannelFactory.java

@@ -22,31 +22,26 @@
 import java.net.ConnectException;
 import java.net.Socket;
 import java.net.SocketTimeoutException;
-import java.net.StandardSocketOptions;
 import java.nio.channels.ByteChannel;
-import java.nio.channels.SocketChannel;
 
 import org.neo4j.driver.internal.security.SecurityPlan;
 import org.neo4j.driver.internal.security.TLSSocketChannel;
 import org.neo4j.driver.v1.Logger;
 
 class ChannelFactory
 {
-    static ByteChannel create( BoltServerAddress address, SecurityPlan securityPlan, int timeoutMillis, Logger log )
-            throws IOException
+    static ByteChannel create( Socket socket, BoltServerAddress address, SecurityPlan securityPlan, int timeoutMillis,
+            Logger log ) throws IOException
     {
-        SocketChannel soChannel = SocketChannel.open();
-        soChannel.setOption( StandardSocketOptions.SO_REUSEADDR, true );
-        soChannel.setOption( StandardSocketOptions.SO_KEEPALIVE, true );
-        connect( soChannel, address, timeoutMillis );
+        connect( socket, address, timeoutMillis );
 
-        ByteChannel channel = soChannel;
+        ByteChannel channel = new UnencryptedSocketChannel( socket );
 
         if ( securityPlan.requiresEncryption() )
         {
             try
             {
-                channel = TLSSocketChannel.create( address, securityPlan, soChannel, log );
+                channel = TLSSocketChannel.create( address, securityPlan, channel, log );
             }
             catch ( Exception e )
             {
@@ -70,10 +65,8 @@ static ByteChannel create( BoltServerAddress address, SecurityPlan securityPlan,
         return channel;
     }
 
-    private static void connect( SocketChannel soChannel, BoltServerAddress address, int timeoutMillis )
-            throws IOException
+    private static void connect( Socket socket, BoltServerAddress address, int timeoutMillis ) throws IOException
     {
-        Socket socket = soChannel.socket();
         try
         {
             socket.connect( address.toSocketAddress(), timeoutMillis );

driver/src/main/java/org/neo4j/driver/internal/net/SocketClient.java

@@ -20,14 +20,19 @@
 
 import java.io.IOException;
 import java.net.ConnectException;
+import java.net.Socket;
+import java.net.SocketAddress;
+import java.net.SocketTimeoutException;
 import java.nio.ByteBuffer;
 import java.nio.channels.ByteChannel;
 import java.util.Queue;
 
 import org.neo4j.driver.internal.messaging.Message;
 import org.neo4j.driver.internal.messaging.MessageFormat;
 import org.neo4j.driver.internal.security.SecurityPlan;
+import org.neo4j.driver.internal.security.TLSSocketChannel;
 import org.neo4j.driver.internal.util.BytePrinter;
+import org.neo4j.driver.v1.Config;
 import org.neo4j.driver.v1.Logger;
 import org.neo4j.driver.v1.exceptions.ClientException;
 import org.neo4j.driver.v1.exceptions.ServiceUnavailableException;
@@ -118,22 +123,31 @@ void blockingWrite( ByteBuffer buf ) throws IOException
 
     public void start()
     {
+        Socket socket = null;
+        boolean connected = false;
         try
         {
             logger.debug( "Connecting to %s, secure: %s", address, securityPlan.requiresEncryption() );
             if( channel == null )
             {
-                setChannel( ChannelFactory.create( address, securityPlan, timeoutMillis, logger ) );
+                socket = newSocket( timeoutMillis );
+                setChannel( ChannelFactory.create( socket, address, securityPlan, timeoutMillis, logger ) );
                 logger.debug( "Connected to %s, secure: %s", address, securityPlan.requiresEncryption() );
             }
 
             logger.debug( "Negotiating protocol with %s", address );
             SocketProtocol protocol = negotiateProtocol();
             setProtocol( protocol );
             logger.debug( "Selected protocol %s with %s", protocol.getClass(), address );
+
+            // reset read timeout (SO_TIMEOUT) to the original value of zero
+            // we do not want to permanently limit amount of time driver waits for database to execute query
+            socket.setSoTimeout( 0 );
+            connected = true;
         }
-        catch ( ConnectException e )
+        catch ( ConnectException | SocketTimeoutException e )
         {
+            // unable to connect socket or TLS/Bolt handshake took too much time
             throw new ServiceUnavailableException( format(
                     "Unable to connect to %s, ensure the database is running and that there is a " +
                     "working network connection to it.", address ), e );
@@ -142,6 +156,19 @@ public void start()
         {
             throw new ServiceUnavailableException( "Unable to process request: " + e.getMessage(), e );
         }
+        finally
+        {
+            if ( !connected && socket != null )
+            {
+                try
+                {
+                    socket.close();
+                }
+                catch ( Throwable ignore )
+                {
+                }
+            }
+        }
     }
 
     public void updateProtocol( String serverVersion )
@@ -301,4 +328,35 @@ public BoltServerAddress address()
     {
         return address;
     }
+
+    /**
+     * Creates new {@link Socket} object with {@link Socket#setSoTimeout(int) read timeout} set to the given value.
+     * Connection to bolt server includes:
+     * <ol>
+     * <li>TCP connect via {@link Socket#connect(SocketAddress, int)}</li>
+     * <li>Optional TLS handshake using {@link TLSSocketChannel}</li>
+     * <li>Bolt handshake</li>
+     * </ol>
+     * We do not want any of these steps to hang infinitely if server does not respond and thus:
+     * <ol>
+     * <li>Use {@link Socket#connect(SocketAddress, int)} with timeout, as configured in
+     * {@link Config#connectionTimeoutMillis()}</li>
+     * <li>Initially set {@link Socket#setSoTimeout(int) read timeout} on the socket. Same connection-timeout value
+     * from {@link Config#connectionTimeoutMillis()} is used. This way blocking reads during TLS and Bolt handshakes
+     * have limited waiting time</li>
+     * </ol>
+     *
+     * @param configuredConnectTimeout user-defined connection timeout to be initially used as read timeout.
+     * @return new socket.
+     * @throws IOException when creation or configuration of the socket fails.
+     */
+    private static Socket newSocket( int configuredConnectTimeout ) throws IOException
+    {
+        Socket socket = new Socket();
+        socket.setReuseAddress( true );
+        socket.setKeepAlive( true );
+        // set read timeout initially
+        socket.setSoTimeout( configuredConnectTimeout );
+        return socket;
+    }
 }

driver/src/main/java/org/neo4j/driver/internal/net/UnencryptedSocketChannel.java

@@ -0,0 +1,65 @@
+/*
+ * Copyright (c) 2002-2017 "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.internal.net;
+
+import java.io.IOException;
+import java.net.Socket;
+import java.nio.ByteBuffer;
+import java.nio.channels.ByteChannel;
+import java.nio.channels.Channels;
+import java.nio.channels.ReadableByteChannel;
+import java.nio.channels.WritableByteChannel;
+
+public class UnencryptedSocketChannel implements ByteChannel
+{
+    private final Socket socket;
+    private final ReadableByteChannel readableChannel;
+    private final WritableByteChannel writableChannel;
+
+    public UnencryptedSocketChannel( Socket socket ) throws IOException
+    {
+        this.socket = socket;
+        this.readableChannel = Channels.newChannel( socket.getInputStream() );
+        this.writableChannel = Channels.newChannel( socket.getOutputStream() );
+    }
+
+    @Override
+    public int read( ByteBuffer dst ) throws IOException
+    {
+        return readableChannel.read( dst );
+    }
+
+    @Override
+    public int write( ByteBuffer src ) throws IOException
+    {
+        return writableChannel.write( src );
+    }
+
+    @Override
+    public boolean isOpen()
+    {
+        return !socket.isClosed();
+    }
+
+    @Override
+    public void close() throws IOException
+    {
+        socket.close();
+    }
+}

driver/src/test/java/org/neo4j/driver/internal/net/SocketClientTest.java

@@ -18,24 +18,25 @@
  */
 package org.neo4j.driver.internal.net;
 
-import org.junit.Ignore;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.ExpectedException;
 
 import java.io.IOException;
 import java.net.ServerSocket;
+import java.net.SocketTimeoutException;
 import java.nio.ByteBuffer;
 import java.nio.channels.ByteChannel;
 import java.util.ArrayList;
 import java.util.List;
 
 import org.neo4j.driver.internal.security.SecurityPlan;
-import org.neo4j.driver.v1.exceptions.ClientException;
 import org.neo4j.driver.v1.exceptions.ServiceUnavailableException;
 
 import static org.hamcrest.CoreMatchers.equalTo;
-import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.instanceOf;
+import static org.junit.Assert.assertThat;
+import static org.junit.Assert.fail;
 import static org.mockito.Matchers.any;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
@@ -49,24 +50,16 @@ public class SocketClientTest
     @Rule
     public ExpectedException exception = ExpectedException.none();
 
-    // TODO: This is not possible with blocking NIO channels, unless we use inputStreams, but then we can't use
-    // off-heap buffers. We need to swap to use selectors, which would allow us to time out.
     @Test
-    @Ignore
-    public void testNetworkTimeout() throws Throwable
+    public void shouldFailWhenProtocolNegotiationTakesTooLong() throws Exception
     {
-        // Given a server that will never reply
-        ServerSocket server = new ServerSocket( 0 );
-        BoltServerAddress address = new BoltServerAddress( "localhost", server.getLocalPort() );
-
-        SocketClient client = dummyClient( address );
-
-        // Expect
-        exception.expect( ClientException.class );
-        exception.expectMessage( "database took longer than network timeout (100ms) to reply." );
+        testReadTimeoutOnConnect( SecurityPlan.insecure() );
+    }
 
-        // When
-        client.start();
+    @Test
+    public void shouldFailWhenTLSHandshakeTakesTooLong() throws Exception
+    {
+        testReadTimeoutOnConnect( SecurityPlan.forAllCertificates() );
     }
 
     @Test
@@ -190,6 +183,26 @@ public void shouldFailIfConnectionFailsWhileWriting() throws IOException
         client.blockingWrite( buffer );
     }
 
+    private static void testReadTimeoutOnConnect( SecurityPlan securityPlan ) throws IOException
+    {
+        try ( ServerSocket server = new ServerSocket( 0 ) ) // server that does not reply
+        {
+            int timeoutMillis = 1_000;
+            BoltServerAddress address = new BoltServerAddress( "localhost", server.getLocalPort() );
+            SocketClient client = new SocketClient( address, securityPlan, timeoutMillis, DEV_NULL_LOGGER );
+
+            try
+            {
+                client.start();
+                fail( "Exception expected" );
+            }
+            catch ( ServiceUnavailableException e )
+            {
+                assertThat( e.getCause(), instanceOf( SocketTimeoutException.class ) );
+            }
+        }
+    }
+
     private static class ByteAtATimeChannel implements ByteChannel
     {