Introduce impersonation support

This update brings support for the new impersonation feature. This feature allows creating sessions with impersonated user, including using default database of impersonated user.

Author: injectives

driver/src/main/java/org/neo4j/driver/SessionConfig.java

@@ -41,13 +41,15 @@ public class SessionConfig
     private final AccessMode defaultAccessMode;
     private final String database;
     private final Optional<Long> fetchSize;
+    private final String impersonatedUser;
 
     private SessionConfig( Builder builder )
     {
         this.bookmarks = builder.bookmarks;
         this.defaultAccessMode = builder.defaultAccessMode;
         this.database = builder.database;
         this.fetchSize = builder.fetchSize;
+        this.impersonatedUser = builder.impersonatedUser;
     }
 
     /**
@@ -116,13 +118,19 @@ public Optional<String> database()
 
     /**
      * This value if set, overrides the default fetch size set on {@link Config#fetchSize()}.
+     *
      * @return an optional value of fetch size.
      */
     public Optional<Long> fetchSize()
     {
         return fetchSize;
     }
 
+    public Optional<String> impersonatedUser()
+    {
+        return Optional.ofNullable( impersonatedUser );
+    }
+
     @Override
     public boolean equals( Object o )
     {
@@ -136,20 +144,20 @@ public boolean equals( Object o )
         }
         SessionConfig that = (SessionConfig) o;
         return Objects.equals( bookmarks, that.bookmarks ) && defaultAccessMode == that.defaultAccessMode && Objects.equals( database, that.database )
