resolve comments

Author: Aditya Abhishek

src/VirtualClient/VirtualClient.Core.UnitTests/KeyVaultManagerTests.cs

@@ -11,6 +11,8 @@
 using Polly;
 using System;
 using System.Net;
+using System.Security.Cryptography;
+using System.Security.Cryptography.X509Certificates;
 using System.Threading;
 using System.Threading.Tasks;
 using VirtualClient.Contracts;
@@ -58,18 +60,24 @@ public void SetupDefaultBehaviors()
                 .ReturnsAsync(Response.FromValue(key, Mock.Of<Response>()));
 
             // Mock the certificates
+
             this.certificateClientMock = new Mock<CertificateClient>(MockBehavior.Strict, new Uri("https://myvault.vault.azure.net/"), new MockTokenCredential());
-            var certificate = CertificateModelFactory.KeyVaultCertificateWithPolicy(
+            var publicKeyCertificate = CertificateModelFactory.KeyVaultCertificateWithPolicy(
                 properties: CertificateModelFactory.CertificateProperties(
                     id: new Uri($"https://myvault.vault.azure.net/certificates/mycert/v3"),
                     name: "mycert",
                     version: "v3",
                     vaultUri: new Uri("https://myvault.vault.azure.net/")),
-                policy: CertificateModelFactory.CertificatePolicy(subject: "CN=mycert"));
+                policy: CertificateModelFactory.CertificatePolicy(subject: "CN=mycert"),
+                cer: this.GenerateTestCertificateBytes());
 
             this.certificateClientMock
                 .Setup(c => c.GetCertificateAsync("mycert", It.IsAny<CancellationToken>()))
-                .ReturnsAsync(Response.FromValue(certificate, Mock.Of<Response>()));
+                .ReturnsAsync(Response.FromValue(publicKeyCertificate, Mock.Of<Response>()));
+
+            this.certificateClientMock
+                .Setup(c => c.DownloadCertificateAsync("mycert", It.IsAny<string>(), It.IsAny<CancellationToken>()))
+                .ReturnsAsync(Response.FromValue(this.GenerateTestCertificateWithPrivateKey(), Mock.Of<Response>()));
 
             // Initialize the KeyVaultManager with the mocked clients
             this.keyVaultManager = new TestKeyVaultManager(
@@ -89,36 +97,38 @@ public void KeyVaultManagerConstructorsValidateRequiredParameters()
         }
 
         [Test]
-        public async Task KeyVaultManagerReturnsExpectedSecretDescriptor()
+        public async Task KeyVaultManagerReturnsExpectedSecretValue()
         {
             var result = await this.keyVaultManager.GetSecretAsync(this.mockDescriptor, CancellationToken.None, Policy.NoOpAsync());
             Assert.IsNotNull(result);
-            Assert.AreEqual("mysecret", result.Name);
-            Assert.AreEqual("secret-value", result.Value);
-            Assert.AreEqual("v1", result.Version);
-            Assert.AreEqual(KeyVaultObjectType.Secret, result.ObjectType);
+            Assert.AreEqual("secret-value", result);
         }
 
         [Test]
-        public async Task KeyVaultManagerReturnsExpectedKeyDescriptor()
+        public async Task KeyVaultManagerReturnsExpectedKey()
         {
             this.mockDescriptor.Name = "mykey";
             var result = await this.keyVaultManager.GetKeyAsync(this.mockDescriptor, CancellationToken.None, Policy.NoOpAsync());
             Assert.IsNotNull(result);
             Assert.AreEqual("mykey", result.Name);
-            Assert.AreEqual("v2", result.Version);
-            Assert.AreEqual(KeyVaultObjectType.Key, result.ObjectType);
         }
 
         [Test]
-        public async Task KeyVaultManagerReturnsExpectedCertificateDescriptor()
+        [TestCase(true)]
+        [TestCase(false)]
+        public async Task KeyVaultManagerReturnsExpectedCertificate(bool retrieveWithPrivateKey)
         {
             this.mockDescriptor.Name = "mycert";
-            var result = await this.keyVaultManager.GetCertificateAsync(this.mockDescriptor, CancellationToken.None, Policy.NoOpAsync());
+            var result = await this.keyVaultManager.GetCertificateAsync(this.mockDescriptor, CancellationToken.None, retrieveWithPrivateKey);
             Assert.IsNotNull(result);
-            Assert.AreEqual("mycert", result.Name);
-            Assert.AreEqual("v3", result.Version);
-            Assert.AreEqual(KeyVaultObjectType.Certificate, result.ObjectType);
+            if (retrieveWithPrivateKey)
+            {
+                Assert.IsTrue(result.HasPrivateKey);
+            }
+            else
+            {
+                Assert.IsFalse(result.HasPrivateKey);
+            }
         }
 
         [Test]
