[Client encryption] Add synchronous initialization to CosmosDataEncryptionKeyProvider

Microsoft.Azure.Cosmos.Encryption.Custom/src/CosmosDataEncryptionKeyProvider.cs

@@ -10,7 +10,14 @@ namespace Microsoft.Azure.Cosmos.Encryption.Custom
     using Microsoft.Data.Encryption.Cryptography;
 
     /// <summary>
-    /// Default implementation for a provider to get a data encryption key - wrapped keys are stored in a Cosmos DB container.
+    /// Default implementation for a provider to get a data encryption key (DEK). Wrapped DEKs are stored as items in a Cosmos DB container.
+    /// Container requirements:
+    ///  - Partition key path must be <c>/id</c> (DEK <c>id</c> used as partition key).
+    ///  - Use a dedicated container to isolate access control and throughput.
+    ///  - Disable TTL to avoid accidental key deletion.
+    /// Usage pattern: construct <see cref="CosmosDataEncryptionKeyProvider"/>, call <see cref="InitializeAsync(Database,string,CancellationToken)"/> or <see cref="Initialize(Container)"/> once at startup, then use <see cref="DataEncryptionKeyContainer"/> for DEK operations.
+    /// Concurrency: initialization is single-assignment (uses <see cref="System.Threading.Interlocked"/>) so concurrent calls after success throw <see cref="InvalidOperationException"/>. Fetch/create operations are safe concurrently.
+    /// Resilience: if the container is deleted or unavailable after initialization, operations surface the underlying exception (for example NotFound). Automatic re-creation is not attempted.
     /// See https://aka.ms/CosmosClientEncryption for more information on client-side encryption support in Azure Cosmos DB.
     /// </summary>
     public sealed class CosmosDataEncryptionKeyProvider : DataEncryptionKeyProvider
@@ -132,18 +139,19 @@ public CosmosDataEncryptionKeyProvider(
         }
 
         /// <summary>
-        /// Initialize Cosmos DB container for CosmosDataEncryptionKeyProvider to store wrapped DEKs
+        /// Ensures the Cosmos DB container for storing wrapped DEKs exists (creating it if needed) and initializes this provider with that container.
+        /// This must be invoked exactly once before any fetch/create operations.
         /// </summary>
-        /// <param name="database">Database</param>
-        /// <param name="containerId">Container id</param>
-        /// <param name="cancellationToken">Cancellation token</param>
-        /// <returns>A task to await on.</returns>
+        /// <param name="database">The Cosmos DB <see cref="Database"/> in which the DEK container should exist.</param>
+        /// <param name="containerId">The identifier of the DEK container. If the container does not exist it will be created with partition key path <c>/id</c>.</param>
+        /// <param name="cancellationToken">An optional cancellation token.</param>
+        /// <returns>A task representing the asynchronous initialization.</returns>
         public async Task InitializeAsync(
             Database database,
             string containerId,
             CancellationToken cancellationToken = default)
         {
-            if (this.container != null)
+            if (Volatile.Read(ref this.container) != null)
             {
                 throw new InvalidOperationException($"{nameof(CosmosDataEncryptionKeyProvider)} has already been initialized.");
             }
@@ -163,7 +171,35 @@ public async Task InitializeAsync(
                     nameof(containerId));
             }
 
-            this.container = containerResponse.Container;
+            this.SetContainer(containerResponse.Container);
+        }
+
+        /// <summary>
+        /// Initializes the provider with an already created Cosmos DB container that meets the required partition key definition (<c>/id</c>).
+        /// </summary>
+        /// <param name="container">Existing Cosmos DB container containing wrapped DEKs or ready to store them.</param>
+        public void Initialize(Container container)
+        {
+            if (container == null)
+            {
+                throw new ArgumentNullException(nameof(container));
+            }
+
+            this.SetContainer(container);
+        }
+
+        /// <summary>
+        /// Sets the backing Cosmos <see cref="Container"/> exactly once.
+        /// Throws if already initialized to prevent accidental reassignment.
+        /// </summary>
+        /// <param name="container">The container to associate with this provider.</param>
+        private void SetContainer(Container container)
+        {
+            Container previous = Interlocked.CompareExchange(ref this.container, container, null);
+            if (previous != null)
+            {
+                throw new InvalidOperationException($"{nameof(CosmosDataEncryptionKeyProvider)} has already been initialized.");
+            }
         }
 
         /// <inheritdoc/>

