com.amazonaws.encryptionsdk.MasterKeyProvider Java Examples

@Test
public void testMixedKeysSingleDecrypt() {
    final SecretKeySpec k1 = new SecretKeySpec(generate(32), "AES");
    final JceMasterKey mk1 = JceMasterKey.getInstance(k1, "jce", "1", WRAPPING_ALG);
    final MockKMSClient kms = new MockKMSClient();
    final String arn2 = kms.createKey().getKeyMetadata().getArn();
    MasterKeyProvider<KmsMasterKey> prov = legacyConstruct(kms);
    KmsMasterKey mk2 = prov.getMasterKey(arn2);
    final MasterKeyProvider<?> mkp = MultipleProviderFactory.buildMultiProvider(mk1, mk2);

    AwsCrypto crypto = new AwsCrypto();
    CryptoResult<byte[], ?> ct = crypto.encryptData(mkp, PLAINTEXT);
    assertEquals(2, ct.getMasterKeyIds().size());

    CryptoResult<byte[], ?> result = crypto.decryptData(mk1, ct.getResult());
    assertArrayEquals(PLAINTEXT, result.getResult());
    // Only the first found key should be used
    assertEquals(1, result.getMasterKeys().size());
    assertEquals(mk1, result.getMasterKeys().get(0));

    result = crypto.decryptData(mk2, ct.getResult());
    assertArrayEquals(PLAINTEXT, result.getResult());
    // Only the first found key should be used
    assertEquals(1, result.getMasterKeys().size());
    assertEquals(mk2, result.getMasterKeys().get(0));
}
 

@SuppressWarnings("unchecked")
@Override
public DataKey<K> decryptDataKey(final CryptoAlgorithm algorithm,
        final Collection<? extends EncryptedDataKey> encryptedDataKeys,
        final Map<String, String> encryptionContext)
        throws UnsupportedProviderException, AwsCryptoException {
    final List<Exception> exceptions = new ArrayList<>();
    for (final MasterKeyProvider<? extends K> prov : providers_) {
        try {
            final DataKey<? extends K> result = prov
                    .decryptDataKey(algorithm, encryptedDataKeys, encryptionContext);
            if (result != null) {
                return (DataKey<K>) result;
            }
        } catch (final Exception ex) {
            exceptions.add(ex);
        }
    }
    throw buildCannotDecryptDksException(exceptions);
}
 

@Override
public K getMasterKey(final String provider, final String keyId) throws UnsupportedProviderException,
        NoSuchMasterKeyException {
    boolean foundProvider = false;
    for (final MasterKeyProvider<? extends K> prov : providers_) {
        if (prov.canProvide(provider)) {
            foundProvider = true;
            try {
                final K result = prov.getMasterKey(provider, keyId);
                if (result != null) {
                    return result;
                }
            } catch (final NoSuchMasterKeyException ex) {
                // swallow and continue
            }
        }
    }
    if (foundProvider) {
        throw new NoSuchMasterKeyException();
    } else {
        throw new UnsupportedProviderException(provider);
    }
}
 

@Test
public void testExplicitCredentials() throws Exception {
    AWSCredentials creds = new AWSCredentials() {
        @Override public String getAWSAccessKeyId() {
            throw new UsedExplicitCredentials();
        }

        @Override public String getAWSSecretKey() {
            throw new UsedExplicitCredentials();
        }
    };

    MasterKeyProvider<KmsMasterKey> mkp = new KmsMasterKeyProvider(creds, "arn:aws:kms:us-east-1:012345678901:key/foo-bar");
    assertExplicitCredentialsUsed(mkp);

    mkp = new KmsMasterKeyProvider(new AWSStaticCredentialsProvider(creds), "arn:aws:kms:us-east-1:012345678901:key/foo-bar");
    assertExplicitCredentialsUsed(mkp);
}
 

