From 78008bc0eeaebb52c84e1654914d58e5eab8382c Mon Sep 17 00:00:00 2001
From: Ritesh H Shukla <kerneltime@gmail.com>
Date: Thu, 19 May 2022 13:20:24 -0700
Subject: [PATCH] HADOOP-18245 Extend KMS related exceptions that get mapped to
 ConnectException (#4329)

---
 .../key/kms/LoadBalancingKMSClientProvider.java      |  6 ++++--
 .../key/kms/TestLoadBalancingKMSClientProvider.java  | 12 ++++++++----
 2 files changed, 12 insertions(+), 6 deletions(-)

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 6f8f4585ee7..f46da1f3186 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
@@ -21,6 +21,7 @@ package org.apache.hadoop.crypto.key.kms;
 import java.io.IOException;
 import java.io.InterruptedIOException;
 import java.net.ConnectException;
+import java.net.SocketException;
 import java.net.URI;
 import java.security.GeneralSecurityException;
 import java.security.NoSuchAlgorithmException;
@@ -29,6 +30,7 @@ import java.util.Collections;
 import java.util.List;
 import java.util.concurrent.atomic.AtomicInteger;
 
+import javax.net.ssl.SSLException;
 import javax.net.ssl.SSLHandshakeException;
 
 import org.apache.hadoop.conf.Configuration;
@@ -182,10 +184,10 @@ public class LoadBalancingKMSClientProvider extends KeyProvider implements
       } catch (IOException ioe) {
         LOG.warn("KMS provider at [{}] threw an IOException: ",
             provider.getKMSUrl(), ioe);
-        // SSLHandshakeException can occur here because of lost connection
+        // SSLException can occur here because of lost connection
         // with the KMS server, creating a ConnectException from it,
         // so that the FailoverOnNetworkExceptionRetry policy will retry
-        if (ioe instanceof SSLHandshakeException) {
+        if (ioe instanceof SSLException || ioe instanceof SocketException) {
           Exception cause = ioe;
           ioe = new ConnectException("SSLHandshakeException: "
               + cause.getMessage());
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 9122389bd6e..886297b745c 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
@@ -31,6 +31,7 @@ import static org.mockito.Mockito.verify;
 import java.io.IOException;
 import java.net.ConnectException;
 import java.net.NoRouteToHostException;
+import java.net.SocketException;
 import java.net.SocketTimeoutException;
 import java.net.URI;
 import java.net.UnknownHostException;
@@ -41,6 +42,7 @@ import java.util.Arrays;
 import java.util.List;
 import java.util.concurrent.TimeUnit;
 
+import javax.net.ssl.SSLException;
 import javax.net.ssl.SSLHandshakeException;
 
 import org.apache.hadoop.conf.Configuration;
@@ -707,16 +709,18 @@ public class TestLoadBalancingKMSClientProvider {
       throws Exception {
     Configuration conf = new Configuration();
     conf.setInt(
-        CommonConfigurationKeysPublic.KMS_CLIENT_FAILOVER_MAX_RETRIES_KEY, 3);
+        CommonConfigurationKeysPublic.KMS_CLIENT_FAILOVER_MAX_RETRIES_KEY, 5);
     final String keyName = "test";
     KMSClientProvider p1 = mock(KMSClientProvider.class);
     when(p1.createKey(Mockito.anyString(), Mockito.any(Options.class)))
         .thenThrow(new SSLHandshakeException("p1"))
+        .thenThrow(new SSLException("p1"))
         .thenReturn(new KMSClientProvider.KMSKeyVersion(keyName, "v1",
             new byte[0]));
     KMSClientProvider p2 = mock(KMSClientProvider.class);
     when(p2.createKey(Mockito.anyString(), Mockito.any(Options.class)))
-        .thenThrow(new ConnectException("p2"));
+        .thenThrow(new ConnectException("p2"))
+        .thenThrow(new SocketException("p1"));
 
     when(p1.getKMSUrl()).thenReturn("p1");
     when(p2.getKMSUrl()).thenReturn("p2");
@@ -725,9 +729,9 @@ public class TestLoadBalancingKMSClientProvider {
         new KMSClientProvider[] {p1, p2}, 0, conf);
 
     kp.createKey(keyName, new Options(conf));
-    verify(p1, Mockito.times(2)).createKey(Mockito.eq(keyName),
+    verify(p1, Mockito.times(3)).createKey(Mockito.eq(keyName),
         Mockito.any(Options.class));
-    verify(p2, Mockito.times(1)).createKey(Mockito.eq(keyName),
+    verify(p2, Mockito.times(2)).createKey(Mockito.eq(keyName),
         Mockito.any(Options.class));
   }