Improve channel tracking for gRPC authentication traces

Main purpose of this PR is to make gRPC authentication tracing to
contain channel information in order to improve debugging.

pr-link: #9316
change-id: cid-bde726e6095f480cf91c1da586140d846c37ff6e

core/common/src/main/java/alluxio/grpc/GrpcChannel.java

@@ -70,7 +70,7 @@ public GrpcChannel(GrpcManagedChannelPool.ChannelKey channelKey, AuthenticatedCh
 
     // Store {@link AuthenticatedChannel::#close) for signaling end of
     // authenticated session during shutdown.
-    mAuthCloseCallback = ((AuthenticatedChannel) channel)::close;
+    mAuthCloseCallback = channel::close;
   }
 
   @Override

core/common/src/main/java/alluxio/security/authentication/ChannelAuthenticator.java

@@ -152,14 +152,14 @@ public void authenticate() throws AlluxioStatusException {
         SaslHandshakeClientHandler handshakeClient =
             new DefaultSaslHandshakeClientHandler(saslClientHandler);
         // Create driver for driving sasl traffic from client side.
-        mClientDriver =
-            new SaslStreamClientDriver(handshakeClient, mAuthenticated, mGrpcAuthTimeoutMs);
+        mClientDriver = new SaslStreamClientDriver(handshakeClient, mAuthenticated, mChannelId,
+            mGrpcAuthTimeoutMs);
         // Start authentication call with the service and update the client driver.
         StreamObserver<SaslMessage> requestObserver =
             SaslAuthenticationServiceGrpc.newStub(mManagedChannel).authenticate(mClientDriver);
         mClientDriver.setServerObserver(requestObserver);
         // Start authentication traffic with the target.
-        mClientDriver.start(mChannelId.toString());
+        mClientDriver.start();
         // Authentication succeeded!
         mManagedChannel.notifyWhenStateChanged(ConnectivityState.READY, () -> {
           mAuthenticated.set(false);

core/common/src/main/java/alluxio/security/authentication/DefaultSaslHandshakeClientHandler.java

@@ -20,6 +20,7 @@
 
 import javax.security.sasl.SaslClient;
 import javax.security.sasl.SaslException;
+import java.util.UUID;
 
 /**
  * Default implementation of {@link SaslHandshakeClientHandler}.
@@ -69,7 +70,7 @@ public SaslMessage handleSaslMessage(SaslMessage message) throws SaslException {
   }
 
   @Override
-  public SaslMessage getInitialMessage(String channelId) throws SaslException {
+  public SaslMessage getInitialMessage(UUID channelId) throws SaslException {
     byte[] initiateSaslResponse = null;
     if (mSaslClientHandler.getSaslClient().hasInitialResponse()) {
       initiateSaslResponse = mSaslClient.evaluateChallenge(S_INITIATE_CHALLENGE);
@@ -80,7 +81,7 @@ public SaslMessage getInitialMessage(String channelId) throws SaslException {
     if (initiateSaslResponse != null) {
       initialResponse.setMessage(ByteString.copyFrom(initiateSaslResponse));
     }
-    initialResponse.setClientId(channelId);
+    initialResponse.setClientId(channelId.toString());
     return initialResponse.build();
   }
 }

core/common/src/main/java/alluxio/security/authentication/SaslHandshakeClientHandler.java

@@ -14,6 +14,7 @@
 import alluxio.grpc.SaslMessage;
 
 import javax.security.sasl.SaslException;
+import java.util.UUID;
 
 /**
  * Interface for providing client-side handshake routines for a particular authentication scheme.
@@ -29,9 +30,9 @@ public interface SaslHandshakeClientHandler {
   public SaslMessage handleSaslMessage(SaslMessage message) throws SaslException;
 
   /**
-   * @param channelId channe for which the authentication is happening
+   * @param channelId channel for which the authentication is happening
    * @return the initial message for Sasl traffic to begin
    * @throws SaslException
    */
-  public SaslMessage getInitialMessage(String channelId) throws SaslException;
+  public SaslMessage getInitialMessage(UUID channelId) throws SaslException;
 }

core/common/src/main/java/alluxio/security/authentication/SaslStreamClientDriver.java

@@ -25,6 +25,7 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.util.UUID;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
@@ -58,6 +59,8 @@ public class SaslStreamClientDriver implements StreamObserver<SaslMessage> {
   private SettableFuture<Boolean> mHandshakeFuture;
   /** Current authenticated state. */
   private AtomicBoolean mAuthenticated;
+  /** Channel Id that has started the authentication. */
+  private UUID mChannelId;
 
   private final long mGrpcAuthTimeoutMs;
 
@@ -66,12 +69,14 @@ public class SaslStreamClientDriver implements StreamObserver<SaslMessage> {
    *
    * @param handshakeClient client handshake handler
    * @param authenticated boolean reference to receive authentication state changes
+   * @param channelId channel Id for authentication
    * @param grpcAuthTimeoutMs authentication timeout in milliseconds
    */
   public SaslStreamClientDriver(SaslHandshakeClientHandler handshakeClient,
-      AtomicBoolean authenticated, long grpcAuthTimeoutMs) {
+      AtomicBoolean authenticated, UUID channelId, long grpcAuthTimeoutMs) {
     mSaslHandshakeClientHandler = handshakeClient;
     mHandshakeFuture = SettableFuture.create();
+    mChannelId = channelId;
     mGrpcAuthTimeoutMs = grpcAuthTimeoutMs;
     mAuthenticated = authenticated;
   }
@@ -88,43 +93,44 @@ public void setServerObserver(StreamObserver<SaslMessage> requestObserver) {
   @Override
   public void onNext(SaslMessage saslMessage) {
     try {
-      LOG.debug("SaslClientDriver received message: {}",
-          saslMessage != null ? saslMessage.getMessageType().toString() : "<NULL>");
+      LOG.debug("SaslClientDriver received message: {} for channel: {}", saslMessage, mChannelId);
       SaslMessage response = mSaslHandshakeClientHandler.handleSaslMessage(saslMessage);
       if (response != null) {
         mRequestObserver.onNext(response);
       } else {
         // {@code null} response means server message was a success.
         mHandshakeFuture.set(true);
       }
-    } catch (SaslException e) {
+    } catch (Exception e) {
+      LOG.debug("Exception while handling SASL message: {} for channel: {}. Error: {}", saslMessage,
+          mChannelId, e);
       mHandshakeFuture.setException(e);
-      mRequestObserver
-          .onError(Status.fromCode(Status.Code.UNAUTHENTICATED).withCause(e).asException());
+      mRequestObserver.onError(e);
     }
   }
 
   @Override
   public void onError(Throwable throwable) {
+    LOG.warn("Received error on client driver for channel: {}. Error: {}", mChannelId, throwable);
     mHandshakeFuture.setException(throwable);
   }
 
   @Override
   public void onCompleted() {
+    LOG.debug("Client authentication closed by server for channel: {}", mChannelId);
     // Server completes the stream when authenticated session is terminated/revoked.
     mAuthenticated.set(false);
   }
 
   /**
    * Starts authentication with the server and wait until completion.
-   * @param channelId channel that is authenticating with the server
    * @throws UnauthenticatedException
    */
-  public void start(String channelId) throws AlluxioStatusException {
+  public void start() throws AlluxioStatusException {
     try {
-      LOG.debug("Starting SASL handshake for ChannelId:{}", channelId);
+      LOG.debug("Starting SASL handshake for channel: {}", mChannelId);
       // Send the server initial message.
-      mRequestObserver.onNext(mSaslHandshakeClientHandler.getInitialMessage(channelId));
+      mRequestObserver.onNext(mSaslHandshakeClientHandler.getInitialMessage(mChannelId));
       // Wait until authentication status changes.
       mAuthenticated.set(mHandshakeFuture.get(mGrpcAuthTimeoutMs, TimeUnit.MILLISECONDS));
     } catch (SaslException se) {
@@ -152,6 +158,7 @@ public void start(String channelId) throws AlluxioStatusException {
    * Stops authenticated session with the server by releasing the long poll.
    */
   public void stop() {
+    LOG.debug("Closing client driver for channel: {}", mChannelId);
     try {
       if (mAuthenticated.get()) {
         mRequestObserver.onCompleted();

core/common/src/main/java/alluxio/security/authentication/SaslStreamServerDriver.java

@@ -22,6 +22,7 @@
 
 import java.io.IOException;
 import java.util.UUID;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 import javax.security.sasl.SaslException;
 
@@ -30,16 +31,21 @@
  */
 public class SaslStreamServerDriver implements StreamObserver<SaslMessage> {
   private static final Logger LOG = LoggerFactory.getLogger(SaslStreamServerDriver.class);
+
+  /** Used to represent uninitialized channel Id. */
+  private static final UUID EMPTY_UUID = new UUID(0L, 0L);
   /** Client's sasl stream. */
   private StreamObserver<SaslMessage> mRequestObserver = null;
   /** Handshake handler for server. */
   private SaslHandshakeServerHandler mSaslHandshakeServerHandler;
   /** Authentication server. */
   private AuthenticationServer mAuthenticationServer;
   /** Id for client-side channel that is authenticating. */
-  private UUID mChannelId;
+  private UUID mChannelId = EMPTY_UUID;
   /** Sasl server handler that will be used for authentication. */
   private SaslServerHandler mSaslServerHandler = null;
+  /** Whether client stream observer is still valid. */
+  private AtomicBoolean mClientStreamValid = new AtomicBoolean(true);
 
   /**
    * Creates {@link SaslStreamServerDriver} for given {@link AuthenticationServer}.
@@ -64,12 +70,11 @@ public void onNext(SaslMessage saslMessage) {
     /** Whether to close the handler after this message.  */
     boolean closeHandler = false;
     try {
-      LOG.debug("SaslServerDriver received message: {}",
-          saslMessage != null ? saslMessage.getMessageType().toString() : "<NULL>");
-
       if (mSaslHandshakeServerHandler == null) {
         // First message received from the client.
         // ChannelId and the AuthenticationName will be set only in the first call.
+        LOG.debug("SaslServerDriver received authentication request of type:{} from channel: {}",
+            saslMessage.getAuthenticationScheme(), saslMessage.getClientId());
         // Initialize this server driver accordingly.
         mChannelId = UUID.fromString(saslMessage.getClientId());
         // Get authentication server to create the Sasl handler for requested scheme.
@@ -79,6 +84,8 @@ public void onNext(SaslMessage saslMessage) {
         // Unregister from registry if in case it was authenticated before.
         mAuthenticationServer.unregisterChannel(mChannelId);
       }
+
+      LOG.debug("SaslServerDriver received message: {} from channel: {}", saslMessage, mChannelId);
       // Respond to client.
       SaslMessage response = mSaslHandshakeServerHandler.handleSaslMessage(saslMessage);
       // Complete the call from server-side before sending success response to client.
@@ -92,15 +99,18 @@ public void onNext(SaslMessage saslMessage) {
       }
       mRequestObserver.onNext(response);
     } catch (SaslException se) {
-      LOG.debug("Exception while handling SASL message: {}", saslMessage, se);
+      LOG.debug("Exception while handling SASL message: {} for channel: {}. Error: {}", saslMessage,
+          mChannelId, se);
       mRequestObserver.onError(new UnauthenticatedException(se).toGrpcStatusException());
       closeHandler = true;
     } catch (UnauthenticatedException ue) {
-      LOG.debug("Exception while handling SASL message: {}", saslMessage, ue);
+      LOG.debug("Exception while handling SASL message: {} for channel: {}. Error: {}", saslMessage,
+          mChannelId, ue);
       mRequestObserver.onError(ue.toGrpcStatusException());
       closeHandler = true;
     } catch (Exception e) {
-      LOG.debug("Exception while handling SASL message: {}", saslMessage, e);
+      LOG.debug("Exception while handling SASL message: {} for channel: {}. Error: {}", saslMessage,
+          mChannelId, e);
       closeHandler = true;
       throw e;
     } finally {
@@ -116,7 +126,11 @@ public void onNext(SaslMessage saslMessage) {
 
   @Override
   public void onError(Throwable throwable) {
-    if (mChannelId != null) {
+    LOG.warn("Error received for channel: {}. Error: {}", mChannelId, throwable);
+    // Error on server invalidates client stream.
+    mClientStreamValid.set(false);
+
+    if (!mChannelId.equals(EMPTY_UUID)) {
       LOG.debug("Closing authenticated channel: {} due to error: {}", mChannelId, throwable);
       if (!mAuthenticationServer.unregisterChannel(mChannelId)) {
         // Channel was not registered. Close driver explicitly.
@@ -140,6 +154,7 @@ public void onCompleted() {
    * Closes the authentication stream.
    */
   public void close() {
+    LOG.debug("Closing server driver for channel: {}", mChannelId);
     // Complete the client stream.
     completeStreamQuietly();
     // Close handler if not already.
@@ -148,7 +163,7 @@ public void close() {
         mSaslServerHandler.close();
       } catch (Exception exc) {
         LogUtils.warnWithException(LOG, "Failed to close server driver for channel: {}.",
-            (mChannelId != null) ? mChannelId : "<NULL>", exc);
+            mChannelId, exc);
       }
     }
   }
@@ -157,11 +172,14 @@ public void close() {
    * Completes the stream with a debug blanket over possible exceptions.
    */
   private void completeStreamQuietly() {
-    if (mRequestObserver != null) {
+    if (mClientStreamValid.get() && mRequestObserver != null) {
       try {
         mRequestObserver.onCompleted();
       } catch (Exception exc) {
-        LOG.debug("Failed to close authentication stream from server.", exc);
+        LOG.debug("Failed to close authentication stream for channel: {}. Error: {}", mChannelId,
+            exc);
+      } finally {
+        mClientStreamValid.set(false);
       }
     }
   }