@Test
public void testMultipleKmsKeys() {
    final MockKMSClient kms = new MockKMSClient();
    final String arn1 = kms.createKey().getKeyMetadata().getArn();
    final String arn2 = kms.createKey().getKeyMetadata().getArn();
    MasterKeyProvider<KmsMasterKey> prov = legacyConstruct(kms, arn1, arn2);
    KmsMasterKey mk1 = prov.getMasterKey(arn1);

    AwsCrypto crypto = new AwsCrypto();
    CryptoResult<byte[], KmsMasterKey> ct = crypto.encryptData(prov, PLAINTEXT);
    assertEquals(2, ct.getMasterKeyIds().size());
    CryptoResult<byte[], KmsMasterKey> result = crypto.decryptData(prov, ct.getResult());
    assertArrayEquals(PLAINTEXT, result.getResult());
    // Only the first found key should be used
    assertEquals(1, result.getMasterKeys().size());
    assertEquals(mk1, result.getMasterKeys().get(0));
}
 

@Test
public void testMixedKeys() {
    final SecretKeySpec k1 = new SecretKeySpec(generate(32), "AES");
    final JceMasterKey mk1 = JceMasterKey.getInstance(k1, "jce", "1", WRAPPING_ALG);
    final MockKMSClient kms = new MockKMSClient();
    final String arn2 = kms.createKey().getKeyMetadata().getArn();
    MasterKeyProvider<KmsMasterKey> prov = legacyConstruct(kms);
    KmsMasterKey mk2 = prov.getMasterKey(arn2);
    final MasterKeyProvider<?> mkp = MultipleProviderFactory.buildMultiProvider(mk1, mk2);

    AwsCrypto crypto = new AwsCrypto();
    CryptoResult<byte[], ?> ct = crypto.encryptData(mkp, PLAINTEXT);
    assertEquals(2, ct.getMasterKeyIds().size());
    CryptoResult<byte[], ?> result = crypto.decryptData(mkp, ct.getResult());
    assertArrayEquals(PLAINTEXT, result.getResult());
    // Only the first found key should be used
    assertEquals(1, result.getMasterKeys().size());
    assertEquals(mk1, result.getMasterKeys().get(0));

    assertMultiReturnsKeys(mkp, mk1, mk2);
}
 

@Test
public void keystoreAndRawProvider() throws GeneralSecurityException, IOException {
    addEntry("key1");
    final SecretKeySpec k1 = new SecretKeySpec(generate(32), "AES");
    final JceMasterKey jcep = JceMasterKey.getInstance(k1, "jce", "1", "AES/GCM/NoPadding");
    final KeyStoreProvider ksp = new KeyStoreProvider(ks, PP, "KeyStore", "RSA/ECB/OAEPWithSHA-256AndMGF1Padding",
            "key1");

    MasterKeyProvider<JceMasterKey> multiProvider = MultipleProviderFactory.buildMultiProvider(JceMasterKey.class,
            jcep, ksp);

    assertEquals(jcep, multiProvider.getMasterKey("jce", "1"));

    final AwsCrypto crypto = new AwsCrypto();
    final CryptoResult<byte[], JceMasterKey> ct = crypto.encryptData(multiProvider, PLAINTEXT);
    assertEquals(2, ct.getMasterKeyIds().size());
    CryptoResult<byte[], JceMasterKey> result = crypto.decryptData(multiProvider, ct.getResult());
    assertArrayEquals(PLAINTEXT, result.getResult());
    assertEquals(jcep, result.getMasterKeys().get(0));

    // Decrypt just using each individually
    assertArrayEquals(PLAINTEXT, crypto.decryptData(jcep, ct.getResult()).getResult());
    assertArrayEquals(PLAINTEXT, crypto.decryptData(ksp, ct.getResult()).getResult());
}
 

@Test
public void testMultipleJceKeys() {
    final SecretKeySpec k1 = new SecretKeySpec(generate(32), "AES");
    final JceMasterKey mk1 = JceMasterKey.getInstance(k1, "jce", "1", WRAPPING_ALG);
    final SecretKeySpec k2 = new SecretKeySpec(generate(32), "AES");
    final JceMasterKey mk2 = JceMasterKey.getInstance(k2, "jce", "2", WRAPPING_ALG);
    final MasterKeyProvider<JceMasterKey> mkp = MultipleProviderFactory.buildMultiProvider(JceMasterKey.class,
            mk1, mk2);

    AwsCrypto crypto = new AwsCrypto();
    CryptoResult<byte[], JceMasterKey> ct = crypto.encryptData(mkp, PLAINTEXT);
    assertEquals(2, ct.getMasterKeyIds().size());
    CryptoResult<byte[], JceMasterKey> result = crypto.decryptData(mkp, ct.getResult());
    assertArrayEquals(PLAINTEXT, result.getResult());
    // Only the first found key should be used
    assertEquals(1, result.getMasterKeys().size());
    assertEquals(mk1, result.getMasterKeys().get(0));

    assertMultiReturnsKeys(mkp, mk1, mk2);
}
 

