neo4j · zhenlineo · Oct 25, 2016 · Oct 25, 2016
diff --git a/driver/src/main/java/org/neo4j/driver/internal/security/SecurityPlan.java b/driver/src/main/java/org/neo4j/driver/internal/security/SecurityPlan.java
@@ -19,20 +19,18 @@
 
 package org.neo4j.driver.internal.security;
 
-import org.neo4j.driver.internal.net.BoltServerAddress;
-import org.neo4j.driver.v1.*;
-
-import javax.net.ssl.KeyManager;
-import javax.net.ssl.SSLContext;
-import javax.net.ssl.TrustManager;
-import javax.net.ssl.TrustManagerFactory;
 import java.io.File;
 import java.io.IOException;
 import java.security.GeneralSecurityException;
 import java.security.KeyManagementException;
 import java.security.KeyStore;
 import java.security.KeyStoreException;
 import java.security.NoSuchAlgorithmException;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import javax.net.ssl.KeyManager;
+import javax.net.ssl.SSLContext;
+import javax.net.ssl.TrustManagerFactory;
 
 import static org.neo4j.driver.internal.util.CertificateTool.loadX509Cert;
 
@@ -68,13 +66,34 @@ public static SecurityPlan forSystemCertificates() throws NoSuchAlgorithmExcepti
     }
 
 
-    public static SecurityPlan forTrustOnFirstUse( File knownHosts, BoltServerAddress address, Logger logger )
+    public static SecurityPlan forTrustOnFirstUse( File knownHosts )
             throws IOException, KeyManagementException, NoSuchAlgorithmException
     {
-        SSLContext sslContext = SSLContext.getInstance( "TLS" );
-        sslContext.init( new KeyManager[0], new TrustManager[]{new TrustOnFirstUseTrustManager( address, knownHosts, logger )}, null );
+        ConcurrentHashMap<String,String> preLoadKnownHostsMap = TrustOnFirstUseTrustManager.createKnownHostsMap( knownHosts );
+        return new TrustOnFirstUseSecurityPlan( knownHosts, preLoadKnownHostsMap );
+    }
 
-        return new SecurityPlan( true, sslContext);
+    public static class TrustOnFirstUseSecurityPlan extends SecurityPlan
+    {
+        private final ConcurrentHashMap<String, String> trustOnFirstUseMap;
+        private final File knownHostsFile;
+
+        private TrustOnFirstUseSecurityPlan( File knownHostsFile, ConcurrentHashMap<String, String> trustOnFirstUseMap )
+        {
+            super( true, null );
+            this.trustOnFirstUseMap = trustOnFirstUseMap;
+            this.knownHostsFile = knownHostsFile;
+        }
+
+        public ConcurrentHashMap<String, String> trustOnFirstUseMap()
+        {
+            return this.trustOnFirstUseMap;
+        }
+
+        public File knownHostFile()
+        {
+            return this.knownHostsFile;
+        }
     }
 
     public static SecurityPlan insecure()
@@ -85,7 +104,7 @@ public static SecurityPlan insecure()
     private final boolean requiresEncryption;
     private final SSLContext sslContext;
 
