Use connect timeout in Bolt and TLS handshake

Previously configured connection timeout has only been applied to actual
`Socket#connect()` call. This is not the only thing driver does to
establish a connection. It also executes TLS handshake (if configured to
use encryption) and Bolt handshake to negotiate protocol version with
the database. Later two steps perform blocking reads without any
timeout.

This could be a problem when database does not respond in time. Also
driver might be simply connecting to something that is not the database
and never responds.

This commit makes driver use configured value of connect timeout as
blocking read timeout for TLS and Bolt handshakes. So both will not hang
forever when other side does not respond. Default value of connection
timeout is 5 seconds. With this commit driver will wait up to 5 seconds
for TLS and Bolt handshake.

Author: lutovich

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
     {