@Test
public void testMultipleJceKeysSingleDecrypt() {
    final SecretKeySpec k1 = new SecretKeySpec(generate(32), "AES");
    final JceMasterKey mk1 = JceMasterKey.getInstance(k1, "jce", "1", WRAPPING_ALG);
    final SecretKeySpec k2 = new SecretKeySpec(generate(32), "AES");
    final JceMasterKey mk2 = JceMasterKey.getInstance(k2, "jce", "2", WRAPPING_ALG);
    final MasterKeyProvider<JceMasterKey> mkp = MultipleProviderFactory.buildMultiProvider(JceMasterKey.class,
            mk1, mk2);

    AwsCrypto crypto = new AwsCrypto();
    CryptoResult<byte[], JceMasterKey> ct = crypto.encryptData(mkp, PLAINTEXT);
    assertEquals(2, ct.getMasterKeyIds().size());

    CryptoResult<byte[], JceMasterKey> result = crypto.decryptData(mk1, ct.getResult());
    assertArrayEquals(PLAINTEXT, result.getResult());
    // Only the first found key should be used
    assertEquals(1, result.getMasterKeys().size());
    assertEquals(mk1, result.getMasterKeys().get(0));

    result = crypto.decryptData(mk2, ct.getResult());
    assertArrayEquals(PLAINTEXT, result.getResult());
    // Only the first found key should be used
    assertEquals(1, result.getMasterKeys().size());
    assertEquals(mk2, result.getMasterKeys().get(0));
}
 

@Test
public void testMixedKeys() {
    final SecretKeySpec k1 = new SecretKeySpec(generate(32), "AES");
    final JceMasterKey mk1 = JceMasterKey.getInstance(k1, "jce", "1", WRAPPING_ALG);
    StaticMasterKey mk2 = new StaticMasterKey("mock1");
    final MasterKeyProvider<?> mkp = MultipleProviderFactory.buildMultiProvider(mk1, mk2);

    AwsCrypto crypto = new AwsCrypto();
    CryptoResult<byte[], ?> ct = crypto.encryptData(mkp, PLAINTEXT);
    assertEquals(2, ct.getMasterKeyIds().size());
    CryptoResult<byte[], ?> result = crypto.decryptData(mkp, ct.getResult());
    assertArrayEquals(PLAINTEXT, result.getResult());
    // Only the first found key should be used
    assertEquals(1, result.getMasterKeys().size());
    assertEquals(mk1, result.getMasterKeys().get(0));

    assertMultiReturnsKeys(mkp, mk1, mk2);
}
 

@Test
public void testMixedKeysSingleDecrypt() {
    final SecretKeySpec k1 = new SecretKeySpec(generate(32), "AES");
    final JceMasterKey mk1 = JceMasterKey.getInstance(k1, "jce", "1", WRAPPING_ALG);
    StaticMasterKey mk2 = new StaticMasterKey("mock1");

    final MasterKeyProvider<?> mkp = MultipleProviderFactory.buildMultiProvider(mk1, mk2);

    AwsCrypto crypto = new AwsCrypto();
    CryptoResult<byte[], ?> ct = crypto.encryptData(mkp, PLAINTEXT);
    assertEquals(2, ct.getMasterKeyIds().size());

    CryptoResult<byte[], ?> result = crypto.decryptData(mk1, ct.getResult());
    assertArrayEquals(PLAINTEXT, result.getResult());
    // Only the first found key should be used
    assertEquals(1, result.getMasterKeys().size());
    assertEquals(mk1, result.getMasterKeys().get(0));

    result = crypto.decryptData(mk2, ct.getResult());
    assertArrayEquals(PLAINTEXT, result.getResult());
    // Only the first found key should be used
    assertEquals(1, result.getMasterKeys().size());
    assertEquals(mk2, result.getMasterKeys().get(0));
}
 

