From 8153fe2bd35fb4df0b64f93ac0046e34d4807ac3 Mon Sep 17 00:00:00 2001
From: Wei-Chiu Chuang <weichiu@apache.org>
Date: Fri, 7 Jul 2017 06:13:10 -0700
Subject: [PATCH] HADOOP-14563.
 LoadBalancingKMSClientProvider#warmUpEncryptedKeys swallows IOException.
 Contributed by Rushabh S Shah.

---
 .../kms/LoadBalancingKMSClientProvider.java   | 12 +++-
 .../TestLoadBalancingKMSClientProvider.java   | 63 +++++++++++++++++++
 2 files changed, 74 insertions(+), 1 deletion(-)

diff --git a/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/crypto/key/kms/LoadBalancingKMSClientProvider.java b/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/crypto/key/kms/LoadBalancingKMSClientProvider.java
index aa2499343d..de9c988dbf 100644
--- a/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/crypto/key/kms/LoadBalancingKMSClientProvider.java
+++ b/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/crypto/key/kms/LoadBalancingKMSClientProvider.java
@@ -39,6 +39,7 @@
 import org.slf4j.LoggerFactory;
 
 import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Preconditions;
 
 /**
  * A simple LoadBalancing KMSClientProvider that round-robins requests
@@ -159,15 +160,24 @@ public Void call(KMSClientProvider provider) throws IOException {
   // This request is sent to all providers in the load-balancing group
   @Override
   public void warmUpEncryptedKeys(String... keyNames) throws IOException {
+    Preconditions.checkArgument(providers.length > 0,
+        "No providers are configured");
+    boolean success = false;
+    IOException e = null;
     for (KMSClientProvider provider : providers) {
       try {
         provider.warmUpEncryptedKeys(keyNames);
+        success = true;
       } catch (IOException ioe) {
+        e = ioe;
         LOG.error(
             "Error warming up keys for provider with url"
-            + "[" + provider.getKMSUrl() + "]");
+            + "[" + provider.getKMSUrl() + "]", ioe);
       }
     }
+    if (!success && e != null) {
+      throw e;
+    }
   }
 
   // This request is sent to all providers in the load-balancing group
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 499b9915d9..d14dd59c77 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
@@ -34,6 +34,7 @@
 import org.apache.hadoop.crypto.key.KeyProvider.Options;
 import org.apache.hadoop.crypto.key.KeyProviderCryptoExtension;
 import org.apache.hadoop.security.authentication.client.AuthenticationException;
+import org.apache.hadoop.security.authorize.AuthorizationException;
 import org.junit.Test;
 import org.mockito.Mockito;
 
@@ -257,4 +258,66 @@ public void testClassCastException() throws Exception {
           "AuthenticationException"));
     }
   }
+
+  /**
+   * tests {@link LoadBalancingKMSClientProvider#warmUpEncryptedKeys(String...)}
+   * error handling in case when all the providers throws {@link IOException}.
+   * @throws Exception
+   */
+  @Test
+  public void testWarmUpEncryptedKeysWhenAllProvidersFail() throws Exception {
+    Configuration conf = new Configuration();
+    KMSClientProvider p1 = mock(KMSClientProvider.class);
+    String keyName = "key1";
+    Mockito.doThrow(new IOException(new AuthorizationException("p1"))).when(p1)
+        .warmUpEncryptedKeys(Mockito.anyString());
+    KMSClientProvider p2 = mock(KMSClientProvider.class);
+    Mockito.doThrow(new IOException(new AuthorizationException("p2"))).when(p2)
+        .warmUpEncryptedKeys(Mockito.anyString());
+
+    when(p1.getKMSUrl()).thenReturn("p1");
+    when(p2.getKMSUrl()).thenReturn("p2");
+
+    LoadBalancingKMSClientProvider kp = new LoadBalancingKMSClientProvider(
+        new KMSClientProvider[] {p1, p2}, 0, conf);
+    try {
+      kp.warmUpEncryptedKeys(keyName);
+      fail("Should fail since both providers threw IOException");
+    } catch (Exception e) {
+      assertTrue(e.getCause() instanceof IOException);
+    }
+    Mockito.verify(p1, Mockito.times(1)).warmUpEncryptedKeys(keyName);
+    Mockito.verify(p2, Mockito.times(1)).warmUpEncryptedKeys(keyName);
+  }
+
+  /**
+   * tests {@link LoadBalancingKMSClientProvider#warmUpEncryptedKeys(String...)}
+   * error handling in case atleast one provider succeeds.
+   * @throws Exception
+   */
+  @Test
+  public void testWarmUpEncryptedKeysWhenOneProviderSucceeds()
+      throws Exception {
+    Configuration conf = new Configuration();
+    KMSClientProvider p1 = mock(KMSClientProvider.class);
+    String keyName = "key1";
+    Mockito.doThrow(new IOException(new AuthorizationException("p1"))).when(p1)
+        .warmUpEncryptedKeys(Mockito.anyString());
+    KMSClientProvider p2 = mock(KMSClientProvider.class);
+    Mockito.doNothing().when(p2)
+        .warmUpEncryptedKeys(Mockito.anyString());
+
+    when(p1.getKMSUrl()).thenReturn("p1");
+    when(p2.getKMSUrl()).thenReturn("p2");
+
+    LoadBalancingKMSClientProvider kp = new LoadBalancingKMSClientProvider(
+        new KMSClientProvider[] {p1, p2}, 0, conf);
+    try {
+      kp.warmUpEncryptedKeys(keyName);
+    } catch (Exception e) {
+      fail("Should not throw Exception since p2 doesn't throw Exception");
+    }
+    Mockito.verify(p1, Mockito.times(1)).warmUpEncryptedKeys(keyName);
+    Mockito.verify(p2, Mockito.times(1)).warmUpEncryptedKeys(keyName);
+  }
 }