Let KeyResolver cache key

Resolves Azure#25254
heaths · Nov 11, 2021 · 88fb971 · 88fb971
1 parent bcb4cf0
commit 88fb971
Show file tree

Hide file tree

Showing 9 changed files with 242 additions and 159 deletions.
diff --git a/sdk/keyvault/Azure.Security.KeyVault.Keys/src/Cryptography/CryptographyClient.cs b/sdk/keyvault/Azure.Security.KeyVault.Keys/src/Cryptography/CryptographyClient.cs
@@ -160,7 +160,8 @@ internal CryptographyClient(Uri keyId, KeyVaultPipeline pipeline)
 
             _pipeline = pipeline;
             _remoteProvider = remoteClient;
-            _provider = remoteClient;
+
+            // Allow _provider to be set upon first request to the service depending on whether the key can be retrieved.
         }
 
         internal ICryptographyProvider RemoteClient => _remoteProvider;

diff --git a/...eyvault/Azure.Security.KeyVault.Keys/src/Cryptography/LocalCryptographyProviderFactory.cs b/...eyvault/Azure.Security.KeyVault.Keys/src/Cryptography/LocalCryptographyProviderFactory.cs
@@ -21,7 +21,7 @@ public static ICryptographyProvider Create(JsonWebKey keyMaterial, KeyProperties
                     return new EcCryptographyProvider(keyMaterial, keyProperties, localOnly);
                 }
 
-                if (keyMaterial.KeyType == KeyType.Oct || keyMaterial.KeyType == KeyType.OctHsm)
+                if ((keyMaterial.KeyType == KeyType.Oct || keyMaterial.KeyType == KeyType.OctHsm) && keyMaterial.HasPrivateKey)
                 {
                     return new AesCryptographyProvider(keyMaterial, keyProperties, localOnly);
                 }

diff --git a/sdk/keyvault/Azure.Security.KeyVault.Keys/src/KeysEventSource.cs b/sdk/keyvault/Azure.Security.KeyVault.Keys/src/KeysEventSource.cs
@@ -40,7 +40,12 @@ public void KeyTypeNotSupported(string operation, KeyVaultKey key)
         {
             if (IsEnabled())
             {
-                string keyType = key?.KeyType.ToString() ?? "(null)";
+                string keyType = "(null)";
+                if (key != null)
+                {
+                    keyType = key.KeyType.ToString();
+                }
+
                 KeyTypeNotSupported(operation, keyType);
             }
         }

diff --git a/sdk/keyvault/Azure.Security.KeyVault.Keys/tests/KeyClientTests.cs b/sdk/keyvault/Azure.Security.KeyVault.Keys/tests/KeyClientTests.cs
@@ -237,10 +237,13 @@ public void GetCryptographyClientValidation()
         [Test]
         public async Task GetCryptographyClientUsesSamePipeline()
         {
+            const string keyContent = @"{""attributes"":{""created"":1626299777,""enabled"":true,""exportable"":false,""updated"":1626299777},""key"":{""key_ops"":[""wrapKey"",""unwrapKey""],""kid"":""https://test.managedhsm.azure.net/keys/test/abcd1234"",""kty"":""oct-HSM""}}";
+
             // Make sure the created CryptographyClient uses the same mock transport as the KeyVault that created it.
             MockTransport transport = new(new[]
             {
-                new MockResponse(200).WithContent(@"{""attributes"":{""created"":1626299777,""enabled"":true,""exportable"":false,""updated"":1626299777},""key"":{""key_ops"":[""wrapKey"",""unwrapKey""],""kid"":""https://test.managedhsm.azure.net/keys/test/abcd1234"",""kty"":""oct-HSM""}}"),
+                new MockResponse(200).WithContent(keyContent), // Key returned after call to create the key.
+                new MockResponse(200).WithContent(keyContent), // Key returned in attempt to cache the key.
                 new MockResponse(200).WithContent(@"{""alg"":""A128KW"",""kid"":""https://test.managedhsm.azure.net/keys/test/abcd1234"",""value"":""dGVzdA""}"),
             });
 

diff --git a/sdk/keyvault/Azure.Security.KeyVault.Keys/tests/LocalCryptographyProviderFactoryTests.cs b/sdk/keyvault/Azure.Security.KeyVault.Keys/tests/LocalCryptographyProviderFactoryTests.cs
@@ -47,6 +47,18 @@ public void NoKeyMaterial()
             Assert.IsNull(provider);
         }
 
+        [Test]
+        public void NoOctKeyMaterial()
+        {
+            JsonWebKey jwk = new(new[] { KeyOperation.WrapKey, KeyOperation.UnwrapKey })
+            {
+                KeyType = KeyType.OctHsm,
+            };
+
+            ICryptographyProvider provider = LocalCryptographyProviderFactory.Create(jwk, null);
+            Assert.IsNull(provider);
+        }
+
         private static IEnumerable<object[]> GetCreateData()
         {
             Aes aes = Aes.Create();

diff --git a/...lt.Keys/tests/SessionRecords/CryptographyClientLiveTests/EncryptDecryptFromKeyClient.json b/...lt.Keys/tests/SessionRecords/CryptographyClientLiveTests/EncryptDecryptFromKeyClient.json