@Test
public void whenBuilderCloned_credentialsAndConfigurationAreRetained() throws Exception {
    AWSCredentialsProvider customProvider1 = spy(new DefaultAWSCredentialsProviderChain());
    AWSCredentialsProvider customProvider2 = spy(new DefaultAWSCredentialsProviderChain());

    KmsMasterKeyProvider.Builder builder = KmsMasterKeyProvider.builder()
            .withCredentials(customProvider1)
            .withKeysForEncryption(KMSTestFixtures.TEST_KEY_IDS[0]);

    KmsMasterKeyProvider.Builder builder2 = builder.clone();

    // This will mutate the first builder to add the new key and change the creds, but leave the clone unchanged.
    MasterKeyProvider<?> mkp2 = builder.withKeysForEncryption(KMSTestFixtures.TEST_KEY_IDS[1]).withCredentials(customProvider2).build();
    MasterKeyProvider<?> mkp1 = builder2.build();

    CryptoResult<byte[], ?> result = new AwsCrypto().encryptData(mkp1, new byte[0]);

    assertEquals(KMSTestFixtures.TEST_KEY_IDS[0], result.getMasterKeyIds().get(0));
    assertEquals(1, result.getMasterKeyIds().size());
    verify(customProvider1, atLeastOnce()).getCredentials();
    verify(customProvider2, never()).getCredentials();

    reset(customProvider1, customProvider2);

    result = new AwsCrypto().encryptData(mkp2, new byte[0]);

    assertTrue(result.getMasterKeyIds().contains(KMSTestFixtures.TEST_KEY_IDS[0]));
    assertTrue(result.getMasterKeyIds().contains(KMSTestFixtures.TEST_KEY_IDS[1]));
    assertEquals(2, result.getMasterKeyIds().size());
    verify(customProvider1, never()).getCredentials();
    verify(customProvider2, atLeastOnce()).getCredentials();
}
 

private void assertExplicitCredentialsUsed(final MasterKeyProvider<KmsMasterKey> mkp) {
    try {
        MasterKeyRequest mkr = MasterKeyRequest.newBuilder()
                                               .setEncryptionContext(Collections.emptyMap())
                                               .setStreaming(true)
                                               .build();
        mkp.getMasterKeysForEncryption(mkr)
           .forEach(mk -> mk.generateDataKey(ALG_AES_128_GCM_IV12_TAG16_NO_KDF, Collections.emptyMap()));

        fail("Expected exception");
    } catch (UsedExplicitCredentials e) {
        // ok
    }
}
 

@Bean("decryptCryptoMaterialsManager")
public CryptoMaterialsManager decryptCryptoMaterialsManager(
    @Value("${cerberus.encryption.cmk.arns}") String cmkArns,
    @Value("${cerberus.encryption.cache.enabled:#{false}}") boolean cacheEnabled,
    @Value("${cerberus.encryption.cache.decrypt.maxSize:1000}") int decryptMaxSize,
    @Value("${cerberus.encryption.cache.decrypt.maxAgeInSeconds:60}") int decryptMaxAge,
    Region currentRegion,
    MetricsService metricsService) {
  MasterKeyProvider<KmsMasterKey> keyProvider = initializeKeyProvider(cmkArns, currentRegion);
  if (cacheEnabled) {
    log.info(
        "Initializing caching decryptCryptoMaterialsManager with CMK: {}, maxSize: {}, maxAge: {}",
        cmkArns,
        decryptMaxSize,
        decryptMaxAge);
    CryptoMaterialsCache cache =
        new MetricReportingCryptoMaterialsCache(decryptMaxAge, metricsService);
    CryptoMaterialsManager cachingCmm =
        CachingCryptoMaterialsManager.newBuilder()
            .withMasterKeyProvider(keyProvider)
            .withCache(cache)
            .withMaxAge(decryptMaxAge, TimeUnit.SECONDS)
            .build();
    return cachingCmm;
  } else {
    log.info("Initializing decryptCryptoMaterialsManager with CMK: {}", cmkArns);
    return new DefaultCryptoMaterialsManager(keyProvider);
  }
}
 

