TestLoadBalancingKMSClientProvider.java
/**
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
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.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertNotEquals;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.junit.jupiter.api.Assertions.fail;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
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;
import java.security.GeneralSecurityException;
import java.security.NoSuchAlgorithmException;
import java.security.PrivilegedExceptionAction;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import javax.net.ssl.SSLException;
import javax.net.ssl.SSLHandshakeException;
import org.apache.hadoop.conf.Configuration;
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.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.Timeout;
import org.mockito.Mockito;
@Timeout(30)
public class TestLoadBalancingKMSClientProvider {
@BeforeAll
public static void setup() throws IOException {
SecurityUtil.setTokenServiceUseIp(false);
}
@Test
public void testCreation() throws Exception {
Configuration conf = new Configuration();
KeyProvider kp = new KMSClientProvider.Factory().createProvider(new URI(
"kms://http@host1:9600/kms/foo"), conf);
assertTrue(kp instanceof LoadBalancingKMSClientProvider);
KMSClientProvider[] providers =
((LoadBalancingKMSClientProvider) kp).getProviders();
assertEquals(1, providers.length);
assertEquals(new HashSet<>(Collections.singleton("http://host1:9600/kms/foo/v1/")),
new HashSet<>(Collections.singleton(providers[0].getKMSUrl())));
kp = new KMSClientProvider.Factory().createProvider(new URI(
"kms://http@host1;host2;host3:9600/kms/foo"), conf);
assertTrue(kp instanceof LoadBalancingKMSClientProvider);
providers =
((LoadBalancingKMSClientProvider) kp).getProviders();
assertEquals(3, providers.length);
assertEquals(new HashSet<>(Arrays.asList("http://host1:9600/kms/foo/v1/",
"http://host2:9600/kms/foo/v1/",
"http://host3:9600/kms/foo/v1/")),
new HashSet<>(Arrays.asList(providers[0].getKMSUrl(),
providers[1].getKMSUrl(),
providers[2].getKMSUrl())));
kp = new KMSClientProvider.Factory().createProvider(new URI(
"kms://http@host1;host2;host3:9600/kms/foo"), conf);
assertTrue(kp instanceof LoadBalancingKMSClientProvider);
providers =
((LoadBalancingKMSClientProvider) kp).getProviders();
assertEquals(3, providers.length);
assertEquals(new HashSet<>(Arrays.asList("http://host1:9600/kms/foo/v1/",
"http://host2:9600/kms/foo/v1/",
"http://host3:9600/kms/foo/v1/")),
new HashSet<>(Arrays.asList(providers[0].getKMSUrl(),
providers[1].getKMSUrl(),
providers[2].getKMSUrl())));
}
@Test
public void testLoadBalancing() throws Exception {
Configuration conf = new Configuration();
KMSClientProvider p1 = mock(KMSClientProvider.class);
when(p1.createKey(Mockito.anyString(), Mockito.any(Options.class)))
.thenReturn(
new KMSClientProvider.KMSKeyVersion("p1", "v1", new byte[0]));
KMSClientProvider p2 = mock(KMSClientProvider.class);
when(p2.createKey(Mockito.anyString(), Mockito.any(Options.class)))
.thenReturn(
new KMSClientProvider.KMSKeyVersion("p2", "v2", new byte[0]));
KMSClientProvider p3 = mock(KMSClientProvider.class);
when(p3.createKey(Mockito.anyString(), Mockito.any(Options.class)))
.thenReturn(
new KMSClientProvider.KMSKeyVersion("p3", "v3", new byte[0]));
KeyProvider kp = new LoadBalancingKMSClientProvider(
new KMSClientProvider[] { p1, p2, p3 }, 0, conf);
assertEquals("p1", kp.createKey("test1", new Options(conf)).getName());
assertEquals("p2", kp.createKey("test2", new Options(conf)).getName());
assertEquals("p3", kp.createKey("test3", new Options(conf)).getName());
assertEquals("p1", kp.createKey("test4", new Options(conf)).getName());
}
@Test
public void testLoadBalancingWithFailure() throws Exception {
Configuration conf = new Configuration();
KMSClientProvider p1 = mock(KMSClientProvider.class);
when(p1.createKey(Mockito.anyString(), Mockito.any(Options.class)))
.thenReturn(
new KMSClientProvider.KMSKeyVersion("p1", "v1", new byte[0]));
when(p1.getKMSUrl()).thenReturn("p1");
// This should not be retried
KMSClientProvider p2 = mock(KMSClientProvider.class);
when(p2.createKey(Mockito.anyString(), Mockito.any(Options.class)))
.thenThrow(new NoSuchAlgorithmException("p2"));
when(p2.getKMSUrl()).thenReturn("p2");
KMSClientProvider p3 = mock(KMSClientProvider.class);
when(p3.createKey(Mockito.anyString(), Mockito.any(Options.class)))
.thenReturn(
new KMSClientProvider.KMSKeyVersion("p3", "v3", new byte[0]));
when(p3.getKMSUrl()).thenReturn("p3");
// This should be retried
KMSClientProvider p4 = mock(KMSClientProvider.class);
when(p4.createKey(Mockito.anyString(), Mockito.any(Options.class)))
.thenThrow(new IOException("p4"));
when(p4.getKMSUrl()).thenReturn("p4");
KeyProvider kp = new LoadBalancingKMSClientProvider(
new KMSClientProvider[] { p1, p2, p3, p4 }, 0, conf);
assertEquals("p1", kp.createKey("test4", new Options(conf)).getName());
// Exceptions other than IOExceptions will not be retried
try {
kp.createKey("test1", new Options(conf)).getName();
fail("Should fail since its not an IOException");
} catch (Exception e) {
assertTrue(e instanceof NoSuchAlgorithmException);
}
assertEquals("p3", kp.createKey("test2", new Options(conf)).getName());
// IOException will trigger retry in next provider
assertEquals("p1", kp.createKey("test3", new Options(conf)).getName());
}
@Test
public void testLoadBalancingWithAllBadNodes() throws Exception {
Configuration conf = new Configuration();
KMSClientProvider p1 = mock(KMSClientProvider.class);
when(p1.createKey(Mockito.anyString(), Mockito.any(Options.class)))
.thenThrow(new IOException("p1"));
KMSClientProvider p2 = mock(KMSClientProvider.class);
when(p2.createKey(Mockito.anyString(), Mockito.any(Options.class)))
.thenThrow(new IOException("p2"));
KMSClientProvider p3 = mock(KMSClientProvider.class);
when(p3.createKey(Mockito.anyString(), Mockito.any(Options.class)))
.thenThrow(new IOException("p3"));
KMSClientProvider p4 = mock(KMSClientProvider.class);
when(p4.createKey(Mockito.anyString(), Mockito.any(Options.class)))
.thenThrow(new IOException("p4"));
when(p1.getKMSUrl()).thenReturn("p1");
when(p2.getKMSUrl()).thenReturn("p2");
when(p3.getKMSUrl()).thenReturn("p3");
when(p4.getKMSUrl()).thenReturn("p4");
KeyProvider kp = new LoadBalancingKMSClientProvider(
new KMSClientProvider[] { p1, p2, p3, p4 }, 0, conf);
try {
kp.createKey("test3", new Options(conf)).getName();
fail("Should fail since all providers threw an IOException");
} catch (Exception e) {
assertTrue(e instanceof IOException);
}
}
// copied from HttpExceptionUtils:
// trick, riding on generics to throw an undeclared exception
private static void throwEx(Throwable ex) {
TestLoadBalancingKMSClientProvider.<RuntimeException>throwException(ex);
}
@SuppressWarnings("unchecked")
private static <E extends Throwable> void throwException(Throwable ex)
throws E {
throw (E) ex;
}
private class MyKMSClientProvider extends KMSClientProvider {
public MyKMSClientProvider(URI uri, Configuration conf) throws IOException {
super(uri, conf);
}
@Override
public EncryptedKeyVersion generateEncryptedKey(
final String encryptionKeyName)
throws IOException, GeneralSecurityException {
throwEx(new AuthenticationException("bar"));
return null;
}
@Override
public KeyVersion decryptEncryptedKey(
final EncryptedKeyVersion encryptedKeyVersion) throws IOException,
GeneralSecurityException {
throwEx(new AuthenticationException("bar"));
return null;
}
@Override
public KeyVersion createKey(final String name, final Options options)
throws NoSuchAlgorithmException, IOException {
throwEx(new AuthenticationException("bar"));
return null;
}
@Override
public KeyVersion rollNewVersion(final String name)
throws NoSuchAlgorithmException, IOException {
throwEx(new AuthenticationException("bar"));
return null;
}
}
@Test
public void testClassCastException() throws Exception {
Configuration conf = new Configuration();
KMSClientProvider p1 = new MyKMSClientProvider(
new URI("kms://http@host1:9600/kms/foo"), conf);
LoadBalancingKMSClientProvider kp = new LoadBalancingKMSClientProvider(
new KMSClientProvider[] {p1}, 0, conf);
try {
kp.generateEncryptedKey("foo");
} catch (IOException ioe) {
assertTrue(ioe.getCause().getClass().getName().contains(
"AuthenticationException"));
}
try {
final KeyProviderCryptoExtension.EncryptedKeyVersion
encryptedKeyVersion =
mock(KeyProviderCryptoExtension.EncryptedKeyVersion.class);
kp.decryptEncryptedKey(encryptedKeyVersion);
} catch (IOException ioe) {
assertTrue(ioe.getCause().getClass().getName().contains(
"AuthenticationException"));
}
try {
final KeyProvider.Options options = KeyProvider.options(conf);
kp.createKey("foo", options);
} catch (IOException ioe) {
assertTrue(ioe.getCause().getClass().getName().contains(
"AuthenticationException"));
}
try {
kp.rollNewVersion("foo");
} catch (IOException ioe) {
assertTrue(ioe.getCause().getClass().getName().contains(
"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);
}
/**
* Tests whether retryPolicy fails immediately on non-idempotent operations,
* after trying each provider once,
* on encountering IOException which is not SocketException.
* @throws Exception
*/
@Test
public void testClientRetriesNonIdempotentOpWithIOExceptionFailsImmediately()
throws Exception {
Configuration conf = new Configuration();
final String keyName = "test";
// Setting total failover attempts to .
conf.setInt(
CommonConfigurationKeysPublic.KMS_CLIENT_FAILOVER_MAX_RETRIES_KEY, 10);
KMSClientProvider p1 = mock(KMSClientProvider.class);
when(p1.createKey(Mockito.anyString(), Mockito.any(Options.class)))
.thenThrow(new IOException("p1"));
KMSClientProvider p2 = mock(KMSClientProvider.class);
when(p2.createKey(Mockito.anyString(), Mockito.any(Options.class)))
.thenThrow(new IOException("p2"));
KMSClientProvider p3 = mock(KMSClientProvider.class);
when(p3.createKey(Mockito.anyString(), Mockito.any(Options.class)))
.thenThrow(new IOException("p3"));
when(p1.getKMSUrl()).thenReturn("p1");
when(p2.getKMSUrl()).thenReturn("p2");
when(p3.getKMSUrl()).thenReturn("p3");
LoadBalancingKMSClientProvider kp = new LoadBalancingKMSClientProvider(
new KMSClientProvider[] {p1, p2, p3}, 0, conf);
try {
kp.createKey(keyName, new Options(conf));
fail("Should fail since all providers threw an IOException");
} catch (Exception e) {
assertTrue(e instanceof IOException);
}
verify(kp.getProviders()[0], Mockito.times(1))
.createKey(Mockito.eq(keyName), Mockito.any(Options.class));
verify(kp.getProviders()[1], Mockito.times(1))
.createKey(Mockito.eq(keyName), Mockito.any(Options.class));
verify(kp.getProviders()[2], Mockito.times(1))
.createKey(Mockito.eq(keyName), Mockito.any(Options.class));
}
/**
* Tests whether retryPolicy retries on idempotent operations
* when encountering IOException.
* @throws Exception
*/
@Test
public void testClientRetriesIdempotentOpWithIOExceptionSucceedsSecondTime()
throws Exception {
Configuration conf = new Configuration();
final String keyName = "test";
final KeyProvider.KeyVersion keyVersion
= new KMSClientProvider.KMSKeyVersion(keyName, "v1",
new byte[0]);
// Setting total failover attempts to .
conf.setInt(
CommonConfigurationKeysPublic.KMS_CLIENT_FAILOVER_MAX_RETRIES_KEY, 10);
KMSClientProvider p1 = mock(KMSClientProvider.class);
when(p1.getCurrentKey(Mockito.anyString()))
.thenThrow(new IOException("p1"))
.thenReturn(keyVersion);
KMSClientProvider p2 = mock(KMSClientProvider.class);
when(p2.getCurrentKey(Mockito.anyString()))
.thenThrow(new IOException("p2"));
KMSClientProvider p3 = mock(KMSClientProvider.class);
when(p3.getCurrentKey(Mockito.anyString()))
.thenThrow(new IOException("p3"));
when(p1.getKMSUrl()).thenReturn("p1");
when(p2.getKMSUrl()).thenReturn("p2");
when(p3.getKMSUrl()).thenReturn("p3");
LoadBalancingKMSClientProvider kp = new LoadBalancingKMSClientProvider(
new KMSClientProvider[] {p1, p2, p3}, 0, conf);
KeyProvider.KeyVersion result = kp.getCurrentKey(keyName);
assertEquals(keyVersion, result);
verify(kp.getProviders()[0], Mockito.times(2))
.getCurrentKey(Mockito.eq(keyName));
verify(kp.getProviders()[1], Mockito.times(1))
.getCurrentKey(Mockito.eq(keyName));
verify(kp.getProviders()[2], Mockito.times(1))
.getCurrentKey(Mockito.eq(keyName));
}
/**
* Tests that client doesn't retry once it encounters AccessControlException
* from first provider.
* This assumes all the kms servers are configured with identical access to
* keys.
* @throws Exception
*/
@Test
public void testClientRetriesWithAccessControlException() throws Exception {
Configuration conf = new Configuration();
conf.setInt(
CommonConfigurationKeysPublic.KMS_CLIENT_FAILOVER_MAX_RETRIES_KEY, 3);
KMSClientProvider p1 = mock(KMSClientProvider.class);
when(p1.createKey(Mockito.anyString(), Mockito.any(Options.class)))
.thenThrow(new AccessControlException("p1"));
KMSClientProvider p2 = mock(KMSClientProvider.class);
when(p2.createKey(Mockito.anyString(), Mockito.any(Options.class)))
.thenThrow(new IOException("p2"));
KMSClientProvider p3 = mock(KMSClientProvider.class);
when(p3.createKey(Mockito.anyString(), Mockito.any(Options.class)))
.thenThrow(new IOException("p3"));
when(p1.getKMSUrl()).thenReturn("p1");
when(p2.getKMSUrl()).thenReturn("p2");
when(p3.getKMSUrl()).thenReturn("p3");
LoadBalancingKMSClientProvider kp = new LoadBalancingKMSClientProvider(
new KMSClientProvider[] {p1, p2, p3}, 0, conf);
try {
kp.createKey("test3", new Options(conf));
fail("Should fail because provider p1 threw an AccessControlException");
} catch (Exception e) {
assertTrue(e instanceof AccessControlException);
}
verify(p1, Mockito.times(1)).createKey(Mockito.eq("test3"),
Mockito.any(Options.class));
verify(p2, Mockito.never()).createKey(Mockito.eq("test3"),
Mockito.any(Options.class));
verify(p3, Mockito.never()).createKey(Mockito.eq("test3"),
Mockito.any(Options.class));
}
/**
* Tests that client doesn't retry once it encounters RunTimeException
* from first provider.
* This assumes all the kms servers are configured with identical access to
* keys.
* @throws Exception
*/
@Test
public void testClientRetriesWithRuntimeException() throws Exception {
Configuration conf = new Configuration();
conf.setInt(
CommonConfigurationKeysPublic.KMS_CLIENT_FAILOVER_MAX_RETRIES_KEY, 3);
KMSClientProvider p1 = mock(KMSClientProvider.class);
when(p1.createKey(Mockito.anyString(), Mockito.any(Options.class)))
.thenThrow(new RuntimeException("p1"));
KMSClientProvider p2 = mock(KMSClientProvider.class);
when(p2.createKey(Mockito.anyString(), Mockito.any(Options.class)))
.thenThrow(new IOException("p2"));
when(p1.getKMSUrl()).thenReturn("p1");
when(p2.getKMSUrl()).thenReturn("p2");
LoadBalancingKMSClientProvider kp = new LoadBalancingKMSClientProvider(
new KMSClientProvider[] {p1, p2}, 0, conf);
try {
kp.createKey("test3", new Options(conf));
fail("Should fail since provider p1 threw RuntimeException");
} catch (Exception e) {
assertTrue(e instanceof RuntimeException);
}
verify(p1, Mockito.times(1)).createKey(Mockito.eq("test3"),
Mockito.any(Options.class));
verify(p2, Mockito.never()).createKey(Mockito.eq("test3"),
Mockito.any(Options.class));
}
/**
* Tests the client retries until it finds a good provider.
* @throws Exception
*/
@Test
public void testClientRetriesWithTimeoutsException() throws Exception {
Configuration conf = new Configuration();
conf.setInt(
CommonConfigurationKeysPublic.KMS_CLIENT_FAILOVER_MAX_RETRIES_KEY, 4);
KMSClientProvider p1 = mock(KMSClientProvider.class);
when(p1.createKey(Mockito.anyString(), Mockito.any(Options.class)))
.thenThrow(new ConnectTimeoutException("p1"));
KMSClientProvider p2 = mock(KMSClientProvider.class);
when(p2.createKey(Mockito.anyString(), Mockito.any(Options.class)))
.thenThrow(new UnknownHostException("p2"));
KMSClientProvider p3 = mock(KMSClientProvider.class);
when(p3.createKey(Mockito.anyString(), Mockito.any(Options.class)))
.thenThrow(new NoRouteToHostException("p3"));
KMSClientProvider p4 = mock(KMSClientProvider.class);
when(p4.createKey(Mockito.anyString(), Mockito.any(Options.class)))
.thenReturn(
new KMSClientProvider.KMSKeyVersion("test3", "v1", new byte[0]));
when(p1.getKMSUrl()).thenReturn("p1");
when(p2.getKMSUrl()).thenReturn("p2");
when(p3.getKMSUrl()).thenReturn("p3");
when(p4.getKMSUrl()).thenReturn("p4");
LoadBalancingKMSClientProvider kp = new LoadBalancingKMSClientProvider(
new KMSClientProvider[] {p1, p2, p3, p4}, 0, conf);
try {
kp.createKey("test3", new Options(conf));
} catch (Exception e) {
fail("Provider p4 should have answered the request.");
}
verify(p1, Mockito.times(1)).createKey(Mockito.eq("test3"),
Mockito.any(Options.class));
verify(p2, Mockito.times(1)).createKey(Mockito.eq("test3"),
Mockito.any(Options.class));
verify(p3, Mockito.times(1)).createKey(Mockito.eq("test3"),
Mockito.any(Options.class));
verify(p4, Mockito.times(1)).createKey(Mockito.eq("test3"),
Mockito.any(Options.class));
}
/**
* Tests the operation succeeds second time after ConnectTimeoutException.
* @throws Exception
*/
@Test
public void testClientRetriesSucceedsSecondTime() throws Exception {
Configuration conf = new Configuration();
conf.setInt(
CommonConfigurationKeysPublic.KMS_CLIENT_FAILOVER_MAX_RETRIES_KEY, 3);
KMSClientProvider p1 = mock(KMSClientProvider.class);
when(p1.createKey(Mockito.anyString(), Mockito.any(Options.class)))
.thenThrow(new ConnectTimeoutException("p1"))
.thenReturn(new KMSClientProvider.KMSKeyVersion("test3", "v1",
new byte[0]));
KMSClientProvider p2 = mock(KMSClientProvider.class);
when(p2.createKey(Mockito.anyString(), Mockito.any(Options.class)))
.thenThrow(new ConnectTimeoutException("p2"));
when(p1.getKMSUrl()).thenReturn("p1");
when(p2.getKMSUrl()).thenReturn("p2");
LoadBalancingKMSClientProvider kp = new LoadBalancingKMSClientProvider(
new KMSClientProvider[] {p1, p2}, 0, conf);
try {
kp.createKey("test3", new Options(conf));
} catch (Exception e) {
fail("Provider p1 should have answered the request second time.");
}
verify(p1, Mockito.times(2)).createKey(Mockito.eq("test3"),
Mockito.any(Options.class));
verify(p2, Mockito.times(1)).createKey(Mockito.eq("test3"),
Mockito.any(Options.class));
}
/**
* Tests whether retryPolicy retries specified number of times.
* @throws Exception
*/
@Test
public void testClientRetriesSpecifiedNumberOfTimes() throws Exception {
Configuration conf = new Configuration();
conf.setInt(
CommonConfigurationKeysPublic.KMS_CLIENT_FAILOVER_MAX_RETRIES_KEY, 10);
KMSClientProvider p1 = mock(KMSClientProvider.class);
when(p1.createKey(Mockito.anyString(), Mockito.any(Options.class)))
.thenThrow(new ConnectTimeoutException("p1"));
KMSClientProvider p2 = mock(KMSClientProvider.class);
when(p2.createKey(Mockito.anyString(), Mockito.any(Options.class)))
.thenThrow(new ConnectTimeoutException("p2"));
when(p1.getKMSUrl()).thenReturn("p1");
when(p2.getKMSUrl()).thenReturn("p2");
LoadBalancingKMSClientProvider kp = new LoadBalancingKMSClientProvider(
new KMSClientProvider[] {p1, p2}, 0, conf);
try {
kp.createKey("test3", new Options(conf));
fail("Should fail");
} catch (Exception e) {
assert (e instanceof ConnectTimeoutException);
}
verify(p1, Mockito.times(6)).createKey(Mockito.eq("test3"),
Mockito.any(Options.class));
verify(p2, Mockito.times(5)).createKey(Mockito.eq("test3"),
Mockito.any(Options.class));
}
/**
* Tests whether retryPolicy retries number of times equals to number of
* providers if conf kms.client.failover.max.attempts is not set.
* @throws Exception
*/
@Test
public void testClientRetriesIfMaxAttemptsNotSet() throws Exception {
Configuration conf = new Configuration();
KMSClientProvider p1 = mock(KMSClientProvider.class);
when(p1.createKey(Mockito.anyString(), Mockito.any(Options.class)))
.thenThrow(new ConnectTimeoutException("p1"));
KMSClientProvider p2 = mock(KMSClientProvider.class);
when(p2.createKey(Mockito.anyString(), Mockito.any(Options.class)))
.thenThrow(new ConnectTimeoutException("p2"));
when(p1.getKMSUrl()).thenReturn("p1");
when(p2.getKMSUrl()).thenReturn("p2");
LoadBalancingKMSClientProvider kp = new LoadBalancingKMSClientProvider(
new KMSClientProvider[] {p1, p2}, 0, conf);
try {
kp.createKey("test3", new Options(conf));
fail("Should fail");
} catch (Exception e) {
assert (e instanceof ConnectTimeoutException);
}
verify(p1, Mockito.times(2)).createKey(Mockito.eq("test3"),
Mockito.any(Options.class));
verify(p2, Mockito.times(1)).createKey(Mockito.eq("test3"),
Mockito.any(Options.class));
}
/**
* Tests that client reties each provider once, when it encounters
* AuthenticationException wrapped in an IOException from first provider.
* @throws Exception
*/
@Test
public void testClientRetriesWithAuthenticationExceptionWrappedinIOException()
throws Exception {
Configuration conf = new Configuration();
conf.setInt(
CommonConfigurationKeysPublic.KMS_CLIENT_FAILOVER_MAX_RETRIES_KEY, 3);
KMSClientProvider p1 = mock(KMSClientProvider.class);
when(p1.createKey(Mockito.anyString(), Mockito.any(Options.class)))
.thenThrow(new IOException(new AuthenticationException("p1")));
KMSClientProvider p2 = mock(KMSClientProvider.class);
when(p2.createKey(Mockito.anyString(), Mockito.any(Options.class)))
.thenThrow(new IOException(new AuthenticationException("p1")));
when(p1.getKMSUrl()).thenReturn("p1");
when(p2.getKMSUrl()).thenReturn("p2");
LoadBalancingKMSClientProvider kp = new LoadBalancingKMSClientProvider(
new KMSClientProvider[] {p1, p2}, 0, conf);
try {
kp.createKey("test3", new Options(conf));
fail("Should fail since provider p1 threw AuthenticationException");
} catch (Exception e) {
assertTrue(e.getCause() instanceof AuthenticationException);
}
verify(p1, Mockito.times(1)).createKey(Mockito.eq("test3"),
Mockito.any(Options.class));
verify(p2, Mockito.times(1)).createKey(Mockito.eq("test3"),
Mockito.any(Options.class));
}
/**
* Tests the operation succeeds second time after SSLHandshakeException.
* @throws Exception
*/
@Test
public void testClientRetriesWithSSLHandshakeExceptionSucceedsSecondTime()
throws Exception {
Configuration conf = new Configuration();
conf.setInt(
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 SocketException("p1"));
when(p1.getKMSUrl()).thenReturn("p1");
when(p2.getKMSUrl()).thenReturn("p2");
LoadBalancingKMSClientProvider kp = new LoadBalancingKMSClientProvider(
new KMSClientProvider[] {p1, p2}, 0, conf);
kp.createKey(keyName, new Options(conf));
verify(p1, Mockito.times(3)).createKey(Mockito.eq(keyName),
Mockito.any(Options.class));
verify(p2, Mockito.times(2)).createKey(Mockito.eq(keyName),
Mockito.any(Options.class));
}
/**
* Tests the operation fails at every attempt after SSLHandshakeException.
* @throws Exception
*/
@Test
public void testClientRetriesWithSSLHandshakeExceptionFailsAtEveryAttempt()
throws Exception {
Configuration conf = new Configuration();
conf.setInt(
CommonConfigurationKeysPublic.KMS_CLIENT_FAILOVER_MAX_RETRIES_KEY, 2);
final String keyName = "test";
final String exceptionMessage = "p1 exception message";
KMSClientProvider p1 = mock(KMSClientProvider.class);
Exception originalSslEx = new SSLHandshakeException(exceptionMessage);
when(p1.createKey(Mockito.anyString(), Mockito.any(Options.class)))
.thenThrow(originalSslEx);
KMSClientProvider p2 = mock(KMSClientProvider.class);
when(p2.createKey(Mockito.anyString(), Mockito.any(Options.class)))
.thenThrow(new ConnectException("p2 exception message"));
when(p1.getKMSUrl()).thenReturn("p1");
when(p2.getKMSUrl()).thenReturn("p2");
LoadBalancingKMSClientProvider kp = new LoadBalancingKMSClientProvider(
new KMSClientProvider[] {p1, p2}, 0, conf);
Exception interceptedEx = intercept(ConnectException.class,
"SSLHandshakeException: " + exceptionMessage,
()-> kp.createKey(keyName, new Options(conf)));
assertEquals(originalSslEx, interceptedEx.getCause());
verify(p1, Mockito.times(2)).createKey(Mockito.eq(keyName),
Mockito.any(Options.class));
verify(p2, Mockito.times(1)).createKey(Mockito.eq(keyName),
Mockito.any(Options.class));
}
/**
* Tests that if an idempotent operation succeeds second time after
* SocketTimeoutException, then the operation is successful.
* @throws Exception
*/
@Test
public void testClientRetriesIdempotentOpWithSocketTimeoutExceptionSucceeds()
throws Exception {
Configuration conf = new Configuration();
conf.setInt(
CommonConfigurationKeysPublic.KMS_CLIENT_FAILOVER_MAX_RETRIES_KEY, 3);
final List<String> keys = Arrays.asList("testKey");
KMSClientProvider p1 = mock(KMSClientProvider.class);
when(p1.getKeys())
.thenThrow(new SocketTimeoutException("p1"))
.thenReturn(keys);
KMSClientProvider p2 = mock(KMSClientProvider.class);
when(p2.getKeys()).thenThrow(new SocketTimeoutException("p2"));
when(p1.getKMSUrl()).thenReturn("p1");
when(p2.getKMSUrl()).thenReturn("p2");
LoadBalancingKMSClientProvider kp = new LoadBalancingKMSClientProvider(
new KMSClientProvider[] {p1, p2}, 0, conf);
List<String> result = kp.getKeys();
assertEquals(keys, result);
verify(p1, Mockito.times(2)).getKeys();
verify(p2, Mockito.times(1)).getKeys();
}
/**
* Tests that if a non idempotent operation fails at every attempt
* after SocketTimeoutException, then SocketTimeoutException is thrown.
* @throws Exception
*/
@Test
public void testClientRetriesIdempotentOpWithSocketTimeoutExceptionFails()
throws Exception {
Configuration conf = new Configuration();
conf.setInt(
CommonConfigurationKeysPublic.KMS_CLIENT_FAILOVER_MAX_RETRIES_KEY, 2);
final String keyName = "test";
final String exceptionMessage = "p1 exception message";
KMSClientProvider p1 = mock(KMSClientProvider.class);
Exception originalEx = new SocketTimeoutException(exceptionMessage);
when(p1.getKeyVersions(Mockito.anyString()))
.thenThrow(originalEx);
KMSClientProvider p2 = mock(KMSClientProvider.class);
when(p2.getKeyVersions(Mockito.anyString()))
.thenThrow(new SocketTimeoutException("p2 exception message"));
when(p1.getKMSUrl()).thenReturn("p1");
when(p2.getKMSUrl()).thenReturn("p2");
LoadBalancingKMSClientProvider kp = new LoadBalancingKMSClientProvider(
new KMSClientProvider[] {p1, p2}, 0, conf);
Exception interceptedEx = intercept(SocketTimeoutException.class,
"SocketTimeoutException: " + exceptionMessage,
()-> kp.getKeyVersions(keyName));
assertEquals(originalEx, interceptedEx);
verify(p1, Mockito.times(2))
.getKeyVersions(Mockito.eq(keyName));
verify(p2, Mockito.times(1))
.getKeyVersions(Mockito.eq(keyName));
}
/**
* Tests whether retryPolicy fails immediately on non-idempotent operations,
* after trying each provider once, on encountering SocketTimeoutException.
* @throws Exception
*/
@Test
public void testClientRetriesNonIdempotentOpWithSocketTimeoutExceptionFails()
throws Exception {
Configuration conf = new Configuration();
final String keyName = "test";
// Setting total failover attempts to .
conf.setInt(
CommonConfigurationKeysPublic.KMS_CLIENT_FAILOVER_MAX_RETRIES_KEY, 10);
KMSClientProvider p1 = mock(KMSClientProvider.class);
when(p1.createKey(Mockito.anyString(), Mockito.any(Options.class)))
.thenThrow(new SocketTimeoutException("p1"));
KMSClientProvider p2 = mock(KMSClientProvider.class);
when(p2.createKey(Mockito.anyString(), Mockito.any(Options.class)))
.thenThrow(new SocketTimeoutException("p2"));
KMSClientProvider p3 = mock(KMSClientProvider.class);
when(p3.createKey(Mockito.anyString(), Mockito.any(Options.class)))
.thenThrow(new SocketTimeoutException("p3"));
when(p1.getKMSUrl()).thenReturn("p1");
when(p2.getKMSUrl()).thenReturn("p2");
when(p3.getKMSUrl()).thenReturn("p3");
LoadBalancingKMSClientProvider kp = new LoadBalancingKMSClientProvider(
new KMSClientProvider[] {p1, p2, p3}, 0, conf);
try {
kp.createKey(keyName, new Options(conf));
fail("Should fail since all providers threw a SocketTimeoutException");
} catch (Exception e) {
assertTrue(e instanceof SocketTimeoutException);
}
verify(kp.getProviders()[0], Mockito.times(1))
.createKey(Mockito.eq(keyName), Mockito.any(Options.class));
verify(kp.getProviders()[1], Mockito.times(1))
.createKey(Mockito.eq(keyName), Mockito.any(Options.class));
verify(kp.getProviders()[2], Mockito.times(1))
.createKey(Mockito.eq(keyName), Mockito.any(Options.class));
}
@Test
public void testTokenServiceCreationWithLegacyFormat() throws Exception {
Configuration conf = new Configuration();
// Create keyprovider with old token format (ip:port)
conf.set(CommonConfigurationKeysPublic.HADOOP_SECURITY_KEY_PROVIDER_PATH,
"kms:/something");
String authority = "host1:9600";
URI kmsUri = URI.create("kms://http@" + authority + "/kms/foo");
KeyProvider kp =
new KMSClientProvider.Factory().createProvider(kmsUri, conf);
assertTrue(kp instanceof LoadBalancingKMSClientProvider);
LoadBalancingKMSClientProvider lbkp = (LoadBalancingKMSClientProvider) kp;
assertEquals(1, lbkp.getProviders().length);
assertEquals(authority, lbkp.getCanonicalServiceName());
for (KMSClientProvider provider : lbkp.getProviders()) {
assertEquals(authority, provider.getCanonicalServiceName());
}
}
@Test
public void testTokenServiceCreationWithUriFormat() throws Exception {
final Configuration conf = new Configuration();
final URI kmsUri = URI.create("kms://http@host1;host2;host3:9600/kms/foo");
final KeyProvider kp =
new KMSClientProvider.Factory().createProvider(kmsUri, conf);
assertTrue(kp instanceof LoadBalancingKMSClientProvider);
final LoadBalancingKMSClientProvider lbkp =
(LoadBalancingKMSClientProvider) kp;
assertEquals(kmsUri.toString(), lbkp.getCanonicalServiceName());
KMSClientProvider[] providers = lbkp.getProviders();
assertEquals(3, providers.length);
for (int i = 0; i < providers.length; i++) {
assertEquals(URI.create(providers[i].getKMSUrl()).getAuthority(),
providers[i].getCanonicalServiceName());
assertNotEquals(kmsUri, providers[i].getCanonicalServiceName());
}
}
private void testTokenSelectionWithConf(Configuration conf) throws Exception {
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";
conf.set(CommonConfigurationKeysPublic.HADOOP_SECURITY_KEY_PROVIDER_PATH,
providerUriString);
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(ugi, actualUgi,
"testTokenSelectionWithConf() should return the" +
" current user, not login user");
}
@Test
public void testTokenSelectionWithKMSUriInConf() throws Exception {
final Configuration conf = new Configuration();
conf.set("hadoop.security.authentication", "kerberos");
// test client with hadoop.security.key.provider.path configured.
String providerUriString = "kms://http@host1;host2;host3:9600/kms/foo";
conf.set(CommonConfigurationKeysPublic.HADOOP_SECURITY_KEY_PROVIDER_PATH,
providerUriString);
testTokenSelectionWithConf(conf);
}
@Test
public void testGetActualUGI() throws Exception {
final Configuration conf = new Configuration();
conf.set("hadoop.security.authentication", "kerberos");
UserGroupInformation.setConfiguration(conf);
// test client without hadoop.security.key.provider.path configured.
testTokenSelectionWithConf(conf);
}
}