Align Key Vault JCA authentication with other KV libraries

sdk/keyvault/azure-security-keyvault-jca/src/main/java/com/azure/security/keyvault/jca/KeyVaultKeyStore.java

@@ -41,7 +41,6 @@
  * @see KeyStoreSpi
  */
 public final class KeyVaultKeyStore extends KeyStoreSpi {
-
     /**
      * Stores the key-store name.
      */
@@ -98,13 +97,13 @@ public final class KeyVaultKeyStore extends KeyStoreSpi {
      * Store the path where the well know certificate is placed
      */
     final String wellKnowPath = Optional.ofNullable(System.getProperty("azure.cert-path.well-known"))
-                                        .orElse("/etc/certs/well-known/");
+        .orElse("/etc/certs/well-known/");
 
     /**
      * Store the path where the custom certificate is placed
      */
     final String customPath = Optional.ofNullable(System.getProperty("azure.cert-path.custom"))
-                                      .orElse("/etc/certs/custom/");
+        .orElse("/etc/certs/custom/");
 
     /**
      * Constructor.
@@ -122,41 +121,50 @@ public final class KeyVaultKeyStore extends KeyStoreSpi {
      */
     public KeyVaultKeyStore() {
         LOGGER.log(FINE, "Constructing KeyVaultKeyStore.");
+
         creationDate = new Date();
         String keyVaultUri = System.getProperty("azure.keyvault.uri");
-        String loginUri = System.getProperty("azure.login.uri");
         String tenantId = System.getProperty("azure.keyvault.tenant-id");
         String clientId = System.getProperty("azure.keyvault.client-id");
         String clientSecret = System.getProperty("azure.keyvault.client-secret");
         String managedIdentity = System.getProperty("azure.keyvault.managed-identity");
+        boolean disableChallengeResourceVerification =
+            Boolean.parseBoolean(System.getProperty("azure.keyvault.disable-challenge-resource-verification"));
         long refreshInterval = getRefreshInterval();
         refreshCertificatesWhenHaveUnTrustCertificate =
             Optional.of("azure.keyvault.jca.refresh-certificates-when-have-un-trust-certificate")
-                    .map(System::getProperty)
-                    .map(Boolean::parseBoolean)
-                    .orElse(false);
+                .map(System::getProperty)
+                .map(Boolean::parseBoolean)
+                .orElse(false);
+
         jreCertificates = JreCertificates.getInstance();
         LOGGER.log(FINE, String.format("Loaded jre certificates: %s.", jreCertificates.getAliases()));
+
         wellKnowCertificates = SpecificPathCertificates.getSpecificPathCertificates(wellKnowPath);
         LOGGER.log(FINE, String.format("Loaded well known certificates: %s.", wellKnowCertificates.getAliases()));
+
         customCertificates = SpecificPathCertificates.getSpecificPathCertificates(customPath);
         LOGGER.log(FINE, String.format("Loaded custom certificates: %s.", customCertificates.getAliases()));
-        keyVaultCertificates = new KeyVaultCertificates(
-            refreshInterval, keyVaultUri, loginUri, tenantId, clientId, clientSecret, managedIdentity);
+
+        keyVaultCertificates = new KeyVaultCertificates(refreshInterval, keyVaultUri, tenantId, clientId, clientSecret,
+            managedIdentity, disableChallengeResourceVerification);
         LOGGER.log(FINE, String.format("Loaded Key Vault certificates: %s.", keyVaultCertificates.getAliases()));
+
         classpathCertificates = new ClasspathCertificates();
         LOGGER.log(FINE, String.format("Loaded classpath certificates: %s.", classpathCertificates.getAliases()));
-        allCertificates = Arrays.asList(
-            jreCertificates, wellKnowCertificates, customCertificates, keyVaultCertificates, classpathCertificates);
+
+        allCertificates = Arrays.asList(jreCertificates, wellKnowCertificates, customCertificates, keyVaultCertificates,
+            classpathCertificates);
     }
 
     Long getRefreshInterval() {
-        return Stream.of("azure.keyvault.jca.certificates-refresh-interval-in-ms", "azure.keyvault.jca.certificates-refresh-interval")
-                     .map(System::getProperty)
-                     .filter(Objects::nonNull)
-                     .map(Long::valueOf)
-                     .findFirst()
-                     .orElse(0L);
+        return Stream.of("azure.keyvault.jca.certificates-refresh-interval-in-ms",
+                "azure.keyvault.jca.certificates-refresh-interval")
+            .map(System::getProperty)
+            .filter(Objects::nonNull)
+            .map(Long::valueOf)
+            .findFirst()
+            .orElse(0L);
     }
 
     /**
@@ -169,17 +177,23 @@ Long getRefreshInterval() {
      * @throws KeyStoreException when no Provider supports a KeyStoreSpi implementation for the specified type
      * @throws IOException when an I/O error occurs.
      */
-    public static KeyStore getKeyVaultKeyStoreBySystemProperty() throws CertificateException, NoSuchAlgorithmException, KeyStoreException, IOException {
+    public static KeyStore getKeyVaultKeyStoreBySystemProperty() throws CertificateException, NoSuchAlgorithmException,
+        KeyStoreException, IOException {
+
         KeyStore keyStore = KeyStore.getInstance(KeyVaultJcaProvider.PROVIDER_NAME);
-        KeyVaultLoadStoreParameter parameter = new KeyVaultLoadStoreParameter(
-            System.getProperty("azure.keyvault.uri"),
-            System.getProperty("azure.login.uri"),
-            System.getProperty("azure.keyvault.tenant-id"),
-            System.getProperty("azure.keyvault.client-id"),
-            System.getProperty("azure.keyvault.client-secret"),
-            System.getProperty("azure.keyvault.managed-identity"));
+        KeyVaultLoadStoreParameter keyVaultLoadStoreParameter =
+            new KeyVaultLoadStoreParameter(
+                System.getProperty("azure.keyvault.uri"),
+                System.getProperty("azure.keyvault.tenant-id"),
+                System.getProperty("azure.keyvault.client-id"),
+                System.getProperty("azure.keyvault.client-secret"),
+                System.getProperty("azure.keyvault.managed-identity"));
+
+        if (Boolean.parseBoolean(System.getProperty("azure.keyvault.disable-challenge-resource-verification"))) {
+            keyVaultLoadStoreParameter.disableChallengeResourceVerification();
+        }
 
-        keyStore.load(parameter);
+        keyStore.load(keyVaultLoadStoreParameter);
 
         return keyStore;
     }
@@ -241,16 +255,17 @@ public boolean engineEntryInstanceOf(String alias, Class<? extends KeyStore.Entr
     @Override
     public Certificate engineGetCertificate(String alias) {
         Certificate certificate = allCertificates.stream()
-                                                 .map(AzureCertificates::getCertificates)
-                                                 .filter(a -> a.containsKey(alias))
-                                                 .findFirst()
-                                                 .map(certificates -> certificates.get(alias))
-                                                 .orElse(null);
+            .map(AzureCertificates::getCertificates)
+            .filter(a -> a.containsKey(alias))
+            .findFirst()
+            .map(certificates -> certificates.get(alias))
+            .orElse(null);
 
         if (refreshCertificatesWhenHaveUnTrustCertificate && certificate == null) {
             keyVaultCertificates.refreshCertificates();
             certificate = keyVaultCertificates.getCertificates().get(alias);
         }
+
         return certificate;
     }
 
@@ -264,20 +279,25 @@ public Certificate engineGetCertificate(String alias) {
     @Override
     public String engineGetCertificateAlias(Certificate cert) {
         String alias = null;
+
         if (cert != null) {
             List<String> aliasList = getAllAliases();
 
             for (String candidateAlias : aliasList) {
                 Certificate certificate = engineGetCertificate(candidateAlias);
+
                 if (certificate.equals(cert)) {
                     alias = candidateAlias;
+
                     break;
                 }
             }
         }
+
         if (refreshCertificatesWhenHaveUnTrustCertificate && alias == null) {
             alias = keyVaultCertificates.refreshAndGetAliasByCertificate(cert);
         }
+
         return alias;
     }
 
@@ -293,10 +313,12 @@ public String engineGetCertificateAlias(Certificate cert) {
     public Certificate[] engineGetCertificateChain(String alias) {
         Certificate[] chain = null;
         Certificate certificate = engineGetCertificate(alias);
+
         if (certificate != null) {
             chain = new Certificate[1];
             chain[0] = certificate;
         }
+
         return chain;
     }
 
@@ -322,7 +344,9 @@ public Date engineGetCreationDate(String alias) {
      * @exception UnrecoverableEntryException if the specified {@code protParam} were insufficient or invalid
      */
     @Override
-    public KeyStore.Entry engineGetEntry(String alias, KeyStore.ProtectionParameter protParam) throws KeyStoreException, NoSuchAlgorithmException, UnrecoverableEntryException {
+    public KeyStore.Entry engineGetEntry(String alias, KeyStore.ProtectionParameter protParam) throws KeyStoreException,
+        NoSuchAlgorithmException, UnrecoverableEntryException {
+
         return super.engineGetEntry(alias, protParam);
     }
 
@@ -336,11 +360,11 @@ public KeyStore.Entry engineGetEntry(String alias, KeyStore.ProtectionParameter
     @Override
     public Key engineGetKey(String alias, char[] password) {
         return allCertificates.stream()
-                              .map(AzureCertificates::getCertificateKeys)
-                              .filter(a -> a.containsKey(alias))
-                              .findFirst()
-                              .map(certificateKeys -> certificateKeys.get(alias))
-                              .orElse(null);
+            .map(AzureCertificates::getCertificateKeys)
+            .filter(a -> a.containsKey(alias))
+            .findFirst()
+            .map(certificateKeys -> certificateKeys.get(alias))
+            .orElse(null);
     }
 
     /**
@@ -368,17 +392,18 @@ public boolean engineIsKeyEntry(String alias) {
     /**
      * Loads the keystore using the given {@code KeyStore.LoadStoreParameter}.
      *
-     * @param param the {@code KeyStore.LoadStoreParameter} that specifies how to load the keystore, which may be
-     * {@code null}.
+     * @param param the {@code KeyStore.LoadStoreParameter}
+     *          that specifies how to load the keystore,
+     *          which may be {@code null}
      */
     @Override
     public void engineLoad(KeyStore.LoadStoreParameter param) {
         if (param instanceof KeyVaultLoadStoreParameter) {
             KeyVaultLoadStoreParameter parameter = (KeyVaultLoadStoreParameter) param;
 
-            keyVaultCertificates.updateKeyVaultClient(
-                parameter.getUri(), parameter.getLoginUri(), parameter.getTenantId(), parameter.getClientId(),
-                parameter.getClientSecret(), parameter.getManagedIdentity());
+            keyVaultCertificates.updateKeyVaultClient(parameter.getUri(), parameter.getTenantId(),
+                parameter.getClientId(), parameter.getClientSecret(), parameter.getManagedIdentity(),
+                parameter.isChallengeResourceVerificationDisabled());
         }
 
         classpathCertificates.loadCertificatesFromClasspath();
@@ -398,6 +423,7 @@ public void engineLoad(InputStream stream, char[] password) {
     private List<String> getAllAliases() {
         List<String> allAliases = new ArrayList<>(jreCertificates.getAliases());
         Map<String, List<String>> aliasLists = new HashMap<>();
+
         aliasLists.put("well known certificates", wellKnowCertificates.getAliases());
         aliasLists.put("custom certificates", customCertificates.getAliases());
         aliasLists.put("key vault certificates", keyVaultCertificates.getAliases());
@@ -410,6 +436,7 @@ private List<String> getAllAliases() {
                 allAliases.add(alias);
             }
         }));
+
         return allAliases;
     }
 
@@ -423,8 +450,10 @@ private List<String> getAllAliases() {
     public void engineSetCertificateEntry(String alias, Certificate certificate) {
         if (getAllAliases().contains(alias)) {
             LOGGER.log(WARNING, "Cannot overwrite own certificate");
+
             return;
         }
+
         classpathCertificates.setCertificateEntry(alias, certificate);
     }
 
@@ -439,7 +468,9 @@ public void engineSetCertificateEntry(String alias, Certificate certificate) {
      * @throws KeyStoreException if this operation fails
      */
     @Override
-    public void engineSetEntry(String alias, KeyStore.Entry entry, KeyStore.ProtectionParameter protParam) throws KeyStoreException {
+    public void engineSetEntry(String alias, KeyStore.Entry entry, KeyStore.ProtectionParameter protParam)
+        throws KeyStoreException {
+
         super.engineSetEntry(alias, entry, protParam);
     }
 
@@ -494,6 +525,4 @@ public void engineStore(OutputStream stream, char[] password) {
     @Override
     public void engineStore(KeyStore.LoadStoreParameter param) {
     }
-
-
 }