-    private SecurityPlan( boolean requiresEncryption, SSLContext sslContext)
+    private SecurityPlan( boolean requiresEncryption, SSLContext sslContext )
     {
         this.requiresEncryption = requiresEncryption;
         this.sslContext = sslContext;

diff --git a/driver/src/main/java/org/neo4j/driver/internal/security/TLSSocketChannel.java b/driver/src/main/java/org/neo4j/driver/internal/security/TLSSocketChannel.java
@@ -22,18 +22,18 @@
 import java.nio.ByteBuffer;
 import java.nio.channels.ByteChannel;
 import java.security.GeneralSecurityException;
+import java.security.KeyManagementException;
+import java.security.NoSuchAlgorithmException;
 import javax.net.ssl.KeyManager;
 import javax.net.ssl.SSLContext;
 import javax.net.ssl.SSLEngine;
 import javax.net.ssl.SSLEngineResult;
 import javax.net.ssl.SSLEngineResult.HandshakeStatus;
 import javax.net.ssl.SSLEngineResult.Status;
+import javax.net.ssl.TrustManager;
 
 import org.neo4j.driver.internal.net.BoltServerAddress;
-import org.neo4j.driver.v1.Logger;
-import org.neo4j.driver.internal.util.BytePrinter;
 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;
 
@@ -70,7 +70,7 @@ public class TLSSocketChannel implements ByteChannel
     public TLSSocketChannel( BoltServerAddress address, SecurityPlan securityPlan, ByteChannel channel, Logger logger )
             throws GeneralSecurityException, IOException
     {
-        this( channel, logger, createSSLEngine( address, securityPlan.sslContext() ) );
+        this( channel, logger, createSSLEngine( address, securityPlan, logger ) );
     }
 
     public TLSSocketChannel( ByteChannel channel, Logger logger, SSLEngine sslEngine ) throws GeneralSecurityException, IOException
@@ -356,10 +356,23 @@ static int bufferCopy( ByteBuffer from, ByteBuffer to )
     /**
      * Create SSLEngine with the SSLContext just created.
      * @param address the host to connect to
-     * @param sslContext the current ssl context
+     * @param securityPlan the security plan which holds the current ssl context
+     * @param logger the logger
      */
-    private static SSLEngine createSSLEngine( BoltServerAddress address, SSLContext sslContext )
+    private static SSLEngine createSSLEngine( BoltServerAddress address, SecurityPlan securityPlan, Logger logger )
+            throws IOException, KeyManagementException, NoSuchAlgorithmException
     {
+        SSLContext sslContext = securityPlan.sslContext();
+        if( securityPlan instanceof SecurityPlan.TrustOnFirstUseSecurityPlan )
+        {
+            // It require a new sslContext for each connection
+            sslContext = SSLContext.getInstance( "TLS" );
+            SecurityPlan.TrustOnFirstUseSecurityPlan plan = (SecurityPlan.TrustOnFirstUseSecurityPlan) securityPlan;
+            sslContext.init( new KeyManager[0], new TrustManager[]{
+                    new TrustOnFirstUseTrustManager( address.toString(), plan.knownHostFile(), plan.trustOnFirstUseMap(), logger )},
+                    null );
+        }
+
         SSLEngine sslEngine = sslContext.createSSLEngine( address.host(), address.port() );
         sslEngine.setUseClientMode( true );
         return sslEngine;

diff --git a/driver/src/main/java/org/neo4j/driver/internal/security/TrustOnFirstUseTrustManager.java b/driver/src/main/java/org/neo4j/driver/internal/security/TrustOnFirstUseTrustManager.java
@@ -28,11 +28,12 @@
 import java.security.NoSuchAlgorithmException;
 import java.security.cert.CertificateException;
 import java.security.cert.X509Certificate;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 import javax.net.ssl.X509TrustManager;
 
-import org.neo4j.driver.internal.net.BoltServerAddress;
-import org.neo4j.driver.v1.Logger;
 import org.neo4j.driver.internal.util.BytePrinter;
+import org.neo4j.driver.v1.Logger;
 
 import static java.lang.String.format;
 import static org.neo4j.driver.internal.util.CertificateTool.X509CertToString;
@@ -50,95 +51,101 @@ public class TrustOnFirstUseTrustManager implements X509TrustManager
      * Then when we try to connect to a known server again, we will authenticate the server by checking if it provides
      * the same certificate as the one saved in this file.
      */
-    private final File knownHosts;
+    private final File knownHostsFile;
 
-    /** The server ip:port (in digits) of the server that we are currently connected to */
-    private final String serverId;
+    /** The map of server ip:port (in digits) and its known certificate we've registered */
+    private final ConcurrentHashMap<String, String> knownHosts;
     private final Logger logger;
+    private final String serverId;
 
-    /** The known certificate we've registered for this server */
-    private String fingerprint;
-
-    TrustOnFirstUseTrustManager( BoltServerAddress address, File knownHosts, Logger logger ) throws IOException
+    TrustOnFirstUseTrustManager( String serverId, File knownHosts, ConcurrentHashMap<String,String> preLoadedKnownHosts, Logger logger )
+            throws IOException
     {
         this.logger = logger;
-        this.serverId = address.toString();
-        this.knownHosts = knownHosts;
-        load();
+        this.knownHostsFile = knownHosts;
+        this.knownHosts = preLoadedKnownHosts;
+        this.serverId = serverId;
+    }
+
+    public static ConcurrentHashMap<String, String> createKnownHostsMap( File knownHostsFile ) throws IOException
+    {
+        ConcurrentHashMap<String, String> knownHosts = new ConcurrentHashMap<>();
+        load( knownHostsFile, knownHosts );
+        return knownHosts;
     }
 
     /**
      * Try to load the certificate form the file if the server we've connected is a known server.
      *
      * @throws IOException
      */
-    private void load() throws IOException
+    private static synchronized void load( File knownHostsFile, Map<String, String> knownHosts ) throws IOException
     {
-        if ( !knownHosts.exists() )
+        if ( !knownHostsFile.exists() )
         {
             return;
         }
 
-        assertKnownHostFileReadable();
+        assertKnownHostFileReadable( knownHostsFile );
 
-        BufferedReader reader = new BufferedReader( new FileReader( knownHosts ) );
+        BufferedReader reader = new BufferedReader( new FileReader( knownHostsFile ) );
         String line;
         while ( (line = reader.readLine()) != null )
         {
             if ( (!line.trim().startsWith( "#" )) )
             {
                 String[] strings = line.split( " " );
-                if ( strings[0].trim().equals( serverId ) )
+                if(strings.length == 2)
                 {
-                    // load the certificate
-                    fingerprint = strings[1].trim();
-                    return;
+                    // we need to load all serverId and finger prints from the file as we do not know which one is
+                    // our current connection.
+                    knownHosts.put( strings[0], strings[1] );
                 }
             }
         }
         reader.close();
     }
 
     /**
-     * Save a new (server_ip, cert) pair into knownHosts file
+     * Save a new (server_ip, cert) pair into knownHostsFile file
      *
      * @param fingerprint the SHA-512 fingerprint of the host certificate
      */
-    private void saveTrustedHost( String fingerprint ) throws IOException
+    private synchronized void saveTrustedHost( String serverId, String fingerprint ) throws IOException
     {
-        this.fingerprint = fingerprint;
 
+        knownHosts.put( serverId, fingerprint );
         logger.warn( "Adding %s as known and trusted certificate for %s.", fingerprint, serverId );
         createKnownCertFileIfNotExists();
 
         assertKnownHostFileWritable();
-        BufferedWriter writer = new BufferedWriter( new FileWriter( knownHosts, true ) );
-        writer.write( serverId + " " + this.fingerprint );
+        BufferedWriter writer = new BufferedWriter( new FileWriter( knownHostsFile, true ) );
+        writer.write( serverId + " " + fingerprint );
         writer.newLine();
         writer.close();
     }
 
 
-    private void assertKnownHostFileReadable() throws IOException
+    private static void assertKnownHostFileReadable( File knownHostsFile ) throws IOException
     {
-        if( !knownHosts.canRead() )
+        if( !knownHostsFile.canRead() )
         {
             throw new IOException( format(
                     "Failed to load certificates from file %s as you have no read permissions to it.\n" +
                     "Try configuring the Neo4j driver to use a file system location you do have read permissions to.",
-                    knownHosts.getAbsolutePath()
+                    knownHostsFile.getAbsolutePath()
             ) );
         }
     }
 
     private void assertKnownHostFileWritable() throws IOException
     {
-        if( !knownHosts.canWrite() )
+        if( !knownHostsFile.canWrite() )
         {
             throw new IOException( format(
                     "Failed to write certificates to file %s as you have no write permissions to it.\n" +
                     "Try configuring the Neo4j driver to use a file system location you do have write permissions to.",
-                    knownHosts.getAbsolutePath()
+                    knownHostsFile.getAbsolutePath()
             ) );
         }
     }
@@ -161,24 +168,25 @@ public void checkServerTrusted( X509Certificate[] chain, String authType )
         X509Certificate certificate = chain[0];
 
         String cert = fingerprint( certificate );
+        String fingerprint = this.knownHosts.get( serverId );
 
-        if ( this.fingerprint == null )
+        if ( fingerprint == null )
         {
             try
             {
-                saveTrustedHost( cert );
+                saveTrustedHost( serverId, cert );
             }
             catch ( IOException e )
             {
                 throw new CertificateException( format(
                         "Failed to save the server ID and the certificate received from the server to file %s.\n" +
                         "Server ID: %s\nReceived cert:\n%s",
-                        knownHosts.getAbsolutePath(), serverId, X509CertToString( cert ) ), e );
+                        knownHostsFile.getAbsolutePath(), serverId, X509CertToString( cert ) ), e );
             }
         }
         else
         {
-            if ( !this.fingerprint.equals( cert ) )
+            if ( !fingerprint.equals( cert ) )
             {
                 throw new CertificateException( format(
                         "Unable to connect to neo4j at `%s`, because the certificate the server uses has changed. " +
@@ -187,8 +195,8 @@ public void checkServerTrusted( X509Certificate[] chain, String authType )
                         "`%s` " +
                         "in the file `%s`.\n" +
                         "The old certificate saved in file is:\n%sThe New certificate received is:\n%s",
-                        serverId, serverId, knownHosts.getAbsolutePath(),
-                        X509CertToString( this.fingerprint ), X509CertToString( cert ) ) );
+                        serverId, serverId, knownHostsFile.getAbsolutePath(),
+                        X509CertToString( fingerprint ), X509CertToString( cert ) ) );
             }
         }
     }