@Test
public void testMultipleRegionKmsKeys() {
    final MockKMSClient us_east_1 = new MockKMSClient();
    us_east_1.setRegion(Region.getRegion(Regions.US_EAST_1));
    final MockKMSClient eu_west_1 = new MockKMSClient();
    eu_west_1.setRegion(Region.getRegion(Regions.EU_WEST_1));
    final String arn1 = us_east_1.createKey().getKeyMetadata().getArn();
    final String arn2 = eu_west_1.createKey().getKeyMetadata().getArn();
    KmsMasterKeyProvider provE = legacyConstruct(us_east_1, Region.getRegion(Regions.US_EAST_1));
    KmsMasterKeyProvider provW = legacyConstruct(eu_west_1, Region.getRegion(Regions.EU_WEST_1));
    KmsMasterKey mk1 = provE.getMasterKey(arn1);
    KmsMasterKey mk2 = provW.getMasterKey(arn2);

    final MasterKeyProvider<KmsMasterKey> mkp = MultipleProviderFactory.buildMultiProvider(KmsMasterKey.class,
                                                                                           mk1, mk2);
    AwsCrypto crypto = new AwsCrypto();
    CryptoResult<byte[], KmsMasterKey> ct = crypto.encryptData(mkp, PLAINTEXT);
    assertEquals(2, ct.getMasterKeyIds().size());

    CryptoResult<byte[], KmsMasterKey> result = crypto.decryptData(mk1, ct.getResult());
    assertArrayEquals(PLAINTEXT, result.getResult());
    assertEquals(1, result.getMasterKeys().size());
    assertEquals(mk1, result.getMasterKeys().get(0));

    result = crypto.decryptData(mk2, ct.getResult());
    assertArrayEquals(PLAINTEXT, result.getResult());
    assertEquals(1, result.getMasterKeys().size());
    assertEquals(mk2, result.getMasterKeys().get(0));

    assertMultiReturnsKeys(mkp, mk1, mk2);

    // Delete one of the two keys and ensure it's still decryptable
    us_east_1.deleteKey(arn1);

    result = crypto.decryptData(mkp, ct.getResult());
    assertArrayEquals(PLAINTEXT, result.getResult());
    // Only the first found key should be used
    assertEquals(1, result.getMasterKeys().size());
    assertEquals(mk2, result.getMasterKeys().get(0));
}
 

@Test
public void testMultipleKmsKeysSingleDecrypt() {
    final MockKMSClient kms = new MockKMSClient();
    final String arn1 = kms.createKey().getKeyMetadata().getArn();
    final String arn2 = kms.createKey().getKeyMetadata().getArn();
    MasterKeyProvider<KmsMasterKey> prov = legacyConstruct(kms, arn1, arn2);
    KmsMasterKey mk1 = prov.getMasterKey(arn1);
    KmsMasterKey mk2 = prov.getMasterKey(arn2);

    AwsCrypto crypto = new AwsCrypto();
    CryptoResult<byte[], KmsMasterKey> ct = crypto.encryptData(prov, PLAINTEXT);
    assertEquals(2, ct.getMasterKeyIds().size());

    CryptoResult<byte[], KmsMasterKey> result = crypto.decryptData(mk1, ct.getResult());
    assertArrayEquals(PLAINTEXT, result.getResult());
    assertEquals(1, result.getMasterKeys().size());
    assertEquals(mk1, result.getMasterKeys().get(0));

    result = crypto.decryptData(mk2, ct.getResult());
    assertArrayEquals(PLAINTEXT, result.getResult());
    assertEquals(1, result.getMasterKeys().size());
    assertEquals(mk2, result.getMasterKeys().get(0));

    // Delete one of the two keys and ensure it's still decryptable
    kms.deleteKey(arn1);

    result = crypto.decryptData(prov, ct.getResult());
    assertArrayEquals(PLAINTEXT, result.getResult());
    // Only the first found key should be used
    assertEquals(1, result.getMasterKeys().size());
    assertEquals(mk2, result.getMasterKeys().get(0));
}
 

