Skip to content

Commit 95d8eb0

Browse files
injectivesinjectives
authored
Add AuthToken to Bolt layer (#1609)
This is an internal improvement in the Bolt layer.
1 parent 54110d2 commit 95d8eb0

File tree

28 files changed

+346
-207
lines changed

28 files changed

+346
-207
lines changed

‎bolt-api-netty/src/main/java/org/neo4j/driver/internal/bolt/basicimpl/NettyBoltConnectionProvider.java‎

Lines changed: 13 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
import java.util.function.Consumer;
3232
import java.util.function.Supplier;
3333
import org.neo4j.driver.internal.bolt.api.AccessMode;
34+
import org.neo4j.driver.internal.bolt.api.AuthToken;
3435
import org.neo4j.driver.internal.bolt.api.BoltAgent;
3536
import org.neo4j.driver.internal.bolt.api.BoltConnection;
3637
import org.neo4j.driver.internal.bolt.api.BoltConnectionProvider;
@@ -44,7 +45,6 @@
4445
import org.neo4j.driver.internal.bolt.api.RoutingContext;
4546
import org.neo4j.driver.internal.bolt.api.SecurityPlan;
4647
import org.neo4j.driver.internal.bolt.api.exception.MinVersionAcquisitionException;
47-
import org.neo4j.driver.internal.bolt.api.values.Value;
4848
import org.neo4j.driver.internal.bolt.api.values.ValueFactory;
4949
import org.neo4j.driver.internal.bolt.basicimpl.impl.BoltConnectionImpl;
5050
import org.neo4j.driver.internal.bolt.basicimpl.impl.ConnectionProvider;
@@ -110,7 +110,7 @@ public CompletionStage<Void> init(
110110
public CompletionStage<BoltConnection> connect(
111111
SecurityPlan securityPlan,
112112
DatabaseName databaseName,
113-
Supplier<CompletionStage<Map<String, Value>>> authMapStageSupplier,
113+
Supplier<CompletionStage<AuthToken>> authTokenStageSupplier,
114114
AccessMode mode,
115115
Set<String> bookmarks,
116116
String impersonatedUser,
@@ -125,17 +125,17 @@ public CompletionStage<BoltConnection> connect(
125125
}
126126

127127
var latestAuthMillisFuture = new CompletableFuture<Long>();
128-
var authMapRef = new AtomicReference<Map<String, Value>>();
129-
return authMapStageSupplier
128+
var authMapRef = new AtomicReference<AuthToken>();
129+
return authTokenStageSupplier
130130
.get()
131-
.thenCompose(authMap -> {
132-
authMapRef.set(authMap);
131+
.thenCompose(authToken -> {
132+
authMapRef.set(authToken);
133133
return this.connectionProvider.acquireConnection(
134134
address,
135135
securityPlan,
136136
routingContext,
137137
databaseName != null ? databaseName.databaseName().orElse(null) : null,
138-
authMap,
138+
authToken.asMap(),
139139
boltAgent,
140140
userAgent,
141141
mode,
@@ -180,11 +180,11 @@ public CompletionStage<BoltConnection> connect(
180180
}
181181

182182
@Override
183-
public CompletionStage<Void> verifyConnectivity(SecurityPlan securityPlan, Map<String, Value> authMap) {
183+
public CompletionStage<Void> verifyConnectivity(SecurityPlan securityPlan, AuthToken authToken) {
184184
return connect(
185185
securityPlan,
186186
null,
187-
() -> CompletableFuture.completedStage(authMap),
187+
() -> CompletableFuture.completedStage(authToken),
188188
AccessMode.WRITE,
189189
Collections.emptySet(),
190190
null,
@@ -196,11 +196,11 @@ public CompletionStage<Void> verifyConnectivity(SecurityPlan securityPlan, Map<S
196196
}
197197

198198
@Override
199-
public CompletionStage<Boolean> supportsMultiDb(SecurityPlan securityPlan, Map<String, Value> authMap) {
199+
public CompletionStage<Boolean> supportsMultiDb(SecurityPlan securityPlan, AuthToken authToken) {
200200
return connect(
201201
securityPlan,
202202
null,
203-
() -> CompletableFuture.completedStage(authMap),
203+
() -> CompletableFuture.completedStage(authToken),
204204
AccessMode.WRITE,
205205
Collections.emptySet(),
206206
null,
@@ -215,11 +215,11 @@ public CompletionStage<Boolean> supportsMultiDb(SecurityPlan securityPlan, Map<S
215215
}
216216

217217
@Override
218-
public CompletionStage<Boolean> supportsSessionAuth(SecurityPlan securityPlan, Map<String, Value> authMap) {
218+
public CompletionStage<Boolean> supportsSessionAuth(SecurityPlan securityPlan, AuthToken authToken) {
219219
return connect(
220220
securityPlan,
221221
null,
222-
() -> CompletableFuture.completedStage(authMap),
222+
() -> CompletableFuture.completedStage(authToken),
223223
AccessMode.WRITE,
224224
Collections.emptySet(),
225225
null,

‎bolt-api-netty/src/main/java/org/neo4j/driver/internal/bolt/basicimpl/impl/BoltConnectionImpl.java‎

Lines changed: 10 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,8 @@
3333
import java.util.function.Function;
3434
import java.util.stream.Collectors;
3535
import org.neo4j.driver.internal.bolt.api.AccessMode;
36-
import org.neo4j.driver.internal.bolt.api.AuthData;
36+
import org.neo4j.driver.internal.bolt.api.AuthInfo;
37+
import org.neo4j.driver.internal.bolt.api.AuthToken;
3738
import org.neo4j.driver.internal.bolt.api.BoltConnection;
3839
import org.neo4j.driver.internal.bolt.api.BoltConnectionState;
3940
import org.neo4j.driver.internal.bolt.api.BoltProtocolVersion;
@@ -82,7 +83,7 @@ public final class BoltConnectionImpl implements BoltConnection {
8283
private final boolean telemetrySupported;
8384
private final boolean serverSideRouting;
8485
private final AtomicReference<BoltConnectionState> stateRef = new AtomicReference<>(BoltConnectionState.OPEN);
85-
private final AtomicReference<CompletableFuture<AuthData>> authDataRef;
86+
private final AtomicReference<CompletableFuture<AuthInfo>> authDataRef;
8687
private final Map<String, Value> routingContext;
8788
private final Queue<Function<ResponseHandler, CompletionStage<Void>>> messageWriters;
8889
private final Clock clock;
@@ -92,7 +93,7 @@ public BoltConnectionImpl(
9293
BoltProtocol protocol,
9394
Connection connection,
9495
EventLoop eventLoop,
95-
Map<String, Value> authMap,
96+
AuthToken authToken,
9697
CompletableFuture<Long> latestAuthMillisFuture,
9798
RoutingContext routingContext,
9899
Clock clock,
@@ -107,7 +108,7 @@ public BoltConnectionImpl(
107108
this.telemetrySupported = connection.isTelemetryEnabled();
108109
this.serverSideRouting = connection.isSsrEnabled();
109110
this.authDataRef = new AtomicReference<>(
110-
CompletableFuture.completedFuture(new AuthDataImpl(authMap, latestAuthMillisFuture.join())));
111+
CompletableFuture.completedFuture(new AuthInfoImpl(authToken, latestAuthMillisFuture.join())));
111112
this.valueFactory = Objects.requireNonNull(valueFactory);
112113
this.routingContext = routingContext.toMap().entrySet().stream()
113114
.collect(Collectors.toUnmodifiableMap(
@@ -369,10 +370,10 @@ public void onSummary(Void summary) {
369370
}
370371

371372
@Override
372-
public CompletionStage<BoltConnection> logon(Map<String, Value> authMap) {
373+
public CompletionStage<BoltConnection> logon(AuthToken authToken) {
373374
return executeInEventLoop(() -> messageWriters.add(handler -> protocol.logon(
374375
connection,
375-
authMap,
376+
authToken.asMap(),
376377
clock,
377378
new MessageHandler<>() {
378379
@Override
@@ -383,7 +384,7 @@ public void onError(Throwable throwable) {
383384

384385
@Override
385386
public void onSummary(Void summary) {
386-
authDataRef.get().complete(new AuthDataImpl(authMap, clock.millis()));
387+
authDataRef.get().complete(new AuthInfoImpl(authToken, clock.millis()));
387388
handler.onLogonSummary(null);
388389
}
389390
},
@@ -498,7 +499,7 @@ public BoltConnectionState state() {
498499
}
499500

500501
@Override
501-
public CompletionStage<AuthData> authData() {
502+
public CompletionStage<AuthInfo> authInfo() {
502503
return authDataRef.get();
503504
}
504505

@@ -572,7 +573,7 @@ private void updateState(Throwable throwable) {
572573
}
573574
}
574575

575-
private record AuthDataImpl(Map<String, Value> authMap, long authAckMillis) implements AuthData {}
576+
private record AuthInfoImpl(AuthToken authToken, long authAckMillis) implements AuthInfo {}
576577

577578
private static class ResponseHandleImpl implements ResponseHandler {
578579
private final ResponseHandler delegate;

‎bolt-api-pooled/src/main/java/org/neo4j/driver/internal/bolt/pooledimpl/PooledBoltConnectionProvider.java‎

Lines changed: 26 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@
3636
import java.util.function.Function;
3737
import java.util.function.Supplier;
3838
import org.neo4j.driver.internal.bolt.api.AccessMode;
39+
import org.neo4j.driver.internal.bolt.api.AuthToken;
3940
import org.neo4j.driver.internal.bolt.api.BasicResponseHandler;
4041
import org.neo4j.driver.internal.bolt.api.BoltAgent;
4142
import org.neo4j.driver.internal.bolt.api.BoltConnection;
@@ -51,7 +52,6 @@
5152
import org.neo4j.driver.internal.bolt.api.SecurityPlan;
5253
import org.neo4j.driver.internal.bolt.api.exception.BoltTransientException;
5354
import org.neo4j.driver.internal.bolt.api.exception.MinVersionAcquisitionException;
54-
import org.neo4j.driver.internal.bolt.api.values.Value;
5555
import org.neo4j.driver.internal.bolt.pooledimpl.impl.PooledBoltConnection;
5656
import org.neo4j.driver.internal.bolt.pooledimpl.impl.util.FutureUtil;
5757

@@ -129,7 +129,7 @@ public CompletionStage<Void> init(
129129
public CompletionStage<BoltConnection> connect(
130130
SecurityPlan securityPlan,
131131
DatabaseName databaseName,
132-
Supplier<CompletionStage<Map<String, Value>>> authMapStageSupplier,
132+
Supplier<CompletionStage<AuthToken>> authTokenStageSupplier,
133133
AccessMode mode,
134134
Set<String> bookmarks,
135135
String impersonatedUser,
@@ -145,7 +145,7 @@ public CompletionStage<BoltConnection> connect(
145145

146146
var acquisitionFuture = new CompletableFuture<PooledBoltConnection>();
147147

148-
authMapStageSupplier.get().whenComplete((authMap, authThrowable) -> {
148+
authTokenStageSupplier.get().whenComplete((authToken, authThrowable) -> {
149149
if (authThrowable != null) {
150150
acquisitionFuture.completeExceptionally(authThrowable);
151151
return;
@@ -168,8 +168,8 @@ public CompletionStage<BoltConnection> connect(
168168
acquisitionFuture,
169169
securityPlan,
170170
databaseName,
171-
authMap,
172-
authMapStageSupplier,
171+
authToken,
172+
authTokenStageSupplier,
173173
mode,
174174
bookmarks,
175175
impersonatedUser,
@@ -191,8 +191,8 @@ private void connect(
191191
CompletableFuture<PooledBoltConnection> acquisitionFuture,
192192
SecurityPlan securityPlan,
193193
DatabaseName databaseName,
194-
Map<String, Value> authMap,
195-
Supplier<CompletionStage<Map<String, Value>>> authMapStageSupplier,
194+
AuthToken authToken,
195+
Supplier<CompletionStage<AuthToken>> authTokenStageSupplier,
196196
AccessMode mode,
197197
Set<String> bookmarks,
198198
String impersonatedUser,
@@ -207,7 +207,7 @@ private void connect(
207207
empty.set(pooledConnectionEntries.isEmpty());
208208
try {
209209
// go over existing entries first
210-
connectionEntryWithMetadata = acquireExistingEntry(authMap, minVersion);
210+
connectionEntryWithMetadata = acquireExistingEntry(authToken, minVersion);
211211
} catch (MinVersionAcquisitionException e) {
212212
acquisitionFuture.completeExceptionally(e);
213213
return;
@@ -284,8 +284,8 @@ private void connect(
284284
acquisitionFuture,
285285
securityPlan,
286286
databaseName,
287-
authMap,
288-
authMapStageSupplier,
287+
authToken,
288+
authTokenStageSupplier,
289289
mode,
290290
bookmarks,
291291
impersonatedUser,
@@ -305,7 +305,7 @@ private void connect(
305305
purge(entry);
306306
metricsListener.afterConnectionReleased(poolId, inUseEvent);
307307
});
308-
reauthStage(entryWithMetadata, authMap).whenComplete((ignored2, throwable2) -> {
308+
reauthStage(entryWithMetadata, authToken).whenComplete((ignored2, throwable2) -> {
309309
if (!acquisitionFuture.complete(pooledConnection)) {
310310
// acquisition timed out
311311
CompletableFuture<PooledBoltConnection> pendingAcquisition;
@@ -336,7 +336,9 @@ private void connect(
336336
.connect(
337337
securityPlan,
338338
databaseName,
339-
empty.get() ? () -> CompletableFuture.completedStage(authMap) : authMapStageSupplier,
339+
empty.get()
340+
? () -> CompletableFuture.completedStage(authToken)
341+
: authTokenStageSupplier,
340342
mode,
341343
bookmarks,
342344
impersonatedUser,
@@ -395,7 +397,7 @@ private void connect(
395397
}
396398

397399
private synchronized ConnectionEntryWithMetadata acquireExistingEntry(
398-
Map<String, Value> authMap, BoltProtocolVersion minVersion) {
400+
AuthToken authToken, BoltProtocolVersion minVersion) {
399401
ConnectionEntryWithMetadata connectionEntryWithMetadata = null;
400402
var iterator = pooledConnectionEntries.iterator();
401403
while (iterator.hasNext()) {
@@ -431,10 +433,10 @@ private synchronized ConnectionEntryWithMetadata acquireExistingEntry(
431433
}
432434

433435
// the pool must not have unauthenticated connections
434-
var authData = connection.authData().toCompletableFuture().getNow(null);
436+
var authInfo = connection.authInfo().toCompletableFuture().getNow(null);
435437

436-
var expiredByError = minAuthTimestamp > 0 && authData.authAckMillis() <= minAuthTimestamp;
437-
var authMatches = authMap.equals(authData.authMap());
438+
var expiredByError = minAuthTimestamp > 0 && authInfo.authAckMillis() <= minAuthTimestamp;
439+
var authMatches = authToken.equals(authInfo.authToken());
438440
var reauthNeeded = expiredByError || !authMatches;
439441

440442
if (reauthNeeded) {
@@ -461,14 +463,14 @@ private synchronized ConnectionEntryWithMetadata acquireExistingEntry(
461463
}
462464

463465
private CompletionStage<Void> reauthStage(
464-
ConnectionEntryWithMetadata connectionEntryWithMetadata, Map<String, Value> authMap) {
466+
ConnectionEntryWithMetadata connectionEntryWithMetadata, AuthToken authToken) {
465467
CompletionStage<Void> stage;
466468
if (connectionEntryWithMetadata.reauthNeeded) {
467469
stage = connectionEntryWithMetadata
468470
.connectionEntry
469471
.connection
470472
.logoff()
471-
.thenCompose(conn -> conn.logon(authMap))
473+
.thenCompose(conn -> conn.logon(authToken))
472474
.handle((ignored, throwable) -> {
473475
if (throwable != null) {
474476
connectionEntryWithMetadata.connectionEntry.connection.close();
@@ -500,11 +502,11 @@ private CompletionStage<Void> livenessCheckStage(ConnectionEntry entry) {
500502
}
501503

502504
@Override
503-
public CompletionStage<Void> verifyConnectivity(SecurityPlan securityPlan, Map<String, Value> authMap) {
505+
public CompletionStage<Void> verifyConnectivity(SecurityPlan securityPlan, AuthToken authToken) {
504506
return connect(
505507
securityPlan,
506508
null,
507-
() -> CompletableFuture.completedStage(authMap),
509+
() -> CompletableFuture.completedStage(authToken),
508510
AccessMode.WRITE,
509511
Collections.emptySet(),
510512
null,
@@ -516,11 +518,11 @@ public CompletionStage<Void> verifyConnectivity(SecurityPlan securityPlan, Map<S
516518
}
517519

518520
@Override
519-
public CompletionStage<Boolean> supportsMultiDb(SecurityPlan securityPlan, Map<String, Value> authMap) {
521+
public CompletionStage<Boolean> supportsMultiDb(SecurityPlan securityPlan, AuthToken authToken) {
520522
return connect(
521523
securityPlan,
522524
null,
523-
() -> CompletableFuture.completedStage(authMap),
525+
() -> CompletableFuture.completedStage(authToken),
524526
AccessMode.WRITE,
525527
Collections.emptySet(),
526528
null,
@@ -535,11 +537,11 @@ public CompletionStage<Boolean> supportsMultiDb(SecurityPlan securityPlan, Map<S
535537
}
536538

537539
@Override
538-
public CompletionStage<Boolean> supportsSessionAuth(SecurityPlan securityPlan, Map<String, Value> authMap) {
540+
public CompletionStage<Boolean> supportsSessionAuth(SecurityPlan securityPlan, AuthToken authToken) {
539541
return connect(
540542
securityPlan,
541543
null,
542-
() -> CompletableFuture.completedStage(authMap),
544+
() -> CompletableFuture.completedStage(authToken),
543545
AccessMode.WRITE,
544546
Collections.emptySet(),
545547
null,

‎bolt-api-pooled/src/main/java/org/neo4j/driver/internal/bolt/pooledimpl/impl/PooledBoltConnection.java‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,8 @@
2323
import java.util.concurrent.CompletableFuture;
2424
import java.util.concurrent.CompletionStage;
2525
import org.neo4j.driver.internal.bolt.api.AccessMode;
26-
import org.neo4j.driver.internal.bolt.api.AuthData;
26+
import org.neo4j.driver.internal.bolt.api.AuthInfo;
27+
import org.neo4j.driver.internal.bolt.api.AuthToken;
2728
import org.neo4j.driver.internal.bolt.api.BasicResponseHandler;
2829
import org.neo4j.driver.internal.bolt.api.BoltConnection;
2930
import org.neo4j.driver.internal.bolt.api.BoltConnectionState;
@@ -162,8 +163,8 @@ public CompletionStage<BoltConnection> logoff() {
162163
}
163164

164165
@Override
165-
public CompletionStage<BoltConnection> logon(Map<String, Value> authMap) {
166-
return delegate.logon(authMap).thenApply(ignored -> this);
166+
public CompletionStage<BoltConnection> logon(AuthToken authToken) {
167+
return delegate.logon(authToken).thenApply(ignored -> this);
167168
}
168169

169170
@Override
@@ -321,8 +322,8 @@ public BoltConnectionState state() {
321322
}
322323

323324
@Override
324-
public CompletionStage<AuthData> authData() {
325-
return delegate.authData();
325+
public CompletionStage<AuthInfo> authInfo() {
326+
return delegate.authInfo();
326327
}
327328

328329
@Override

0 commit comments

Comments
 (0)