LoadBalancingKMSClientProvider extends KeyProvider implements
    CryptoExtension,
    KeyProviderDelegationTokenExtension.DelegationTokenExtension { @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() + "]", ioe); } } if (!success && e != null) { throw e; } } LoadBalancingKMSClientProvider(URI providerUri,
      KMSClientProvider[] providers, Configuration conf); @VisibleForTesting  LoadBalancingKMSClientProvider(KMSClientProvider[] providers, long seed,
      Configuration conf); private  LoadBalancingKMSClientProvider(URI uri,
      KMSClientProvider[] providers, long seed, Configuration conf); @VisibleForTesting KMSClientProvider[] getProviders(); @Override org.apache.hadoop.security.token.Token<? extends TokenIdentifier> selectDelegationToken(Credentials creds); @Override String getCanonicalServiceName(); @Override Token<?> getDelegationToken(String renewer); @Override long renewDelegationToken(final Token<?> token); @Override Void cancelDelegationToken(final Token<?> token); @Override void warmUpEncryptedKeys(String... keyNames); @Override void drain(String keyName); @Override void invalidateCache(String keyName); @Override EncryptedKeyVersion generateEncryptedKey(final String encryptionKeyName); @Override KeyVersion decryptEncryptedKey(final EncryptedKeyVersion encryptedKeyVersion); @Override EncryptedKeyVersion reencryptEncryptedKey(
      final EncryptedKeyVersion ekv); @Override void reencryptEncryptedKeys(final List<EncryptedKeyVersion> ekvs); @Override KeyVersion getKeyVersion(final String versionName); @Override List<String> getKeys(); @Override Metadata[] getKeysMetadata(final String... names); @Override List<KeyVersion> getKeyVersions(final String name); @Override KeyVersion getCurrentKey(final String name); @Override Metadata getMetadata(final String name); @Override KeyVersion createKey(final String name, final byte[] material,
      final Options options); @Override KeyVersion createKey(final String name, final Options options); @Override void deleteKey(final String name); @Override KeyVersion rollNewVersion(final String name, final byte[] material); @Override KeyVersion rollNewVersion(final String name); @Override void close(); @Override void flush(); static Logger LOG; }
@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); }
@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); }