org.apache.hadoop.crypto.key.KeyProvider Java Examples

@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());
}
 

@GET
@Path(KMSRESTConstants.KEY_RESOURCE + "/{name:.*}/" +
    KMSRESTConstants.METADATA_SUB_RESOURCE)
@Produces(MediaType.APPLICATION_JSON)
public Response getMetadata(@PathParam("name") final String name)
    throws Exception {
  UserGroupInformation user = HttpUserGroupInformation.get();
  KMSClientProvider.checkNotEmpty(name, "name");
  KMSWebApp.getAdminCallsMeter().mark();
  assertAccess(KMSACLs.Type.GET_METADATA, user, KMSOp.GET_METADATA, name);

  KeyProvider.Metadata metadata = user.doAs(
      new PrivilegedExceptionAction<KeyProvider.Metadata>() {
        @Override
        public KeyProvider.Metadata run() throws Exception {
          return provider.getMetadata(name);
        }
      }
  );

  Object json = KMSServerJSONUtils.toJSON(name, metadata);
  kmsAudit.ok(user, KMSOp.GET_METADATA, name, "");
  return Response.ok().type(MediaType.APPLICATION_JSON).entity(json).build();
}
 

@GET
@Path(KMSRESTConstants.KEYS_METADATA_RESOURCE)
@Produces(MediaType.APPLICATION_JSON)
public Response getKeysMetadata(@QueryParam(KMSRESTConstants.KEY)
    List<String> keyNamesList) throws Exception {
  KMSWebApp.getAdminCallsMeter().mark();
  UserGroupInformation user = HttpUserGroupInformation.get();
  final String[] keyNames = keyNamesList.toArray(
      new String[keyNamesList.size()]);
  assertAccess(KMSACLs.Type.GET_METADATA, user, KMSOp.GET_KEYS_METADATA);

  KeyProvider.Metadata[] keysMeta = user.doAs(
      new PrivilegedExceptionAction<KeyProvider.Metadata[]>() {
        @Override
        public KeyProvider.Metadata[] run() throws Exception {
          return provider.getKeysMetadata(keyNames);
        }
      }
  );

  Object json = KMSServerJSONUtils.toJSON(keyNames, keysMeta);
  kmsAudit.ok(user, KMSOp.GET_KEYS_METADATA, "");
  return Response.ok().type(MediaType.APPLICATION_JSON).entity(json).build();
}
 

public KeyProviderCache(long expiryMs) {
  cache = CacheBuilder.newBuilder()
      .expireAfterAccess(expiryMs, TimeUnit.MILLISECONDS)
      .removalListener(new RemovalListener<URI, KeyProvider>() {
        @Override
        public void onRemoval(
            RemovalNotification<URI, KeyProvider> notification) {
          try {
            notification.getValue().close();
          } catch (Throwable e) {
            LOG.error(
                "Error closing KeyProvider with uri ["
                    + notification.getKey() + "]", e);
            ;
          }
        }
      })
      .build();
}
 

private static KeyProvider.KeyVersion decryptEncryptedDataEncryptionKey(DistributedFileSystem dfs, FileEncryptionInfo feInfo) throws IOException {
    KeyProvider provider = dfs.dfs.getKeyProvider();
    if (provider == null) {
        throw new IOException("No KeyProvider is configured, cannot access" +
                " an encrypted file");
    }
    KeyProviderCryptoExtension.EncryptedKeyVersion ekv = KeyProviderCryptoExtension.EncryptedKeyVersion.createForDecryption(
            feInfo.getKeyName(), feInfo.getEzKeyVersionName(), feInfo.getIV(),
            feInfo.getEncryptedDataEncryptionKey());
    try {
        KeyProviderCryptoExtension cryptoProvider = KeyProviderCryptoExtension
                .createKeyProviderCryptoExtension(provider);
        return cryptoProvider.decryptEncryptedKey(ekv);
    } catch (GeneralSecurityException e) {
        throw new IOException(e);
    }
}
 

@SuppressWarnings("unchecked")
public static Map toJSON(String keyName, KeyProvider.Metadata meta) {
  Map json = new LinkedHashMap();
  if (meta != null) {
    json.put(KMSRESTConstants.NAME_FIELD, keyName);
    json.put(KMSRESTConstants.CIPHER_FIELD, meta.getCipher());
    json.put(KMSRESTConstants.LENGTH_FIELD, meta.getBitLength());
    json.put(KMSRESTConstants.DESCRIPTION_FIELD, meta.getDescription());
    json.put(KMSRESTConstants.ATTRIBUTES_FIELD, meta.getAttributes());
    json.put(KMSRESTConstants.CREATED_FIELD,
        meta.getCreated().getTime());
    json.put(KMSRESTConstants.VERSIONS_FIELD,
        (long) meta.getVersions());
  }
  return json;
}
 

