From 13414be1a1d502e608818a9d7840880d794c8ab7 Mon Sep 17 00:00:00 2001
From: Zhe Zhang <zhz@apache.org>
Date: Fri, 4 Mar 2016 09:48:05 -0800
Subject: [PATCH] HDFS-9888. Allow reseting KerberosName in unit tests.
 Contributed by Xiao Chen.

(cherry picked from commit 3e8099a45a4cfd4c5c0e3dce4370514cb2c90da9)
(cherry picked from commit 00ff3d737c71e1a228ede64a6b9abd5227ef2117)
(cherry picked from commit 9c1c48c0d44a45c4f96e8d60cfa6453a85523905)
---
 .../authentication/util/KerberosName.java       | 17 ++++++++++++++---
 .../hadoop/security/UserGroupInformation.java   |  2 +-
 .../hdfs/server/balancer/TestBalancer.java      | 12 +++++++-----
 3 files changed, 22 insertions(+), 9 deletions(-)

diff --git a/hadoop-common-project/hadoop-auth/src/main/java/org/apache/hadoop/security/authentication/util/KerberosName.java b/hadoop-common-project/hadoop-auth/src/main/java/org/apache/hadoop/security/authentication/util/KerberosName.java
index 7ae8ab2672e..0bc1109e512 100644
--- a/hadoop-common-project/hadoop-auth/src/main/java/org/apache/hadoop/security/authentication/util/KerberosName.java
+++ b/hadoop-common-project/hadoop-auth/src/main/java/org/apache/hadoop/security/authentication/util/KerberosName.java
@@ -25,6 +25,7 @@
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
+import com.google.common.annotations.VisibleForTesting;
 import org.apache.hadoop.classification.InterfaceAudience;
 import org.apache.hadoop.classification.InterfaceStability;
 import org.slf4j.Logger;
@@ -90,6 +91,16 @@ public class KerberosName {
     }
   }
 
+  @VisibleForTesting
+  public static void resetDefaultRealm() {
+    try {
+      defaultRealm = KerberosUtil.getDefaultRealm();
+    } catch (Exception ke) {
+      LOG.debug("resetting default realm failed, "
+          + "current default realm will still be used.", ke);
+    }
+  }
+
   /**
    * Create a name from the full Kerberos principal name.
    * @param name full Kerberos principal name.
@@ -412,16 +423,16 @@ public static String getRules() {
     }
     return ruleString;
   }
-  
+
   /**
    * Indicates if the name rules have been set.
-   * 
+   *
    * @return if the name rules have been set.
    */
   public static boolean hasRulesBeenSet() {
     return rules != null;
   }
-  
+
   static void printRules() throws IOException {
     int i = 0;
     for(Rule r: rules) {
diff --git a/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/security/UserGroupInformation.java b/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/security/UserGroupInformation.java
index 7f4ae968e7a..059779599fc 100644
--- a/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/security/UserGroupInformation.java
+++ b/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/security/UserGroupInformation.java
@@ -314,7 +314,7 @@ public static void setConfiguration(Configuration conf) {
   
   @InterfaceAudience.Private
   @VisibleForTesting
-  static void reset() {
+  public static void reset() {
     authenticationMethod = null;
     conf = null;
     groups = null;
diff --git a/hadoop-hdfs-project/hadoop-hdfs/src/test/java/org/apache/hadoop/hdfs/server/balancer/TestBalancer.java b/hadoop-hdfs-project/hadoop-hdfs/src/test/java/org/apache/hadoop/hdfs/server/balancer/TestBalancer.java
index 5ddda0b9328..6713f3e4286 100644
--- a/hadoop-hdfs-project/hadoop-hdfs/src/test/java/org/apache/hadoop/hdfs/server/balancer/TestBalancer.java
+++ b/hadoop-hdfs-project/hadoop-hdfs/src/test/java/org/apache/hadoop/hdfs/server/balancer/TestBalancer.java
@@ -93,6 +93,7 @@
 import org.apache.hadoop.http.HttpConfig;
 import org.apache.hadoop.io.IOUtils;
 import org.apache.hadoop.minikdc.MiniKdc;
+import org.apache.hadoop.security.authentication.util.KerberosName;
 import org.apache.hadoop.security.SecurityUtil;
 import org.apache.hadoop.security.UserGroupInformation;
 import org.apache.hadoop.security.ssl.KeyStoreTestUtil;
@@ -139,6 +140,7 @@ static void initSecureConf(Configuration conf) throws Exception {
     SecurityUtil.setAuthenticationMethod(
         UserGroupInformation.AuthenticationMethod.KERBEROS, conf);
     UserGroupInformation.setConfiguration(conf);
+    KerberosName.resetDefaultRealm();
     assertTrue("Expected configuration to enable security",
         UserGroupInformation.isSecurityEnabled());
 
@@ -270,7 +272,7 @@ static Block[][] distributeBlocks(ExtendedBlock[] blocks,
     long[] usedSpace = new long[distribution.length];
     System.arraycopy(distribution, 0, usedSpace, 0, distribution.length);
 
-    List<List<Block>> blockReports = 
+    List<List<Block>> blockReports =
       new ArrayList<List<Block>>(usedSpace.length);
     Block[][] results = new Block[usedSpace.length][];
     for(int i=0; i<usedSpace.length; i++) {
@@ -327,7 +329,7 @@ private void testUnevenDistribution(Configuration conf,
         blocks, (short)(numDatanodes-1), distribution);
 
     // restart the cluster: do NOT format the cluster
-    conf.set(DFSConfigKeys.DFS_NAMENODE_SAFEMODE_THRESHOLD_PCT_KEY, "0.0f"); 
+    conf.set(DFSConfigKeys.DFS_NAMENODE_SAFEMODE_THRESHOLD_PCT_KEY, "0.0f");
     cluster = new MiniDFSCluster.Builder(conf).numDataNodes(numDatanodes)
                                               .format(false)
                                               .racks(racks)
@@ -1556,10 +1558,10 @@ public void testManyBalancerSimultaneously() throws Exception {
   @Test(timeout = 300000)
   public void testBalancerWithKeytabs() throws Exception {
     final Configuration conf = new HdfsConfiguration();
-    initSecureConf(conf);
-    final UserGroupInformation ugi = UserGroupInformation.loginUserFromKeytabAndReturnUGI(
-        principal, keytabFile.getAbsolutePath());
     try {
+      initSecureConf(conf);
+      final UserGroupInformation ugi = UserGroupInformation.loginUserFromKeytabAndReturnUGI(
+          principal, keytabFile.getAbsolutePath());
       ugi.doAs(new PrivilegedExceptionAction<Void>() {
         @Override
         public Void run() throws Exception {