From 9dad697dbcaf6f4ac92da0d176668bca7265c097 Mon Sep 17 00:00:00 2001
From: Tsz-Wo Nicholas Sze <szetszwo@apache.org>
Date: Sat, 20 Jul 2024 15:18:22 +0800
Subject: [PATCH] HADOOP-19227. ipc.Server accelerate token negotiation only
 for the default mechanism. (#6949)

---
 .../src/main/java/org/apache/hadoop/ipc/Server.java   |  4 +++-
 .../org/apache/hadoop/security/SaslConstants.java     |  2 +-
 .../org/apache/hadoop/security/SaslRpcClient.java     | 11 ++++++++++-
 3 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/ipc/Server.java b/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/ipc/Server.java
index 0fb902a095..0d9e7296d2 100644
--- a/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/ipc/Server.java
+++ b/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/ipc/Server.java
@@ -106,6 +106,7 @@
 import org.apache.hadoop.ipc.protobuf.RpcHeaderProtos.RPCTraceInfoProto;
 import org.apache.hadoop.net.NetUtils;
 import org.apache.hadoop.security.AccessControlException;
+import org.apache.hadoop.security.SaslConstants;
 import org.apache.hadoop.security.SaslPropertiesResolver;
 import org.apache.hadoop.security.SaslRpcServer;
 import org.apache.hadoop.security.SaslRpcServer.AuthMethod;
@@ -2604,7 +2605,8 @@ private RpcSaslProto buildSaslNegotiateResponse()
       RpcSaslProto negotiateMessage = negotiateResponse;
       // accelerate token negotiation by sending initial challenge
       // in the negotiation response
-      if (enabledAuthMethods.contains(AuthMethod.TOKEN)) {
+      if (enabledAuthMethods.contains(AuthMethod.TOKEN)
+          && SaslConstants.SASL_MECHANISM_DEFAULT.equals(AuthMethod.TOKEN.getMechanismName())) {
         saslServer = createSaslServer(AuthMethod.TOKEN);
         byte[] challenge = saslServer.evaluateResponse(new byte[0]);
         RpcSaslProto.Builder negotiateBuilder =
diff --git a/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/security/SaslConstants.java b/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/security/SaslConstants.java
index 71e4b44873..6c6a2383ff 100644
--- a/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/security/SaslConstants.java
+++ b/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/security/SaslConstants.java
@@ -32,7 +32,7 @@ public class SaslConstants {
 
   private static final String SASL_MECHANISM_ENV = "HADOOP_SASL_MECHANISM";
   public static final String SASL_MECHANISM;
-  private static final String SASL_MECHANISM_DEFAULT = "DIGEST-MD5";
+  public static final String SASL_MECHANISM_DEFAULT = "DIGEST-MD5";
 
   static {
     final String mechanism = System.getenv(SASL_MECHANISM_ENV);
diff --git a/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/security/SaslRpcClient.java b/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/security/SaslRpcClient.java
index ce7878480e..aabe2ed4e0 100644
--- a/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/security/SaslRpcClient.java
+++ b/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/security/SaslRpcClient.java
@@ -39,6 +39,7 @@
 import javax.security.auth.callback.PasswordCallback;
 import javax.security.auth.callback.UnsupportedCallbackException;
 import javax.security.auth.kerberos.KerberosPrincipal;
+import javax.security.sasl.AuthorizeCallback;
 import javax.security.sasl.RealmCallback;
 import javax.security.sasl.RealmChoiceCallback;
 import javax.security.sasl.Sasl;
@@ -681,9 +682,17 @@ public void handle(Callback[] callbacks)
           pc = (PasswordCallback) callback;
         } else if (callback instanceof RealmCallback) {
           rc = (RealmCallback) callback;
+        } else if (callback instanceof AuthorizeCallback) {
+          final AuthorizeCallback ac = (AuthorizeCallback) callback;
+          final String authId = ac.getAuthenticationID();
+          final String authzId = ac.getAuthorizationID();
+          ac.setAuthorized(authId.equals(authzId));
+          if (ac.isAuthorized()) {
+            ac.setAuthorizedID(authzId);
+          }
         } else {
           throw new UnsupportedCallbackException(callback,
-              "Unrecognized SASL client callback");
+              "Unrecognized SASL client callback " + callback.getClass());
         }
       }
       if (nc != null) {