public KeyProvider get(final Configuration conf) {
  URI kpURI = createKeyProviderURI(conf);
  if (kpURI == null) {
    return null;
  }
  try {
    return cache.get(kpURI, new Callable<KeyProvider>() {
      @Override
      public KeyProvider call() throws Exception {
        return DFSUtil.createKeyProvider(conf);
      }
    });
  } catch (Exception e) {
    LOG.error("Could not create KeyProvider for DFSClient !!", e.getCause());
    return null;
  }
}
 

public HdfsDataInputStream createWrappedInputStream(DFSInputStream dfsis)
        throws IOException {
    final FileEncryptionInfo feInfo = dfsis.getFileEncryptionInfo();
    if (feInfo != null) {
        // File is encrypted, wrap the stream in a crypto stream.
        // Currently only one version, so no special logic based on the version #
        getCryptoProtocolVersion(feInfo);
        final CryptoCodec codec = getCryptoCodec(getConfiguration(), feInfo);
        final KeyProvider.KeyVersion decrypted = decryptEncryptedDataEncryptionKey(dfsis, feInfo);
        final CryptoInputStream cryptoIn =
                new CryptoInputStream(dfsis, codec, decrypted.getMaterial(),
                        feInfo.getIV());
        return new HdfsDataInputStream(cryptoIn);
    } else {
        // No FileEncryptionInfo so no encryption.
        return new HdfsDataInputStream(dfsis);
    }
}
 

public static KeyProvider.KeyVersion decryptEncryptedDataEncryptionKey(
    FileEncryptionInfo feInfo, KeyProvider keyProvider) throws IOException {
  if (keyProvider == null) {
    throw new IOException("No KeyProvider is configured, " +
        "cannot access an encrypted file");
  } else {
    EncryptedKeyVersion ekv = EncryptedKeyVersion.createForDecryption(
        feInfo.getKeyName(), feInfo.getEzKeyVersionName(), feInfo.getIV(),
        feInfo.getEncryptedDataEncryptionKey());

    try {
      KeyProviderCryptoExtension cryptoProvider = KeyProviderCryptoExtension
          .createKeyProviderCryptoExtension(keyProvider);
      return cryptoProvider.decryptEncryptedKey(ekv);
    } catch (GeneralSecurityException gse) {
      throw new IOException(gse);
    }
  }
}
 

@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());
}
 

/**
 * Decrypts a EDEK by consulting the KeyProvider.
 */
private KeyVersion decryptEncryptedDataEncryptionKey(FileEncryptionInfo
    feInfo) throws IOException {
  TraceScope scope = Trace.startSpan("decryptEDEK", traceSampler);
  try {
    KeyProvider provider = getKeyProvider();
    if (provider == null) {
      throw new IOException("No KeyProvider is configured, cannot access" +
          " an encrypted file");
    }
    EncryptedKeyVersion ekv = EncryptedKeyVersion.createForDecryption(
        feInfo.getKeyName(), feInfo.getEzKeyVersionName(), feInfo.getIV(),
        feInfo.getEncryptedDataEncryptionKey());
    try {
      KeyProviderCryptoExtension cryptoProvider = KeyProviderCryptoExtension
          .createKeyProviderCryptoExtension(provider);
      return cryptoProvider.decryptEncryptedKey(ekv);
    } catch (GeneralSecurityException e) {
      throw new IOException(e);
    }
  } finally {
    scope.close();
  }
}
 

private KeyProvider createProvider(URI providerUri, Configuration conf,
    URL origUrl, int port, String hostsPart) throws IOException {
  String[] hosts = hostsPart.split(";");
  if (hosts.length == 1) {
    return new KMSClientProvider(providerUri, conf);
  } else {
    KMSClientProvider[] providers = new KMSClientProvider[hosts.length];
    for (int i = 0; i < hosts.length; i++) {
      try {
        providers[i] =
            new KMSClientProvider(
                new URI("kms", origUrl.getProtocol(), hosts[i], port,
                    origUrl.getPath(), null, null), conf);
      } catch (URISyntaxException e) {
        throw new IOException("Could not instantiate KMSProvider..", e);
      }
    }
    return new LoadBalancingKMSClientProvider(providers, conf);
  }
}
 