@@ -186,6 +196,42 @@ public void KeyVaultManagerAppliesRetryPolicyOnTransientErrors()
             Assert.AreEqual(3, attempts);
         }
 
+        // Create a dummy, self-signed certificate for testing
+        private byte[] GenerateTestCertificateBytes()
+        {
+            var distinguishedName = new X500DistinguishedName("CN=TestCert");
+
+            using var rsa = RSA.Create(2048);
+            var request = new CertificateRequest(distinguishedName, rsa, HashAlgorithmName.SHA256, RSASignaturePadding.Pkcs1);
+
+            var certificate = request.CreateSelfSigned(
+                DateTimeOffset.UtcNow.AddDays(-1),
+                DateTimeOffset.UtcNow.AddYears(1));
+
+            // Export to DER format (byte[]) – matches cert.Cer in Azure Key Vault
+            return certificate.Export(X509ContentType.Cert);
+        }
+
+        private X509Certificate2 GenerateTestCertificateWithPrivateKey()
+        {
+            var distinguishedName = new X500DistinguishedName("CN=TestWithPrivateKey");
+
+            using var rsa = RSA.Create(2048);
+            var request = new CertificateRequest(distinguishedName, rsa, HashAlgorithmName.SHA256, RSASignaturePadding.Pkcs1);
+
+            var certificate = request.CreateSelfSigned(
+                DateTimeOffset.UtcNow.AddDays(-1),
+                DateTimeOffset.UtcNow.AddYears(1));
+
+            // Export and import to ensure HasPrivateKey is true in unit tests
+            var bytes = certificate.Export(X509ContentType.Pfx, "");
+#pragma warning disable SYSLIB0057
+            // using this obsolete method just for Unit Testing
+            var cert = new X509Certificate2(bytes, "", X509KeyStorageFlags.Exportable);
+#pragma warning restore SYSLIB0057
+            return cert;
+        }
+
         private class TestKeyVaultManager : KeyVaultManager
         {
             private readonly SecretClient secretClient;

src/VirtualClient/VirtualClient.Core/IKeyVaultManager.cs

@@ -3,8 +3,10 @@
 
 namespace VirtualClient
 {
+    using System.Security.Cryptography.X509Certificates;
     using System.Threading;
     using System.Threading.Tasks;
+    using Azure.Security.KeyVault.Keys;
     using Polly;
     using VirtualClient.Contracts;
 
@@ -24,8 +26,8 @@ public interface IKeyVaultManager
         /// <param name="descriptor"> Provides the details for the secret to retrieve (requires "SecretName" and "VaultUri"). </param>
         /// <param name="cancellationToken"> A token that can be used to cancel the operation. </param>
         /// <param name="retryPolicy"> A policy to use for handling retries when transient errors/failures happen. </param>
-        /// <returns> A <see cref="KeyVaultDescriptor"/> containing the secret value and metadata. </returns>
-        Task<KeyVaultDescriptor> GetSecretAsync(
+        /// <returns> A <see cref="string"/> containing the secret value. </returns>
+        Task<string> GetSecretAsync(
             KeyVaultDescriptor descriptor,
             CancellationToken cancellationToken,
             IAsyncPolicy retryPolicy = null);
@@ -36,8 +38,8 @@ Task<KeyVaultDescriptor> GetSecretAsync(
         /// <param name="descriptor"> Provides the details for the key to retrieve (requires "KeyName" and "VaultUri"). </param>
         /// <param name="cancellationToken"> A token that can be used to cancel the operation. </param>
         /// <param name="retryPolicy"> A policy to use for handling retries when transient errors/failures happen. </param>
-        /// <returns> A <see cref="KeyVaultDescriptor"/> containing the key metadata. </returns>
-        Task<KeyVaultDescriptor> GetKeyAsync(
+        /// <returns> A <see cref="KeyVaultKey"/> containing the key properties. </returns>
+        Task<KeyVaultKey> GetKeyAsync(
             KeyVaultDescriptor descriptor,
             CancellationToken cancellationToken,
             IAsyncPolicy retryPolicy = null);
@@ -47,11 +49,13 @@ Task<KeyVaultDescriptor> GetKeyAsync(
         /// </summary>
         /// <param name="descriptor"> Provides the details for the certificate to retrieve (requires "CertificateName" and "VaultUri"). </param>
         /// <param name="cancellationToken"> A token that can be used to cancel the operation. </param>
+        /// <param name="retrieveWithPrivateKey"> Indicates whether the private key should be retrieved along with the certificate. Default is false </param>
         /// <param name="retryPolicy"> A policy to use for handling retries when transient errors/failures happen. </param>
-        /// <returns> A <see cref="KeyVaultDescriptor"/> containing the certificate metadata. </returns>
-        Task<KeyVaultDescriptor> GetCertificateAsync(
+        /// <returns> A <see cref="X509Certificate2"/> containing the requested certificate. </returns>
+        Task<X509Certificate2> GetCertificateAsync(
             KeyVaultDescriptor descriptor,
             CancellationToken cancellationToken,
+            bool retrieveWithPrivateKey = false,
             IAsyncPolicy retryPolicy = null);
     }
 }

src/VirtualClient/VirtualClient.Core/KeyVaultManager.cs

@@ -5,15 +5,14 @@ namespace VirtualClient
 {
     using System;
     using System.Net;
-    using System.Runtime.ConstrainedExecution;
+    using System.Security.Cryptography.X509Certificates;
     using System.Threading;
     using System.Threading.Tasks;
     using Azure;
     using Azure.Core;
     using Azure.Security.KeyVault.Certificates;
     using Azure.Security.KeyVault.Keys;
     using Azure.Security.KeyVault.Secrets;
-    using Microsoft.CodeAnalysis.CSharp.Syntax;
     using Polly;
     using VirtualClient.Common.Extensions;
     using VirtualClient.Contracts;
@@ -59,12 +58,12 @@ public KeyVaultManager(DependencyKeyVaultStore storeDescription)
         /// <param name="cancellationToken">A token that can be used to cancel the operation.</param>
         /// <param name="retryPolicy">A policy to use for handling retries when transient errors/failures happen.</param>
         /// <returns>
-        /// A <see cref="KeyVaultDescriptor"/> containing the secret value and metadata.
+        /// A <see cref="string"/> containing the secret value.
         /// </returns>
         /// <exception cref="DependencyException">
         /// Thrown if the secret is not found, access is denied, or another error occurs.
         /// </exception>
-        public async Task<KeyVaultDescriptor> GetSecretAsync(
+        public async Task<string> GetSecretAsync(
             KeyVaultDescriptor descriptor,
             CancellationToken cancellationToken,
             IAsyncPolicy retryPolicy = null)
@@ -87,16 +86,7 @@ public async Task<KeyVaultDescriptor> GetSecretAsync(
                 return await (retryPolicy ?? KeyVaultManager.DefaultRetryPolicy).ExecuteAsync(async () =>
                 {
                     KeyVaultSecret secret = await client.GetSecretAsync(secretName, cancellationToken: cancellationToken);
-                    KeyVaultDescriptor result = new KeyVaultDescriptor(descriptor)
-                    {
-                        Value = secret.Value,
-                        Version = secret.Properties.Version,
-                        Name = secretName,
-                        VaultUri = vaultUri.ToString(),
-                        ObjectId = secret.Id?.ToString(),
-                        ObjectType = KeyVaultObjectType.Secret
-                    };
-                    return result;
+                    return secret.Value;
                 }).ConfigureAwait(false);
             }
             catch (RequestFailedException ex) when (ex.Status == (int)HttpStatusCode.Forbidden)
@@ -136,12 +126,12 @@ public async Task<KeyVaultDescriptor> GetSecretAsync(
         /// <param name="cancellationToken">A token that can be used to cancel the operation.</param>
         /// <param name="retryPolicy">A policy to use for handling retries when transient errors/failures happen.</param>
         /// <returns>
-        /// A <see cref="KeyVaultDescriptor"/> containing the key metadata.
+        /// A <see cref="KeyVaultKey"/> containing the key.
         /// </returns>
         /// <exception cref="DependencyException">
         /// Thrown if the key is not found, access is denied, or another error occurs.
         /// </exception>
-        public async Task<KeyVaultDescriptor> GetKeyAsync(
+        public async Task<KeyVaultKey> GetKeyAsync(
             KeyVaultDescriptor descriptor,
             CancellationToken cancellationToken,
             IAsyncPolicy retryPolicy = null)
@@ -162,15 +152,7 @@ public async Task<KeyVaultDescriptor> GetKeyAsync(
                 return await (retryPolicy ?? KeyVaultManager.DefaultRetryPolicy).ExecuteAsync(async () =>
                 {
                     KeyVaultKey key = await client.GetKeyAsync(keyName, cancellationToken: cancellationToken);
-                    KeyVaultDescriptor result = new KeyVaultDescriptor(descriptor)
-                    {
-                        ObjectType = KeyVaultObjectType.Key,
-                        Name = keyName,
-                        VaultUri = vaultUri.ToString(),
-                        Version = key.Properties.Version,
-                        ObjectId = key.Id.ToString()
-                    };
-                    return result;
+                    return key;
                 }).ConfigureAwait(false);
             }
             catch (RequestFailedException ex) when (ex.Status == (int)HttpStatusCode.Forbidden)
@@ -208,16 +190,18 @@ public async Task<KeyVaultDescriptor> GetKeyAsync(
         /// </summary>
         /// <param name="descriptor">Provides the details for the certificate to retrieve (requires "CertificateName" and "VaultUri").</param>
         /// <param name="cancellationToken">A token that can be used to cancel the operation.</param>
+        /// <param name="retrieveWithPrivateKey">flag to decode whether to retrieve certificate with private key</param>
         /// <param name="retryPolicy">A policy to use for handling retries when transient errors/failures happen.</param>
         /// <returns>
-        /// A <see cref="KeyVaultDescriptor"/> containing the certificate metadata.
+        /// A <see cref="X509Certificate2"/> containing the certificate 
         /// </returns>
         /// <exception cref="DependencyException">
         /// Thrown if the certificate is not found, access is denied, or another error occurs.
         /// </exception>
-        public async Task<KeyVaultDescriptor> GetCertificateAsync(
+        public async Task<X509Certificate2> GetCertificateAsync(
             KeyVaultDescriptor descriptor,
             CancellationToken cancellationToken,
+            bool retrieveWithPrivateKey = false,
             IAsyncPolicy retryPolicy = null)
         {
             this.ValidateKeyVaultStore();
@@ -235,17 +219,26 @@ public async Task<KeyVaultDescriptor> GetCertificateAsync(
             {
                 return await (retryPolicy ?? KeyVaultManager.DefaultRetryPolicy).ExecuteAsync(async () =>
                 {
-                    KeyVaultCertificateWithPolicy cert = await client.GetCertificateAsync(certName, cancellationToken: cancellationToken);
-                    KeyVaultDescriptor result = new KeyVaultDescriptor(descriptor)
+                    // Get the full certificate with private key (PFX) if requested
+                    if (retrieveWithPrivateKey)
+                    {
+                        X509Certificate2 privateKeyCert = await client
+                            .DownloadCertificateAsync(certName, cancellationToken: cancellationToken)
+                            .ConfigureAwait(false);
+
+                        if (privateKeyCert is null || !privateKeyCert.HasPrivateKey)
+                        {
+                            throw new DependencyException("Failed to retrieve certificate content with private key.");
+                        }
+
+                        return privateKeyCert;
+                    }
+                    else
                     {
-                        ObjectType = KeyVaultObjectType.Certificate,
-                        Name = certName,
-                        VaultUri = vaultUri.ToString(),
-                        Version = cert.Properties.Version,
-                        ObjectId = cert.Id.ToString(),
-                        Policy = cert.Policy?.ToString()
-                    };
-                    return result;
+                        // If private key not needed, load cert from PublicBytes
+                        KeyVaultCertificateWithPolicy cert = await client.GetCertificateAsync(certName, cancellationToken: cancellationToken);
+                        return X509CertificateLoader.LoadCertificate(cert.Cer);
+                    }
                 }).ConfigureAwait(false);
             }
             catch (RequestFailedException ex) when (ex.Status == (int)HttpStatusCode.Forbidden)