-                && Objects.equals( fetchSize, that.fetchSize );
+               && Objects.equals( fetchSize, that.fetchSize ) && Objects.equals( impersonatedUser, that.impersonatedUser );
     }
 
     @Override
     public int hashCode()
     {
-        return Objects.hash( bookmarks, defaultAccessMode, database );
+        return Objects.hash( bookmarks, defaultAccessMode, database, impersonatedUser );
     }
 
     @Override
     public String toString()
     {
         return "SessionParameters{" + "bookmarks=" + bookmarks + ", defaultAccessMode=" + defaultAccessMode + ", database='" + database + '\'' +
-                ", fetchSize=" + fetchSize + '}';
+               ", fetchSize=" + fetchSize + "impersonatedUser=" + impersonatedUser + '}';
     }
 
     /**
@@ -161,6 +169,7 @@ public static class Builder
         private Iterable<Bookmark> bookmarks = null;
         private AccessMode defaultAccessMode = AccessMode.WRITE;
         private String database = null;
+        private String impersonatedUser = null;
 
         private Builder()
         {
@@ -268,6 +277,31 @@ public Builder withFetchSize( long size )
             return this;
         }
 
+        /**
+         * Set the impersonated user that the newly created session is going to use for query execution.
+         * <p>
+         * The driver must have the necessary permissions to impersonate and run queries as the impersonated user.
+         * <p>
+         * When {@link #withDatabase(String)} is not used, the driver will discover the default database name of the impersonated user on first session usage.
+         * From that moment, the discovered database name will be used as the default database name for the whole lifetime of the new session.
+         * <p>
+         * <b>Compatible with 4.4+ only.</b> You MUST have all servers running 4.4 version or above and communicating over Bolt 4.4 or above.
+         *
+         * @param impersonatedUser the user to impersonate. Provided value should not be {@code null}.
+         * @return this builder
+         */
+        public Builder withImpersonatedUser( String impersonatedUser )
+        {
+            requireNonNull( impersonatedUser, "Impersonated user should not be null." );
+            if ( impersonatedUser.isEmpty() )
+            {
+                // Empty string is an illegal user. Fail fast on client.
+                throw new IllegalArgumentException( String.format( "Illegal impersonated user '%s'.", impersonatedUser ) );
+            }
+            this.impersonatedUser = impersonatedUser;
+            return this;
+        }
+
         public SessionConfig build()
         {
             return new SessionConfig( this );

driver/src/main/java/org/neo4j/driver/internal/DirectConnectionProvider.java

@@ -18,19 +18,21 @@
  */
 package org.neo4j.driver.internal;
 
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
 
 import org.neo4j.driver.internal.async.ConnectionContext;
 import org.neo4j.driver.internal.async.connection.DirectConnection;
 import org.neo4j.driver.internal.spi.Connection;
 import org.neo4j.driver.internal.spi.ConnectionPool;
 import org.neo4j.driver.internal.spi.ConnectionProvider;
+import org.neo4j.driver.internal.util.Futures;
 
+import static org.neo4j.driver.internal.async.ConnectionContext.PENDING_DATABASE_NAME_EXCEPTION_SUPPLIER;
 import static org.neo4j.driver.internal.messaging.request.MultiDatabaseUtil.supportsMultiDatabase;
 
 /**
- * Simple {@link ConnectionProvider connection provider} that obtains connections form the given pool only for
- * the given address.
+ * Simple {@link ConnectionProvider connection provider} that obtains connections form the given pool only for the given address.
  */
 public class DirectConnectionProvider implements ConnectionProvider
 {
@@ -46,7 +48,12 @@ public class DirectConnectionProvider implements ConnectionProvider
     @Override
     public CompletionStage<Connection> acquireConnection( ConnectionContext context )
     {
-        return acquireConnection().thenApply( connection -> new DirectConnection( connection, context.databaseName(), context.mode() ) );
+        CompletableFuture<DatabaseName> databaseNameFuture = context.databaseNameFuture();
+        databaseNameFuture.complete( DatabaseNameUtil.defaultDatabase() );
+        return acquireConnection().thenApply(
+                connection -> new DirectConnection( connection,
+                                                    Futures.joinNowOrElseThrow( databaseNameFuture, PENDING_DATABASE_NAME_EXCEPTION_SUPPLIER ),
+                                                    context.mode(), context.impersonatedUser() ) );
     }
 
     @Override

driver/src/main/java/org/neo4j/driver/internal/ImpersonationUtil.java

@@ -0,0 +1,46 @@
+/*
+ * Copyright (c) "Neo4j"
+ * Neo4j Sweden AB [http://neo4j.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;
+
+import org.neo4j.driver.exceptions.ClientException;
+import org.neo4j.driver.internal.messaging.v44.BoltProtocolV44;
+import org.neo4j.driver.internal.spi.Connection;
+import org.neo4j.driver.internal.util.ServerVersion;
+
+public class ImpersonationUtil
+{
+    public static final String IMPERSONATION_UNSUPPORTED_ERROR_MESSAGE =
+            "Detected connection that does not support impersonation, please make sure to have all servers running 4.4 version or above and communicating" +
+            " over Bolt version 4.4 or above when using impersonation feature";
+
+    public static Connection ensureImpersonationSupport( Connection connection, String impersonatedUser )
+    {
+        if ( impersonatedUser != null && !supportsImpersonation( connection ) )
+        {
+            throw new ClientException( IMPERSONATION_UNSUPPORTED_ERROR_MESSAGE );
+        }
+        return connection;
+    }
+
+    private static boolean supportsImpersonation( Connection connection )
+    {
+        return connection.serverVersion().greaterThanOrEqual( ServerVersion.v4_4_0 ) &&
+               connection.protocol().version().compareTo( BoltProtocolV44.VERSION ) >= 0;
+    }
+}

driver/src/main/java/org/neo4j/driver/internal/SessionFactoryImpl.java

@@ -52,7 +52,8 @@ public NetworkSession newInstance( SessionConfig sessionConfig )
     {
         BookmarkHolder bookmarkHolder = new DefaultBookmarkHolder( InternalBookmark.from( sessionConfig.bookmarks() ) );
         return createSession( connectionProvider, retryLogic, parseDatabaseName( sessionConfig ),
-                sessionConfig.defaultAccessMode(), bookmarkHolder, parseFetchSize( sessionConfig ), logging );
+                              sessionConfig.defaultAccessMode(), bookmarkHolder, parseFetchSize( sessionConfig ),
+                              sessionConfig.impersonatedUser().orElse( null ), logging );
     }
 
     private long parseFetchSize( SessionConfig sessionConfig )
@@ -98,10 +99,10 @@ public ConnectionProvider getConnectionProvider()
     }
 
     private NetworkSession createSession( ConnectionProvider connectionProvider, RetryLogic retryLogic, DatabaseName databaseName, AccessMode mode,
-            BookmarkHolder bookmarkHolder, long fetchSize, Logging logging )
+                                          BookmarkHolder bookmarkHolder, long fetchSize, String impersonatedUser, Logging logging )
     {
         return leakedSessionsLoggingEnabled
-               ? new LeakLoggingNetworkSession( connectionProvider, retryLogic, databaseName, mode, bookmarkHolder, fetchSize, logging )
-               : new NetworkSession( connectionProvider, retryLogic, databaseName, mode, bookmarkHolder, fetchSize, logging );
+               ? new LeakLoggingNetworkSession( connectionProvider, retryLogic, databaseName, mode, bookmarkHolder, impersonatedUser, fetchSize, logging )
+               : new NetworkSession( connectionProvider, retryLogic, databaseName, mode, bookmarkHolder, impersonatedUser, fetchSize, logging );
     }
 }

driver/src/main/java/org/neo4j/driver/internal/async/ConnectionContext.java

@@ -18,6 +18,9 @@
  */
 package org.neo4j.driver.internal.async;
 
+import java.util.concurrent.CompletableFuture;
+import java.util.function.Supplier;
+
 import org.neo4j.driver.AccessMode;
 import org.neo4j.driver.Bookmark;
 import org.neo4j.driver.internal.DatabaseName;
@@ -28,9 +31,13 @@
  */
 public interface ConnectionContext
 {
-    DatabaseName databaseName();
+    Supplier<IllegalStateException> PENDING_DATABASE_NAME_EXCEPTION_SUPPLIER = () -> new IllegalStateException( "Pending database name encountered." );
+
+    CompletableFuture<DatabaseName> databaseNameFuture();
 
     AccessMode mode();
 
     Bookmark rediscoveryBookmark();
+
+    String impersonatedUser();
 }

driver/src/main/java/org/neo4j/driver/internal/async/ImmutableConnectionContext.java

@@ -18,6 +18,8 @@
  */
 package org.neo4j.driver.internal.async;
 
+import java.util.concurrent.CompletableFuture;
+
 import org.neo4j.driver.AccessMode;
 import org.neo4j.driver.Bookmark;
 import org.neo4j.driver.internal.DatabaseName;
@@ -32,24 +34,26 @@
  */
 public class ImmutableConnectionContext implements ConnectionContext
 {
-    private static final ConnectionContext SINGLE_DB_CONTEXT = new ImmutableConnectionContext( defaultDatabase(), empty(), AccessMode.READ );
-    private static final ConnectionContext MULTI_DB_CONTEXT = new ImmutableConnectionContext( systemDatabase(), empty(), AccessMode.READ );
+    private static final ConnectionContext SINGLE_DB_CONTEXT = new ImmutableConnectionContext( defaultDatabase(), empty(), AccessMode.READ, null );
+    private static final ConnectionContext MULTI_DB_CONTEXT = new ImmutableConnectionContext( systemDatabase(), empty(), AccessMode.READ, null );
 
-    private final DatabaseName databaseName;
+    private final CompletableFuture<DatabaseName> databaseNameFuture;
     private final AccessMode mode;
     private final Bookmark rediscoveryBookmark;
+    private final String impersonatedUser;
 
-    public ImmutableConnectionContext( DatabaseName databaseName, Bookmark bookmark, AccessMode mode )
+    public ImmutableConnectionContext( DatabaseName databaseName, Bookmark bookmark, AccessMode mode, String impersonatedUser )
     {
-        this.databaseName = databaseName;
+        this.databaseNameFuture = CompletableFuture.completedFuture( databaseName );
         this.rediscoveryBookmark = bookmark;
         this.mode = mode;
+        this.impersonatedUser = impersonatedUser;
     }
 
     @Override
-    public DatabaseName databaseName()
+    public CompletableFuture<DatabaseName> databaseNameFuture()
     {
-        return databaseName;
+        return databaseNameFuture;
     }
 
     @Override
@@ -64,10 +68,15 @@ public Bookmark rediscoveryBookmark()
         return rediscoveryBookmark;
     }
 
+    @Override
+    public String impersonatedUser()
+    {
+        return impersonatedUser;
+    }
+
     /**
-     * A simple context is used to test connectivity with a remote server/cluster.
-     * As long as there is a read only service, the connection shall be established successfully.
-     * Depending on whether multidb is supported or not, this method returns different context for routing table discovery.
+     * A simple context is used to test connectivity with a remote server/cluster. As long as there is a read only service, the connection shall be established
+     * successfully. Depending on whether multidb is supported or not, this method returns different context for routing table discovery.
      */
     public static ConnectionContext simple( boolean supportsMultiDb )
     {

driver/src/main/java/org/neo4j/driver/internal/async/LeakLoggingNetworkSession.java

@@ -33,9 +33,9 @@ public class LeakLoggingNetworkSession extends NetworkSession
     private final String stackTrace;
 
     public LeakLoggingNetworkSession( ConnectionProvider connectionProvider, RetryLogic retryLogic, DatabaseName databaseName, AccessMode mode,
-            BookmarkHolder bookmarkHolder, long fetchSize, Logging logging )
+                                      BookmarkHolder bookmarkHolder, String impersonatedUser, long fetchSize, Logging logging )
     {
-        super( connectionProvider, retryLogic, databaseName, mode, bookmarkHolder, fetchSize, logging );
+        super( connectionProvider, retryLogic, databaseName, mode, bookmarkHolder, impersonatedUser, fetchSize, logging );
         this.stackTrace = captureStackTrace();
     }