@GET
@Path(KMSRESTConstants.KEYS_METADATA_RESOURCE)
@Produces(MediaType.APPLICATION_JSON)
public Response getKeysMetadata(@QueryParam(KMSRESTConstants.KEY)
    List<String> keyNamesList) throws Exception {
  KMSWebApp.getAdminCallsMeter().mark();
  UserGroupInformation user = HttpUserGroupInformation.get();
  final String[] keyNames = keyNamesList.toArray(
      new String[keyNamesList.size()]);
  assertAccess(KMSACLs.Type.GET_METADATA, user, KMSOp.GET_KEYS_METADATA);

  KeyProvider.Metadata[] keysMeta = user.doAs(
      new PrivilegedExceptionAction<KeyProvider.Metadata[]>() {
        @Override
        public KeyProvider.Metadata[] run() throws Exception {
          return provider.getKeysMetadata(keyNames);
        }
      }
  );

  Object json = KMSServerJSONUtils.toJSON(keyNames, keysMeta);
  kmsAudit.ok(user, KMSOp.GET_KEYS_METADATA, "");
  return Response.ok().type(MediaType.APPLICATION_JSON).entity(json).build();
}
 

@GET
@Path(KMSRESTConstants.KEY_RESOURCE + "/{name:.*}/" +
    KMSRESTConstants.METADATA_SUB_RESOURCE)
@Produces(MediaType.APPLICATION_JSON)
public Response getMetadata(@PathParam("name") final String name)
    throws Exception {
  UserGroupInformation user = HttpUserGroupInformation.get();
  KMSClientProvider.checkNotEmpty(name, "name");
  KMSWebApp.getAdminCallsMeter().mark();
  assertAccess(KMSACLs.Type.GET_METADATA, user, KMSOp.GET_METADATA, name);

  KeyProvider.Metadata metadata = user.doAs(
      new PrivilegedExceptionAction<KeyProvider.Metadata>() {
        @Override
        public KeyProvider.Metadata run() throws Exception {
          return provider.getMetadata(name);
        }
      }
  );

  Object json = KMSServerJSONUtils.toJSON(name, metadata);
  kmsAudit.ok(user, KMSOp.GET_METADATA, name, "");
  return Response.ok().type(MediaType.APPLICATION_JSON).entity(json).build();
}
 

@SuppressWarnings("unchecked")
public static Map toJSON(String keyName, KeyProvider.Metadata meta) {
  Map json = new LinkedHashMap();
  if (meta != null) {
    json.put(KMSRESTConstants.NAME_FIELD, keyName);
    json.put(KMSRESTConstants.CIPHER_FIELD, meta.getCipher());
    json.put(KMSRESTConstants.LENGTH_FIELD, meta.getBitLength());
    json.put(KMSRESTConstants.DESCRIPTION_FIELD, meta.getDescription());
    json.put(KMSRESTConstants.ATTRIBUTES_FIELD, meta.getAttributes());
    json.put(KMSRESTConstants.CREATED_FIELD,
        meta.getCreated().getTime());
    json.put(KMSRESTConstants.VERSIONS_FIELD,
        (long) meta.getVersions());
  }
  return json;
}
 

@SuppressWarnings("unchecked")
public static List toJSON(List<KeyProvider.KeyVersion> keyVersions) {
  List json = new ArrayList();
  if (keyVersions != null) {
    for (KeyProvider.KeyVersion version : keyVersions) {
      json.add(toJSON(version));
    }
  }
  return json;
}
 

@SuppressWarnings("unchecked")
public static Map toJSON(KeyProvider.KeyVersion keyVersion) {
  Map json = new LinkedHashMap();
  if (keyVersion != null) {
    json.put(KMSRESTConstants.NAME_FIELD,
        keyVersion.getName());
    json.put(KMSRESTConstants.VERSION_NAME_FIELD,
        keyVersion.getVersionName());
    json.put(KMSRESTConstants.MATERIAL_FIELD,
        Base64.encodeBase64URLSafeString(
            keyVersion.getMaterial()));
  }
  return json;
}
 

@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);
  }
}
 

@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);
  }
}
 

