From d3fdbbc2ab1dcefcee49f8468f33b0e694b305cf Mon Sep 17 00:00:00 2001
From: Sunil G <sunilg@apache.org>
Date: Fri, 4 Jan 2019 21:49:06 +0530
Subject: [PATCH] HADOOP-15997. KMS client uses wrong UGI after HADOOP-14445.
 Contributed by Wei-Chiu Chuang.

(cherry picked from commit 51427cbdfb39cb6f5774b7b70009d7ee4388edfc)
---
 .../crypto/key/kms/KMSClientProvider.java     | 13 +++---
 .../TestLoadBalancingKMSClientProvider.java   | 45 +++++++++++++++++++
 2 files changed, 51 insertions(+), 7 deletions(-)

diff --git a/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/crypto/key/kms/KMSClientProvider.java b/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/crypto/key/kms/KMSClientProvider.java
index f082a15ff8e..689442fbf70 100644
--- a/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/crypto/key/kms/KMSClientProvider.java
+++ b/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/crypto/key/kms/KMSClientProvider.java
@@ -1147,17 +1147,16 @@ public class KMSClientProvider extends KeyProvider implements CryptoExtension,
     // Add existing credentials from the UGI, since provider is cached.
     Credentials creds = ugi.getCredentials();
     if (!creds.getAllTokens().isEmpty()) {
-      LOG.debug("Searching for token that matches service: {}", dtService);
-      org.apache.hadoop.security.token.Token<? extends TokenIdentifier>
-          dToken = creds.getToken(dtService);
-      if (dToken != null) {
-        return true;
-      }
+      LOG.debug("Searching for KMS delegation token in user {}'s credentials",
+          ugi);
+      return clientTokenProvider.selectDelegationToken(creds) != null;
     }
+
     return false;
   }
 
-  private UserGroupInformation getActualUgi() throws IOException {
+  @VisibleForTesting
+  UserGroupInformation getActualUgi() throws IOException {
     final UserGroupInformation currentUgi = UserGroupInformation
         .getCurrentUser();
 
diff --git a/hadoop-common-project/hadoop-common/src/test/java/org/apache/hadoop/crypto/key/kms/TestLoadBalancingKMSClientProvider.java b/hadoop-common-project/hadoop-common/src/test/java/org/apache/hadoop/crypto/key/kms/TestLoadBalancingKMSClientProvider.java
index c99c63dfffd..259fedad8fe 100644
--- a/hadoop-common-project/hadoop-common/src/test/java/org/apache/hadoop/crypto/key/kms/TestLoadBalancingKMSClientProvider.java
+++ b/hadoop-common-project/hadoop-common/src/test/java/org/apache/hadoop/crypto/key/kms/TestLoadBalancingKMSClientProvider.java
@@ -18,6 +18,7 @@
 package org.apache.hadoop.crypto.key.kms;
 
 import static org.apache.hadoop.crypto.key.KeyProviderCryptoExtension.EncryptedKeyVersion;
+import static org.apache.hadoop.crypto.key.kms.KMSDelegationToken.TOKEN_KIND;
 import static org.apache.hadoop.test.LambdaTestUtils.intercept;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotEquals;
@@ -35,6 +36,7 @@ import java.net.URI;
 import java.net.UnknownHostException;
 import java.security.GeneralSecurityException;
 import java.security.NoSuchAlgorithmException;
+import java.security.PrivilegedExceptionAction;
 import java.util.Arrays;
 import java.util.List;
 
@@ -45,11 +47,15 @@ import org.apache.hadoop.crypto.key.KeyProvider;
 import org.apache.hadoop.crypto.key.KeyProvider.Options;
 import org.apache.hadoop.crypto.key.KeyProviderCryptoExtension;
 import org.apache.hadoop.fs.CommonConfigurationKeysPublic;
+import org.apache.hadoop.io.Text;
 import org.apache.hadoop.net.ConnectTimeoutException;
 import org.apache.hadoop.security.AccessControlException;
+import org.apache.hadoop.security.Credentials;
 import org.apache.hadoop.security.SecurityUtil;
+import org.apache.hadoop.security.UserGroupInformation;
 import org.apache.hadoop.security.authentication.client.AuthenticationException;
 import org.apache.hadoop.security.authorize.AuthorizationException;
+import org.apache.hadoop.security.token.Token;
 import org.junit.BeforeClass;
 import org.junit.Rule;
 import org.junit.Test;
@@ -909,4 +915,43 @@ public class TestLoadBalancingKMSClientProvider {
       assertNotEquals(kmsUri, providers[i].getCanonicalServiceName());
     }
   }
+
+  @Test
+  public void testGetActualUGI() throws Exception {
+    // enable security
+    final Configuration conf = new Configuration();
+    conf.set("hadoop.security.authentication", "kerberos");
+    UserGroupInformation.setConfiguration(conf);
+
+    UserGroupInformation ugi = UserGroupInformation.createUserForTesting(
+        "foo", new String[] {"hadoop"});
+
+    String providerUriString = "kms://http@host1;host2;host3:9600/kms/foo";
+    final URI kmsUri = URI.create(providerUriString);
+    // create a fake kms dt
+    final Token token = new Token();
+    token.setKind(TOKEN_KIND);
+    token.setService(new Text(providerUriString));
+    // call getActualUgi() with the current user.
+    UserGroupInformation actualUgi =
+        ugi.doAs(new PrivilegedExceptionAction<UserGroupInformation>(){
+          @Override
+          public UserGroupInformation run() throws Exception {
+            final KeyProvider kp =
+                new KMSClientProvider.Factory().createProvider(kmsUri, conf);
+            final LoadBalancingKMSClientProvider lbkp =
+                (LoadBalancingKMSClientProvider) kp;
+            final Credentials creds = new Credentials();
+            creds.addToken(token.getService(), token);
+            UserGroupInformation.getCurrentUser().addCredentials(creds);
+
+            KMSClientProvider[] providers = lbkp.getProviders();
+            return providers[0].getActualUgi();
+          }
+        });
+    // make sure getActualUgi() returns the current user, not login user.
+    assertEquals(
+        "getActualUgi() should return the current user, not login user",
+        ugi, actualUgi);
+  }
 }
\ No newline at end of file