Support connect to multi keyvault for keyless (Azure#24718)

Author: zhichengliu12581

eng/code-quality-reports/src/main/resources/checkstyle/checkstyle-suppressions.xml

@@ -45,6 +45,9 @@
   <!-- Fields tenantId, clientId and clientSecret are not set in all constructors. -->
   <suppress checks="com.azure.tools.checkstyle.checks.EnforceFinalFieldsCheck" files="com.azure.security.keyvault.jca.implementation.KeyVaultClient"/>
 
+  <!-- We should move KeyVaultPrivateKey into implementation. -->
+  <suppress checks="com.azure.tools.checkstyle.checks.NoImplInPublicAPI" files="com.azure.security.keyvault.jca.KeyVaultPrivateKey"/>
+
   <!-- Fields defaultTrustManager and keyStore might be set more than once. -->
   <suppress checks="com.azure.tools.checkstyle.checks.EnforceFinalFieldsCheck" files="com.azure.security.keyvault.jca.KeyVaultTrustManager"/>
 

eng/code-quality-reports/src/main/resources/spotbugs/spotbugs-exclude.xml

@@ -2414,6 +2414,10 @@
     <Method name="~(get|post)"/>
     <Bug pattern="RCN_REDUNDANT_NULLCHECK_WOULD_HAVE_BEEN_A_NPE"/>
   </Match>
+  <Match>
+    <Class name="com.azure.security.keyvault.jca.KeyVaultPrivateKey"/>
+    <Bug pattern="SE_BAD_FIELD"/>
+  </Match>
 
   <!-- Returning a new copy of the object is not necessary for queuing system use cases.
   Client do not change content of message once message is created. Array copy would degrade the performance.

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

@@ -3,6 +3,8 @@
 
 package com.azure.security.keyvault.jca;
 
+import com.azure.security.keyvault.jca.implementation.KeyVaultClient;
+
 import javax.crypto.SecretKey;
 import java.security.PrivateKey;
 
@@ -20,14 +22,35 @@ public class KeyVaultPrivateKey implements PrivateKey, SecretKey {
 
     private String algorithm;
 
+    private final KeyVaultClient keyVaultClient;
+
     /**
      * Builder for key vault private key
      * @param algorithm algorithm
      * @param kid The key id
+     * @param keyVaultClient related keyVaultClient
      */
-    public KeyVaultPrivateKey(String algorithm, String kid) {
+    public KeyVaultPrivateKey(String algorithm, String kid, KeyVaultClient keyVaultClient) {
         this.algorithm = algorithm;
         this.kid = kid;
+        this.keyVaultClient = keyVaultClient;
+    }
+
+    /**
+     * Builder for key vault private key
+     * @param algorithm algorithm
+     * @param kid The key id
+     */
+    public KeyVaultPrivateKey(String algorithm, String kid) {
+        this(algorithm, kid, null);
+    }
+
+    /**
+     * Get related keyVaultClient, which will be used when signature
+     * @return related keyVaultClient
+     */
+    public KeyVaultClient getKeyVaultClient() {
+        return keyVaultClient;
     }
 
     /**

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

@@ -330,7 +330,7 @@ public Key getKey(String alias, char[] password) {
             LOGGER.exiting("KeyVaultClient", "getKey", null);
             return Optional.ofNullable(certificateBundle)
                            .map(CertificateBundle::getKid)
-                           .map(kid -> new KeyVaultPrivateKey(keyType, kid))
+                           .map(kid -> new KeyVaultPrivateKey(keyType, kid, this))
                            .orElse(null);
         }
         String certificateSecretUri = certificateBundle.getSid();

sdk/keyvault/azure-security-keyvault-jca/src/main/java/com/azure/security/keyvault/jca/implementation/signature/AbstractKeyVaultKeyLessSignature.java

@@ -15,7 +15,6 @@
 import java.security.AlgorithmParameters;
 import java.security.InvalidAlgorithmParameterException;
 import java.security.spec.AlgorithmParameterSpec;
-import static com.azure.security.keyvault.jca.implementation.KeyVaultClient.createKeyVaultClientBySystemProperty;
 
 /**
  * KeyVault Signature to key less sign
@@ -35,14 +34,6 @@ public abstract class AbstractKeyVaultKeyLessSignature extends SignatureSpi {
      */
     public abstract String getAlgorithmName();
 
-    public AbstractKeyVaultKeyLessSignature() {
-        this.keyVaultClient = createKeyVaultClientBySystemProperty();
-    }
-
-    void setKeyVaultClient(KeyVaultClient keyVaultClient) {
-        this.keyVaultClient = keyVaultClient;
-    }
-
     // After throw UnsupportedOperationException, other methods will be called.
     // such as RSAPSSSignature#engineInitVerify.
     @Override
@@ -81,6 +72,7 @@ protected void engineSetParameter(String param, Object value)
     protected void engineInitSign(PrivateKey privateKey, SecureRandom random) {
         if (privateKey instanceof KeyVaultPrivateKey) {
             keyId = ((KeyVaultPrivateKey) privateKey).getKid();
+            keyVaultClient = ((KeyVaultPrivateKey) privateKey).getKeyVaultClient();
         } else {
             throw new UnsupportedOperationException("engineInitSign() not supported which private key is not instance of KeyVaultPrivateKey");
         }

sdk/keyvault/azure-security-keyvault-jca/src/main/java/com/azure/security/keyvault/jca/implementation/signature/KeyVaultKeyLessECSignature.java

@@ -35,8 +35,11 @@ public abstract class KeyVaultKeyLessECSignature extends AbstractKeyVaultKeyLess
     protected byte[] engineSign() {
         byte[] mHash = getDigestValue();
         String encode = Base64.getEncoder().encodeToString(mHash);
-        byte[] encrypted = keyVaultClient.getSignedWithPrivateKey(keyVaultDigestName, encode, keyId);
-        return encodeByte(encrypted);
+        if (keyVaultClient != null) {
+            byte[] encrypted = keyVaultClient.getSignedWithPrivateKey(keyVaultDigestName, encode, keyId);
+            return encodeByte(encrypted);
+        }
+        return new byte[0];
     }
 
 }

sdk/keyvault/azure-security-keyvault-jca/src/main/java/com/azure/security/keyvault/jca/implementation/signature/KeyVaultKeyLessRsaSignature.java

@@ -28,7 +28,10 @@ protected byte[] engineSign() {
         byte[] mHash = getDigestValue();
         String encode = Base64.getEncoder().encodeToString(mHash);
         //For all RSA type certificate in keyVault, we can use PS256 to encrypt.
-        return keyVaultClient.getSignedWithPrivateKey("PS256", encode, keyId);
+        if (keyVaultClient != null) {
+            return keyVaultClient.getSignedWithPrivateKey("PS256", encode, keyId);
+        }
+        return new byte[0];
     }
 
     @Override

sdk/keyvault/azure-security-keyvault-jca/src/test/java/com/azure/security/keyvault/jca/implementation/signature/KeyVaultKeyLessECSignatureTest.java

@@ -4,6 +4,7 @@
 package com.azure.security.keyvault.jca.implementation.signature;
 
 import com.azure.security.keyvault.jca.KeyVaultEncode;
+import com.azure.security.keyvault.jca.KeyVaultPrivateKey;
 import com.azure.security.keyvault.jca.implementation.KeyVaultClient;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.BeforeEach;
@@ -24,6 +25,8 @@ public class KeyVaultKeyLessECSignatureTest {
 
     private final KeyVaultClient keyVaultClient = mock(KeyVaultClient.class);
 
+    private final KeyVaultPrivateKey keyVaultPrivateKey = mock(KeyVaultPrivateKey.class);
+
     private final byte[] signedWithES256 = "fake256Value".getBytes();
     private final byte[] signedWithES384 = "fake384Value".getBytes();
 
@@ -98,11 +101,12 @@ public void engineSetParameterTest() {
     public void setDigestNameAndEngineSignTest() {
         keyVaultKeyLessECSignature = new KeyVaultKeyLessEcSha256Signature();
         when(keyVaultClient.getSignedWithPrivateKey(ArgumentMatchers.eq("ES256"), anyString(), ArgumentMatchers.eq(null))).thenReturn(signedWithES256);
-        keyVaultKeyLessECSignature.setKeyVaultClient(keyVaultClient);
+        when(keyVaultPrivateKey.getKeyVaultClient()).thenReturn(keyVaultClient);
+        keyVaultKeyLessECSignature.engineInitSign(keyVaultPrivateKey, null);
         Assertions.assertArrayEquals(KeyVaultEncode.encodeByte(signedWithES256), keyVaultKeyLessECSignature.engineSign());
 
         keyVaultKeyLessECSignature = new KeyVaultKeyLessEcSha384Signature();
-        keyVaultKeyLessECSignature.setKeyVaultClient(keyVaultClient);
+        keyVaultKeyLessECSignature.engineInitSign(keyVaultPrivateKey, null);
         when(keyVaultClient.getSignedWithPrivateKey(ArgumentMatchers.eq("ES384"), anyString(), ArgumentMatchers.eq(null))).thenReturn(signedWithES384);
         assertArrayEquals(KeyVaultEncode.encodeByte(signedWithES384), keyVaultKeyLessECSignature.engineSign());
     }

sdk/keyvault/azure-security-keyvault-jca/src/test/java/com/azure/security/keyvault/jca/implementation/signature/KeyVaultKeyLessRsaSignatureTest.java

@@ -3,6 +3,7 @@
 
 package com.azure.security.keyvault.jca.implementation.signature;
 
+import com.azure.security.keyvault.jca.KeyVaultPrivateKey;
 import com.azure.security.keyvault.jca.implementation.KeyVaultClient;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -29,6 +30,8 @@ public class KeyVaultKeyLessRsaSignatureTest {
 
     private final KeyVaultClient keyVaultClient = mock(KeyVaultClient.class);
 
+    private final KeyVaultPrivateKey keyVaultPrivateKey = mock(KeyVaultPrivateKey.class);
+
     @BeforeEach
     public void before() {
         System.setProperty("azure.keyvault.uri", KEY_VAULT_TEST_URI_GLOBAL);
@@ -97,7 +100,8 @@ public void engineSetParameterTest() {
     @Test
     public void setDigestNameAndEngineSignTest() throws InvalidAlgorithmParameterException {
         keyVaultKeyLessRsaSignature = new KeyVaultKeyLessRsaSignature();
-        keyVaultKeyLessRsaSignature.setKeyVaultClient(keyVaultClient);
+        when(keyVaultPrivateKey.getKeyVaultClient()).thenReturn(keyVaultClient);
+        keyVaultKeyLessRsaSignature.engineInitSign(keyVaultPrivateKey, null);
         keyVaultKeyLessRsaSignature.engineSetParameter(new PSSParameterSpec("SHA-1", "MGF1", MGF1ParameterSpec.SHA1, 20, 1));
         when(keyVaultClient.getSignedWithPrivateKey(ArgumentMatchers.eq("PS256"), anyString(), ArgumentMatchers.eq(null))).thenReturn("fakeValue".getBytes());
         assertArrayEquals("fakeValue".getBytes(), keyVaultKeyLessRsaSignature.engineSign());