@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 testCreation() throws Exception {
  Configuration conf = new Configuration();
  KeyProvider kp = new KMSClientProvider.Factory().createProvider(new URI(
      "kms://http@host1/kms/foo"), conf);
  assertTrue(kp instanceof KMSClientProvider);
  assertEquals("http://host1/kms/foo/v1/",
      ((KMSClientProvider) kp).getKMSUrl());

  kp = new KMSClientProvider.Factory().createProvider(new URI(
      "kms://http@host1;host2;host3/kms/foo"), conf);
  assertTrue(kp instanceof LoadBalancingKMSClientProvider);
  KMSClientProvider[] providers =
      ((LoadBalancingKMSClientProvider) kp).getProviders();
  assertEquals(3, providers.length);
  assertEquals(Sets.newHashSet("http://host1/kms/foo/v1/",
      "http://host2/kms/foo/v1/",
      "http://host3/kms/foo/v1/"),
      Sets.newHashSet(providers[0].getKMSUrl(),
          providers[1].getKMSUrl(),
          providers[2].getKMSUrl()));

  kp = new KMSClientProvider.Factory().createProvider(new URI(
      "kms://http@host1;host2;host3:16000/kms/foo"), conf);
  assertTrue(kp instanceof LoadBalancingKMSClientProvider);
  providers =
      ((LoadBalancingKMSClientProvider) kp).getProviders();
  assertEquals(3, providers.length);
  assertEquals(Sets.newHashSet("http://host1:16000/kms/foo/v1/",
      "http://host2:16000/kms/foo/v1/",
      "http://host3:16000/kms/foo/v1/"),
      Sets.newHashSet(providers[0].getKMSUrl(),
          providers[1].getKMSUrl(),
          providers[2].getKMSUrl()));
}
 

/**
 * This provider expects URIs in the following form :
 * kms://<PROTO>@<AUTHORITY>/<PATH>
 *
 * where :
 * - PROTO = http or https
 * - AUTHORITY = <HOSTS>[:<PORT>]
 * - HOSTS = <HOSTNAME>[;<HOSTS>]
 * - HOSTNAME = string
 * - PORT = integer
 *
 * If multiple hosts are provider, the Factory will create a
 * {@link LoadBalancingKMSClientProvider} that round-robins requests
 * across the provided list of hosts.
 */
@Override
public KeyProvider createProvider(URI providerUri, Configuration conf)
    throws IOException {
  if (SCHEME_NAME.equals(providerUri.getScheme())) {
    URL origUrl = new URL(extractKMSPath(providerUri).toString());
    String authority = origUrl.getAuthority();
    // check for ';' which delimits the backup hosts
    if (Strings.isNullOrEmpty(authority)) {
      throw new IOException(
          "No valid authority in kms uri [" + origUrl + "]");
    }
    // Check if port is present in authority
    // In the current scheme, all hosts have to run on the same port
    int port = -1;
    String hostsPart = authority;
    if (authority.contains(":")) {
      String[] t = authority.split(":");
      try {
        port = Integer.parseInt(t[1]);
      } catch (Exception e) {
        throw new IOException(
            "Could not parse port in kms uri [" + origUrl + "]");
      }
      hostsPart = t[0];
    }
    return createProvider(providerUri, conf, origUrl, port, hostsPart);
  }
  return null;
}
 

@Override
public DelegationTokenIssuer[] getAdditionalTokenIssuers()
    throws IOException {
  KeyProvider keyProvider;
  try {
    keyProvider = getKeyProvider();
  } catch (IOException ioe) {
    LOG.debug("Error retrieving KeyProvider.", ioe);
    return null;
  }
  if (keyProvider instanceof DelegationTokenIssuer) {
    return new DelegationTokenIssuer[]{(DelegationTokenIssuer)keyProvider};
  }
  return null;
}
 

@Override
public KeyProvider createProvider(URI providerName, Configuration conf)
    throws IOException {
  if ("dummy".equals(providerName.getScheme())) {
    return new DummyKeyProvider(conf);
  }
  return null;
}
 

@SuppressWarnings("unchecked")
public static List toJSON(List<KeyProvider.KeyVersion> keyVersions) {
  List json = new ArrayList();
  if (keyVersions != null) {
    for (KeyProvider.KeyVersion version : keyVersions) {
      json.add(toJSON(version));
    }
  }
  return json;
}
 

@SuppressWarnings("unchecked")
public static List toJSON(List<KeyProvider.KeyVersion> keyVersions) {
  List json = new ArrayList();
  if (keyVersions != null) {
    for (KeyProvider.KeyVersion version : keyVersions) {
      json.add(KMSUtil.toJSON(version));
    }
  }
  return json;
}
 