Microsoft.Azure.Cosmos.Encryption.Custom/tests/Microsoft.Azure.Cosmos.Encryption.Custom.Tests/Contracts/DotNetSDKEncryptionCustomAPI.net6.json

@@ -76,6 +76,11 @@
           "Type": "Constructor",
           "Attributes": [],
           "MethodInfo": "Void .ctor(Microsoft.Data.Encryption.Cryptography.EncryptionKeyStoreProvider, System.Nullable`1[System.TimeSpan])"
+        },
+        "Void Initialize(Microsoft.Azure.Cosmos.Container)": {
+          "Type": "Method",
+          "Attributes": [],
+          "MethodInfo": "Void Initialize(Microsoft.Azure.Cosmos.Container);IsAbstract:False;IsStatic:False;IsVirtual:False;IsGenericMethod:False;IsConstructor:False;IsFinal:False;"
         }
       },
       "NestedTypes": {}
@@ -457,6 +462,11 @@
               "Type": "Constructor",
               "Attributes": [],
               "MethodInfo": "Void .ctor(Microsoft.Data.Encryption.Cryptography.EncryptionKeyStoreProvider, System.Nullable`1[System.TimeSpan])"
+            },
+            "Void Initialize(Microsoft.Azure.Cosmos.Container)": {
+              "Type": "Method",
+              "Attributes": [],
+              "MethodInfo": "Void Initialize(Microsoft.Azure.Cosmos.Container);IsAbstract:False;IsStatic:False;IsVirtual:False;IsGenericMethod:False;IsConstructor:False;IsFinal:False;"
             }
           },
           "NestedTypes": {}

Microsoft.Azure.Cosmos.Encryption.Custom/tests/Microsoft.Azure.Cosmos.Encryption.Custom.Tests/Contracts/DotNetSDKEncryptionCustomAPI.net8.json

@@ -76,6 +76,11 @@
           "Type": "Constructor",
           "Attributes": [],
           "MethodInfo": "Void .ctor(Microsoft.Data.Encryption.Cryptography.EncryptionKeyStoreProvider, System.Nullable`1[System.TimeSpan])"
+        },
+        "Void Initialize(Microsoft.Azure.Cosmos.Container)": {
+          "Type": "Method",
+          "Attributes": [],
+          "MethodInfo": "Void Initialize(Microsoft.Azure.Cosmos.Container);IsAbstract:False;IsStatic:False;IsVirtual:False;IsGenericMethod:False;IsConstructor:False;IsFinal:False;"
         }
       },
       "NestedTypes": {}
@@ -457,6 +462,11 @@
               "Type": "Constructor",
               "Attributes": [],
               "MethodInfo": "Void .ctor(Microsoft.Data.Encryption.Cryptography.EncryptionKeyStoreProvider, System.Nullable`1[System.TimeSpan])"
+            },
+            "Void Initialize(Microsoft.Azure.Cosmos.Container)": {
+              "Type": "Method",
+              "Attributes": [],
+              "MethodInfo": "Void Initialize(Microsoft.Azure.Cosmos.Container);IsAbstract:False;IsStatic:False;IsVirtual:False;IsGenericMethod:False;IsConstructor:False;IsFinal:False;"
             }
           },
           "NestedTypes": {}

Microsoft.Azure.Cosmos.Encryption.Custom/tests/Microsoft.Azure.Cosmos.Encryption.Custom.Tests/CosmosDataEncryptionKeyProviderTests.cs