@Test
public void testNoKeyMKP() throws Exception {
    AWSCredentials creds = new ThrowingCredentials();

    MasterKeyRequest mkr = MasterKeyRequest.newBuilder()
                                           .setEncryptionContext(Collections.emptyMap())
                                           .setStreaming(true)
                                           .build();

    MasterKeyProvider<KmsMasterKey> mkp = new KmsMasterKeyProvider(creds);
    assertTrue(mkp.getMasterKeysForEncryption(mkr).isEmpty());

    mkp = new KmsMasterKeyProvider(new AWSStaticCredentialsProvider(creds));
    assertTrue(mkp.getMasterKeysForEncryption(mkr).isEmpty());
}
 

@Override
public K getMasterKey(final String keyId) throws UnsupportedProviderException, NoSuchMasterKeyException {
    for (final MasterKeyProvider<? extends K> prov : providers_) {
        try {
            final K result = prov.getMasterKey(keyId);
            if (result != null) {
                return result;
            }
        } catch (final NoSuchMasterKeyException ex) {
            // swallow and continue
        }
    }
    throw new NoSuchMasterKeyException();
}
 

private MultiProvider(final List<? extends MasterKeyProvider<? extends K>> providers) {
    Utils.assertNonNull(providers, "providers");
    if (providers.isEmpty()) {
        throw new IllegalArgumentException("providers must not be empty");
    }
    providers_ = new ArrayList<>(providers);
}
 

@Override
public List<K> getMasterKeysForEncryption(final MasterKeyRequest request) {
    final List<K> result = new ArrayList<>();
    for (final MasterKeyProvider<? extends K> prov : providers_) {
        result.addAll(prov.getMasterKeysForEncryption(request));
    }
    return result;
}
 

@Bean("encryptCryptoMaterialsManager")
public CryptoMaterialsManager encryptCryptoMaterialsManager(
    @Value("${cerberus.encryption.cmk.arns}") String cmkArns,
    @Value("${cerberus.encryption.cache.enabled:false}") boolean cacheEnabled,
    @Value("${cerberus.encryption.cache.encrypt.maxSize:100}") int encryptMaxSize,
    @Value("${cerberus.encryption.cache.encrypt.maxAgeInSeconds:60}") int encryptMaxAge,
    @Value("${cerberus.encryption.cache.encrypt.messageUseLimit:100}") int encryptMessageUseLimit,
    Region currentRegion,
    MetricsService metricsService) {
  MasterKeyProvider<KmsMasterKey> keyProvider = initializeKeyProvider(cmkArns, currentRegion);
  if (cacheEnabled) {
    log.info(
        "Initializing caching encryptCryptoMaterialsManager with CMK: {}, maxSize: {}, maxAge: {}, "
            + "messageUseLimit: {}",
        cmkArns,
        encryptMaxSize,
        encryptMaxAge,
        encryptMessageUseLimit);
    CryptoMaterialsCache cache =
        new MetricReportingCryptoMaterialsCache(encryptMaxSize, metricsService);
    CryptoMaterialsManager cachingCmm =
        CachingCryptoMaterialsManager.newBuilder()
            .withMasterKeyProvider(keyProvider)
            .withCache(cache)
            .withMaxAge(encryptMaxAge, TimeUnit.SECONDS)
            .withMessageUseLimit(encryptMessageUseLimit)
            .build();
    return cachingCmm;
  } else {
    log.info("Initializing encryptCryptoMaterialsManager with CMK: {}", cmkArns);
    return new DefaultCryptoMaterialsManager(keyProvider);
  }
}
 

private static void standardEncrypt(final String kmsArn, final String fileName) throws Exception {
    // Encrypt with the KMS CMK and the escrowed public key
    // 1. Instantiate the SDK
    final AwsCrypto crypto = new AwsCrypto();

    // 2. Instantiate a KMS master key provider
    final KmsMasterKeyProvider kms = new KmsMasterKeyProvider(kmsArn);
    
    // 3. Instantiate a JCE master key provider
    // Because the user does not have access to the private escrow key,
    // they pass in "null" for the private key parameter.
    final JceMasterKey escrowPub = JceMasterKey.getInstance(publicEscrowKey, null, "Escrow", "Escrow",
            "RSA/ECB/OAEPWithSHA-512AndMGF1Padding");

    // 4. Combine the providers into a single master key provider
    final MasterKeyProvider<?> provider = MultipleProviderFactory.buildMultiProvider(kms, escrowPub);

    // 5. Encrypt the file
    // To simplify the code, we omit the encryption context. Production code should always 
    // use an encryption context. For an example, see the other SDK samples.
    final FileInputStream in = new FileInputStream(fileName);
    final FileOutputStream out = new FileOutputStream(fileName + ".encrypted");
    final CryptoOutputStream<?> encryptingStream = crypto.createEncryptingStream(provider, out);

    IOUtils.copy(in, encryptingStream);
    in.close();
    encryptingStream.close();
}
 