@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 testCreation() throws Exception {
  Configuration conf = new Configuration();
  KeyProvider kp = new KMSClientProvider.Factory().createProvider(new URI(
      "kms://http@host1/kms/foo"), conf);
  assertTrue(kp instanceof KMSClientProvider);
  assertEquals("http://host1/kms/foo/v1/",
      ((KMSClientProvider) kp).getKMSUrl());

  kp = new KMSClientProvider.Factory().createProvider(new URI(
      "kms://http@host1;host2;host3/kms/foo"), conf);
  assertTrue(kp instanceof LoadBalancingKMSClientProvider);
  KMSClientProvider[] providers =
      ((LoadBalancingKMSClientProvider) kp).getProviders();
  assertEquals(3, providers.length);
  assertEquals(Sets.newHashSet("http://host1/kms/foo/v1/",
      "http://host2/kms/foo/v1/",
      "http://host3/kms/foo/v1/"),
      Sets.newHashSet(providers[0].getKMSUrl(),
          providers[1].getKMSUrl(),
          providers[2].getKMSUrl()));

  kp = new KMSClientProvider.Factory().createProvider(new URI(
      "kms://http@host1;host2;host3:16000/kms/foo"), conf);
  assertTrue(kp instanceof LoadBalancingKMSClientProvider);
  providers =
      ((LoadBalancingKMSClientProvider) kp).getProviders();
  assertEquals(3, providers.length);
  assertEquals(Sets.newHashSet("http://host1:16000/kms/foo/v1/",
      "http://host2:16000/kms/foo/v1/",
      "http://host3:16000/kms/foo/v1/"),
      Sets.newHashSet(providers[0].getKMSUrl(),
          providers[1].getKMSUrl(),
          providers[2].getKMSUrl()));
}
 

/**
 * This provider expects URIs in the following form :
 * kms://<PROTO>@<AUTHORITY>/<PATH>
 *
 * where :
 * - PROTO = http or https
 * - AUTHORITY = <HOSTS>[:<PORT>]
 * - HOSTS = <HOSTNAME>[;<HOSTS>]
 * - HOSTNAME = string
 * - PORT = integer
 *
 * If multiple hosts are provider, the Factory will create a
 * {@link LoadBalancingKMSClientProvider} that round-robins requests
 * across the provided list of hosts.
 */
@Override
public KeyProvider createProvider(URI providerUri, Configuration conf)
    throws IOException {
  if (SCHEME_NAME.equals(providerUri.getScheme())) {
    URL origUrl = new URL(extractKMSPath(providerUri).toString());
    String authority = origUrl.getAuthority();
    // check for ';' which delimits the backup hosts
    if (Strings.isNullOrEmpty(authority)) {
      throw new IOException(
          "No valid authority in kms uri [" + origUrl + "]");
    }
    // Check if port is present in authority
    // In the current scheme, all hosts have to run on the same port
    int port = -1;
    String hostsPart = authority;
    if (authority.contains(":")) {
      String[] t = authority.split(":");
      try {
        port = Integer.parseInt(t[1]);
      } catch (Exception e) {
        throw new IOException(
            "Could not parse port in kms uri [" + origUrl + "]");
      }
      hostsPart = t[0];
    }
    return createProvider(providerUri, conf, origUrl, port, hostsPart);
  }
  return null;
}
 

/**
 * Tests obtaining delegation token from stored key
 */
@Test(timeout = 120000)
public void testDelegationToken() throws Exception {
  UserGroupInformation.createRemoteUser("JobTracker");
  DistributedFileSystem dfs = cluster.getFileSystem();
  KeyProvider keyProvider = Mockito.mock(KeyProvider.class,
      withSettings().extraInterfaces(
          DelegationTokenExtension.class,
          CryptoExtension.class));
  Mockito.when(keyProvider.getConf()).thenReturn(conf);
  byte[] testIdentifier = "Test identifier for delegation token".getBytes();

  Token<?> testToken = new Token(testIdentifier, new byte[0],
      new Text(), new Text());
  Mockito.when(((DelegationTokenExtension)keyProvider).
      addDelegationTokens(anyString(), (Credentials)any())).
      thenReturn(new Token<?>[] { testToken });

  dfs.getClient().setKeyProvider(keyProvider);

  Credentials creds = new Credentials();
  final Token<?> tokens[] = dfs.addDelegationTokens("JobTracker", creds);
  DistributedFileSystem.LOG.debug("Delegation tokens: " +
      Arrays.asList(tokens));
  Assert.assertEquals(2, tokens.length);
  Assert.assertEquals(tokens[1], testToken);
  Assert.assertEquals(1, creds.numberOfTokens());
}