@@ -213,42 +221,46 @@ public static String fingerprint( X509Certificate cert ) throws CertificateExcep
 
     private File createKnownCertFileIfNotExists() throws IOException
     {
-        if ( !knownHosts.exists() )
+        if ( !knownHostsFile.exists() )
         {
-            File parentDir = knownHosts.getParentFile();
+            File parentDir = knownHostsFile.getParentFile();
             try
             {
                 if ( parentDir != null && !parentDir.exists() )
                 {
                     if ( !parentDir.mkdirs() )
                     {
-                        throw new IOException( "Failed to create directories for the known hosts file in " + knownHosts.getAbsolutePath() +
+                        throw new IOException( "Failed to create directories for the known hosts file in " + knownHostsFile
+
+                                .getAbsolutePath() +
                                                ". This is usually because you do not have write permissions to the directory. " +
                                                "Try configuring the Neo4j driver to use a file system location you do have write permissions to." );
                     }
                 }
-                if ( !knownHosts.createNewFile() )
+                if ( !knownHostsFile.createNewFile() )
                 {
-                    throw new IOException( "Failed to create a known hosts file at " + knownHosts.getAbsolutePath() +
+                    throw new IOException( "Failed to create a known hosts file at " + knownHostsFile
+                            .getAbsolutePath() +
                                            ". This is usually because you do not have write permissions to the directory. " +
                                            "Try configuring the Neo4j driver to use a file system location you do have write permissions to." );
                 }
             }
             catch( SecurityException e )
             {
-                throw new IOException( "Failed to create known host file and/or parent directories at " + knownHosts.getAbsolutePath() +
+                throw new IOException( "Failed to create known host file and/or parent directories at " + knownHostsFile
+                        .getAbsolutePath() +
                                        ". This is usually because you do not have write permission to the directory. " +
                                        "Try configuring the Neo4j driver to use a file location you have write permissions to." );
             }
-            BufferedWriter writer = new BufferedWriter( new FileWriter( knownHosts ) );
+            BufferedWriter writer = new BufferedWriter( new FileWriter( knownHostsFile ) );
             writer.write( "# This file contains trusted certificates for Neo4j servers, it's created by Neo4j drivers." );
             writer.newLine();
             writer.write( "# You can configure the location of this file in `org.neo4j.driver.Config`" );
             writer.newLine();
             writer.close();
         }
 
-        return knownHosts;
+        return knownHostsFile;
     }
 
     /**

diff --git a/driver/src/main/java/org/neo4j/driver/v1/GraphDatabase.java b/driver/src/main/java/org/neo4j/driver/v1/GraphDatabase.java
@@ -34,7 +34,6 @@
 import org.neo4j.driver.internal.spi.ConnectionPool;
 import org.neo4j.driver.internal.util.Clock;
 import org.neo4j.driver.v1.exceptions.ClientException;
-import org.neo4j.driver.v1.util.BiFunction;
 import org.neo4j.driver.v1.util.Function;
 
 import static java.lang.String.format;
@@ -222,8 +221,7 @@ private static SecurityPlan createSecurityPlan( BoltServerAddress address, Confi
             case TRUST_CUSTOM_CA_SIGNED_CERTIFICATES:
                 return SecurityPlan.forSignedCertificates( config.trustStrategy().certFile() );
             case TRUST_ON_FIRST_USE:
-                return SecurityPlan.forTrustOnFirstUse( config.trustStrategy().certFile(),
-                        address, logger );
+                return SecurityPlan.forTrustOnFirstUse( config.trustStrategy().certFile() );
             default:
                 throw new ClientException(
                         "Unknown TLS authentication strategy: " + config.trustStrategy().strategy().name() );