@@ -0,0 +1,180 @@
+//------------------------------------------------------------
+// Copyright (c) Microsoft Corporation.  All rights reserved.
+//------------------------------------------------------------
+
+namespace Microsoft.Azure.Cosmos.Encryption.Tests
+{
+    using System;
+    using System.Threading;
+    using System.Threading.Tasks;
+    using Microsoft.Azure.Cosmos;
+    using Microsoft.Azure.Cosmos.Encryption.Custom;
+    using Microsoft.VisualStudio.TestTools.UnitTesting;
+    using Moq;
+    using Microsoft.Data.Encryption.Cryptography;
+
+    [TestClass]
+    public class CosmosDataEncryptionKeyProviderTests
+    {
+        private const string ContainerId = "dekContainer";
+
+        [TestMethod]
+        public async Task InitializeAsync_WithValidContainer_CreatesAndSetsContainer()
+        {
+            Mock<Container> mockContainer = new(MockBehavior.Strict);
+            Mock<ContainerResponse> mockContainerResponse = new(MockBehavior.Strict);
+            mockContainerResponse.Setup(r => r.Container).Returns(mockContainer.Object);
+            mockContainerResponse.Setup(r => r.Resource).Returns(new ContainerProperties(ContainerId, partitionKeyPath: "/id"));
+
+            Mock<Database> mockDatabase = new(MockBehavior.Strict);
+            mockDatabase
+                .Setup(db => db.CreateContainerIfNotExistsAsync(
+                    It.Is<string>(s => s == ContainerId),
+                    It.Is<string>(pk => pk == "/id"),
+                    It.IsAny<int?>(),
+                    It.IsAny<RequestOptions>(),
+                    It.IsAny<CancellationToken>()))
+                .ReturnsAsync(mockContainerResponse.Object);
+
+            CosmosDataEncryptionKeyProvider provider = CreateProvider();
+
+            await provider.InitializeAsync(mockDatabase.Object, ContainerId);
+
+            Assert.AreSame(mockContainer.Object, provider.Container);
+
+            mockDatabase.VerifyAll();
+            mockContainerResponse.VerifyAll();
+        }
+
+        [TestMethod]
+        public async Task InitializeAsync_WithWrongPartitionKey_Throws()
+        {
+            Mock<Container> mockContainer = new(MockBehavior.Strict);
+            Mock<ContainerResponse> mockContainerResponse = new(MockBehavior.Strict);
+            mockContainerResponse.Setup(r => r.Container).Returns(mockContainer.Object);
+            mockContainerResponse.Setup(r => r.Resource).Returns(new ContainerProperties("dekBad", partitionKeyPath: "/different-id"));
+
+            Mock<Database> mockDatabase = new(MockBehavior.Strict);
+            mockDatabase
+                .Setup(db => db.CreateContainerIfNotExistsAsync(
+                    It.Is<string>(s => s == "dekBad"),
+                    It.Is<string>(pk => pk == "/id"),
+                    It.IsAny<int?>(),
+                    It.IsAny<RequestOptions>(),
+                    It.IsAny<CancellationToken>()))
+                .ReturnsAsync(mockContainerResponse.Object);
+
+            CosmosDataEncryptionKeyProvider provider = CreateProvider();
+
+            ArgumentException ex = await Assert.ThrowsExceptionAsync<ArgumentException>(() => provider.InitializeAsync(mockDatabase.Object, "dekBad"));
+
+            StringAssert.Contains(ex.Message, "partition key definition");
+            Assert.AreEqual("containerId", ex.ParamName);
+
+            mockDatabase.VerifyAll();
+        }
+
+        [TestMethod]
+        public void Initialize_WithContainer_Succeeds()
+        {
+            Mock<Container> mockContainer = new(MockBehavior.Strict);
+
+            CosmosDataEncryptionKeyProvider provider = CreateProvider();
+            provider.Initialize(mockContainer.Object);
+
+            Assert.AreSame(mockContainer.Object, provider.Container);
+        }
+
+        [TestMethod]
+        public void Initialize_WithNullContainer_Throws()
+        {
+            CosmosDataEncryptionKeyProvider provider = CreateProvider();
+
+            ArgumentNullException ex = Assert.ThrowsException<ArgumentNullException>(() => provider.Initialize(null));
+
+            Assert.AreEqual("container", ex.ParamName);
+        }
+
+        [TestMethod]
+        public void Initialize_Twice_Throws()
+        {
+            Mock<Container> mockContainer = new(MockBehavior.Strict);
+            CosmosDataEncryptionKeyProvider provider = CreateProvider();
+            provider.Initialize(mockContainer.Object);
+
+            InvalidOperationException ex = Assert.ThrowsException<InvalidOperationException>(() => provider.Initialize(mockContainer.Object));
+
+            StringAssert.Contains(ex.Message, nameof(CosmosDataEncryptionKeyProvider));
+        }
+
+        [TestMethod]
+        public async Task InitializeAsync_AfterInitializeContainer_Throws()
+        {
+            Mock<Container> mockContainer = new(MockBehavior.Strict);
+            Mock<Database> mockDatabase = new(MockBehavior.Strict);
+            CosmosDataEncryptionKeyProvider provider = CreateProvider();
+
+            provider.Initialize(mockContainer.Object);
+
+            await Assert.ThrowsExceptionAsync<InvalidOperationException>(() => provider.InitializeAsync(mockDatabase.Object, "ignored"));
+        }
+
+        [TestMethod]
+        public void AccessContainer_BeforeInitialization_Throws()
+        {
+            CosmosDataEncryptionKeyProvider provider = CreateProvider();
+
+            InvalidOperationException ex = Assert.ThrowsException<InvalidOperationException>(() => _ = provider.Container);
+
+            StringAssert.Contains(ex.Message, nameof(CosmosDataEncryptionKeyProvider));
+        }
+
+        [TestMethod]
+        public void Constructor_EncryptionKeyWrapProvider_SetsProperties()
+        {
+#pragma warning disable CS0618
+            Mock<EncryptionKeyWrapProvider> wrapProviderMock = new(MockBehavior.Strict);
+
+            CosmosDataEncryptionKeyProvider provider = new(wrapProviderMock.Object);
+
+            Assert.AreSame(wrapProviderMock.Object, provider.EncryptionKeyWrapProvider);
+            Assert.IsNotNull(provider.DataEncryptionKeyContainer);
+            Assert.IsNotNull(provider.DekCache);
+#pragma warning restore CS0618
+        }
+
+        [TestMethod]
+        public void Constructor_EncryptionKeyStoreProvider_SetsMdePropertiesAndTtl_DefaultInfinite()
+        {
+            TestEncryptionKeyStoreProvider keyStoreProvider = new();
+
+            CosmosDataEncryptionKeyProvider provider = new(keyStoreProvider);
+
+            Assert.AreSame(keyStoreProvider, provider.EncryptionKeyStoreProvider);
+            Assert.IsNotNull(provider.DekCache);
+            Assert.IsNotNull(provider.DataEncryptionKeyContainer);
+            Assert.IsTrue(provider.PdekCacheTimeToLive.HasValue);
+            Assert.IsTrue(provider.PdekCacheTimeToLive.Value > TimeSpan.Zero);
+        }
+
+        [TestMethod]
+        public void Constructor_KeyStoreProvider_SetsMdePropertiesAndTtl_Custom()
+        {
+            TimeSpan ttl = TimeSpan.FromMinutes(15);
+            TestEncryptionKeyStoreProvider keyStoreProvider = new()
+            {
+                DataEncryptionKeyCacheTimeToLive = ttl
+            };
+
+            CosmosDataEncryptionKeyProvider provider = new(keyStoreProvider);
+
+            Assert.AreSame(keyStoreProvider, provider.EncryptionKeyStoreProvider);
+            Assert.AreEqual(ttl, provider.PdekCacheTimeToLive);
+        }
+
+        private static CosmosDataEncryptionKeyProvider CreateProvider()
+        {
+            return new CosmosDataEncryptionKeyProvider(new TestEncryptionKeyStoreProvider());
+        }
+    }
+}