private static void standardDecrypt(final String kmsArn, final String fileName) throws Exception {
    // Decrypt with the KMS CMK and the escrow public key. You can use a combined provider,
    // as shown here, or just the KMS master key provider.

    // 1. Instantiate the SDK
    final AwsCrypto crypto = new AwsCrypto();

    // 2. Instantiate a KMS master key provider
    final KmsMasterKeyProvider kms = new KmsMasterKeyProvider(kmsArn);
    
    // 3. Instantiate a JCE master key provider
    // Because the user does not have access to the private 
    // escrow key, they pass in "null" for the private key parameter.
    final JceMasterKey escrowPub = JceMasterKey.getInstance(publicEscrowKey, null, "Escrow", "Escrow",
            "RSA/ECB/OAEPWithSHA-512AndMGF1Padding");

    // 4. Combine the providers into a single master key provider
    final MasterKeyProvider<?> provider = MultipleProviderFactory.buildMultiProvider(kms, escrowPub);

    // 5. Decrypt the file
    // To simplify the code, we omit the encryption context. Production code should always 
    // use an encryption context. For an example, see the other SDK samples.
    final FileInputStream in = new FileInputStream(fileName + ".encrypted");
    final FileOutputStream out = new FileOutputStream(fileName + ".decrypted");
    final CryptoOutputStream<?> decryptingStream = crypto.createDecryptingStream(provider, out);
    IOUtils.copy(in, decryptingStream);
    in.close();
    decryptingStream.close();
}
 

private void assertMultiReturnsKeys(MasterKeyProvider<?> mkp, MasterKey<?>... mks) {
    for (MasterKey<?> mk : mks) {
        assertEquals(mk, mkp.getMasterKey(mk.getKeyId()));
        assertEquals(mk, mkp.getMasterKey(mk.getProviderId(), mk.getKeyId()));
    }
}
 

static KmsMasterKey getInstance(final Supplier<AWSKMS> kms, final String id,
        final MasterKeyProvider<KmsMasterKey> provider) {
    return new KmsMasterKey(kms, id, provider);
}
 

private KmsMasterKey(final Supplier<AWSKMS> kms, final String id, final MasterKeyProvider<KmsMasterKey> provider) {
    kms_ = kms;
    id_ = id;
    sourceProvider_ = provider;
}
 

@Test
public void testGrantTokenPassthrough_usingMKPWithers() throws Exception {
    MockKMSClient client = spy(new MockKMSClient());

    RegionalClientSupplier supplier = mock(RegionalClientSupplier.class);
    when(supplier.getClient(any())).thenReturn(client);

    String key1 = client.createKey().getKeyMetadata().getArn();
    String key2 = client.createKey().getKeyMetadata().getArn();

    KmsMasterKeyProvider mkp0 = KmsMasterKeyProvider.builder()
                                                    .withDefaultRegion("us-west-2")
                                                    .withCustomClientFactory(supplier)
                                                    .withKeysForEncryption(key1, key2)
                                                    .build();

    MasterKeyProvider<?> mkp = mkp0.withGrantTokens("foo");

    byte[] ciphertext = new AwsCrypto().encryptData(mkp, new byte[0]).getResult();

    ArgumentCaptor<GenerateDataKeyRequest> gdkr = ArgumentCaptor.forClass(GenerateDataKeyRequest.class);
    verify(client, times(1)).generateDataKey(gdkr.capture());

    assertEquals(key1, gdkr.getValue().getKeyId());
    assertEquals(1, gdkr.getValue().getGrantTokens().size());
    assertEquals("foo", gdkr.getValue().getGrantTokens().get(0));

    ArgumentCaptor<EncryptRequest> er = ArgumentCaptor.forClass(EncryptRequest.class);
    verify(client, times(1)).encrypt(er.capture());

    assertEquals(key2, er.getValue().getKeyId());
    assertEquals(1, er.getValue().getGrantTokens().size());
    assertEquals("foo", er.getValue().getGrantTokens().get(0));

    mkp = mkp0.withGrantTokens(Arrays.asList("bar"));

    new AwsCrypto().decryptData(mkp, ciphertext);

    ArgumentCaptor<DecryptRequest> decrypt = ArgumentCaptor.forClass(DecryptRequest.class);
    verify(client, times(1)).decrypt(decrypt.capture());

    assertEquals(1, decrypt.getValue().getGrantTokens().size());
    assertEquals("bar", decrypt.getValue().getGrantTokens().get(0));

    verify(supplier, atLeastOnce()).getClient("us-west-2");
    verifyNoMoreInteractions(supplier);
}
 

