review feedback

deephaven · Sep 10, 2024 · fa314e4 · fa314e4
1 parent 2963204
commit fa314e4
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 5 deletions.
diff --git a/server/src/main/java/io/deephaven/server/arrow/FlightServiceGrpcImpl.java b/server/src/main/java/io/deephaven/server/arrow/FlightServiceGrpcImpl.java
@@ -136,7 +136,12 @@ public void onNext(final Flight.HandshakeRequest value) {
                         try {
                             UUID uuid = UuidCreator.fromString(req.getPayload().toString(StandardCharsets.US_ASCII));
                             SessionState session = sessionService.getSessionForToken(uuid);
-                            respondWithAuthTokenBin(session);
+                            if (session != null) {
+                                SessionService.TokenExpiration expiration = session.getExpiration();
+                                if (expiration != null) {
+                                    respondWithAuthTokenBin(expiration);
+                                }
+                            }
                             return;
                         } catch (IllegalArgumentException | InvalidUuidException ignored) {
                         }
@@ -157,7 +162,7 @@ public void onNext(final Flight.HandshakeRequest value) {
             }
 
             SessionState session = sessionService.newSession(auth.get());
-            respondWithAuthTokenBin(session);
+            respondWithAuthTokenBin(session.getExpiration());
         }
 
         private Optional<AuthContext> login(String type, long version, ByteString payload,
@@ -171,10 +176,10 @@ private Optional<AuthContext> login(String type, long version, ByteString payloa
         }
 
         /** send the bearer token as an AuthTokenBin, as headers might have already been sent */
-        private void respondWithAuthTokenBin(SessionState session) {
+        private void respondWithAuthTokenBin(SessionService.TokenExpiration expiration) {
             isComplete = true;
             responseObserver.onNext(Flight.HandshakeResponse.newBuilder()
-                    .setPayload(session.getExpiration().getTokenAsByteString())
+                    .setPayload(expiration.getTokenAsByteString())
                     .build());
             responseObserver.onCompleted();
         }

diff --git a/server/src/main/java/io/deephaven/server/session/SessionServiceGrpcImpl.java b/server/src/main/java/io/deephaven/server/session/SessionServiceGrpcImpl.java
@@ -3,6 +3,7 @@
 //
 package io.deephaven.server.session;
 
+import com.github.f4b6a3.uuid.exception.InvalidUuidException;
 import com.google.protobuf.ByteString;
 import com.google.rpc.Code;
 import io.deephaven.auth.AuthContext;
@@ -345,7 +346,7 @@ public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(final ServerCall<Re
             if (altToken != null) {
                 try {
                     session = service.getSessionForToken(UUID.fromString(new String(altToken)));
-                } catch (IllegalArgumentException ignored) {
+                } catch (IllegalArgumentException | InvalidUuidException ignored) {
                 }
             }