@Test
public void testGrantTokenPassthrough_usingMKsetCall() throws Exception {
    MockKMSClient client = spy(new MockKMSClient());

    RegionalClientSupplier supplier = mock(RegionalClientSupplier.class);
    when(supplier.getClient(any())).thenReturn(client);

    String key1 = client.createKey().getKeyMetadata().getArn();
    String key2 = client.createKey().getKeyMetadata().getArn();

    KmsMasterKeyProvider mkp0 = KmsMasterKeyProvider.builder()
                                                   .withDefaultRegion("us-west-2")
                                                   .withCustomClientFactory(supplier)
                                                   .withKeysForEncryption(key1, key2)
                                                   .build();
    KmsMasterKey mk1 = mkp0.getMasterKey(key1);
    KmsMasterKey mk2 = mkp0.getMasterKey(key2);

    mk1.setGrantTokens(singletonList("foo"));
    mk2.setGrantTokens(singletonList("foo"));

    MasterKeyProvider<?> mkp = buildMultiProvider(mk1, mk2);

    byte[] ciphertext = new AwsCrypto().encryptData(mkp, new byte[0]).getResult();

    ArgumentCaptor<GenerateDataKeyRequest> gdkr = ArgumentCaptor.forClass(GenerateDataKeyRequest.class);
    verify(client, times(1)).generateDataKey(gdkr.capture());

    assertEquals(key1, gdkr.getValue().getKeyId());
    assertEquals(1, gdkr.getValue().getGrantTokens().size());
    assertEquals("foo", gdkr.getValue().getGrantTokens().get(0));

    ArgumentCaptor<EncryptRequest> er = ArgumentCaptor.forClass(EncryptRequest.class);
    verify(client, times(1)).encrypt(er.capture());

    assertEquals(key2, er.getValue().getKeyId());
    assertEquals(1, er.getValue().getGrantTokens().size());
    assertEquals("foo", er.getValue().getGrantTokens().get(0));

    new AwsCrypto().decryptData(mkp, ciphertext);

    ArgumentCaptor<DecryptRequest> decrypt = ArgumentCaptor.forClass(DecryptRequest.class);
    verify(client, times(1)).decrypt(decrypt.capture());

    assertEquals(1, decrypt.getValue().getGrantTokens().size());
    assertEquals("foo", decrypt.getValue().getGrantTokens().get(0));

    verify(supplier, atLeastOnce()).getClient("us-west-2");
    verifyNoMoreInteractions(supplier);
}
 

private void assertMultiReturnsKeys(MasterKeyProvider<?> mkp, MasterKey<?>... mks) {
    for (MasterKey<?> mk : mks) {
        assertEquals(mk, mkp.getMasterKey(mk.getKeyId()));
        assertEquals(mk, mkp.getMasterKey(mk.getProviderId(), mk.getKeyId()));
    }
}
 

public static <K extends MasterKey<K>> MasterKeyProvider<K> buildMultiProvider(final Class<K> masterKeyClass,
        final List<? extends MasterKeyProvider<? extends K>> providers) {
    return new MultiProvider<K>(providers);
}