[Client encryption]: Implement EncryptionContainerStream

Microsoft.Azure.Cosmos.Encryption.Custom/src/Common/CosmosJsonDotNetSerializer.cs

@@ -7,6 +7,8 @@ namespace Microsoft.Azure.Cosmos.Encryption.Custom
     using System;
     using System.IO;
     using System.Text;
+    using System.Threading;
+    using System.Threading.Tasks;
     using Newtonsoft.Json;
 
     /// <summary>
@@ -40,6 +42,18 @@ internal CosmosJsonDotNetSerializer(JsonSerializerSettings jsonSerializerSetting
         /// <returns>The object representing the deserialized stream</returns>
         public T FromStream<T>(Stream stream)
         {
+            return this.FromStream<T>(stream, false);
+        }
+
+        /// <summary>
+        /// Convert a Stream to the passed in type.
+        /// </summary>
+        /// <typeparam name="T">The type of object that should be deserialized</typeparam>
+        /// <param name="stream">An open stream that is readable that contains JSON</param>
+        /// <param name="leaveOpen">True if input stream shouldn't be disposed</param>
+        /// <returns>The object representing the deserialized stream</returns>
+        public T FromStream<T>(Stream stream, bool leaveOpen)
+        {
 #if NET8_0_OR_GREATER
             ArgumentNullException.ThrowIfNull(stream);
 #else
@@ -54,7 +68,7 @@ public T FromStream<T>(Stream stream)
                 return (T)(object)stream;
             }
 
-            using (StreamReader sr = new (stream))
+            using (StreamReader sr = new (stream, Encoding.UTF8, true, 1024, leaveOpen))
             using (JsonTextReader jsonTextReader = new (sr))
             {
                 jsonTextReader.ArrayPool = JsonArrayPool.Instance;
@@ -72,19 +86,30 @@ public T FromStream<T>(Stream stream)
         public MemoryStream ToStream<T>(T input)
         {
             MemoryStream streamPayload = new ();
-            using (StreamWriter streamWriter = new (streamPayload, encoding: CosmosJsonDotNetSerializer.DefaultEncoding, bufferSize: 1024, leaveOpen: true))
+#pragma warning disable VSTHRD002 // Avoid problematic synchronous waits
+            this.ToStreamAsync(input, streamPayload, CancellationToken.None).GetAwaiter().GetResult();
+#pragma warning restore VSTHRD002 // Avoid problematic synchronous waits
+            return streamPayload;
+        }
+
+        public async Task ToStreamAsync<T>(T input, Stream output, CancellationToken cancellationToken)
+        {
+            using (StreamWriter streamWriter = new (output, encoding: CosmosJsonDotNetSerializer.DefaultEncoding, bufferSize: 1024, leaveOpen: true))
             using (JsonTextWriter writer = new (streamWriter))
             {
                 writer.ArrayPool = JsonArrayPool.Instance;
                 writer.Formatting = Newtonsoft.Json.Formatting.None;
                 JsonSerializer jsonSerializer = this.GetSerializer();
                 jsonSerializer.Serialize(writer, input);
                 writer.Flush();
-                streamWriter.Flush();
+#if NET8_0_OR_GREATER
+                await streamWriter.FlushAsync(cancellationToken);
+#else
+                await streamWriter.FlushAsync();
+#endif
             }
 
-            streamPayload.Position = 0;
-            return streamPayload;
+            output.Position = 0;
         }
 
         /// <summary>

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

@@ -7,6 +7,7 @@ namespace Microsoft.Azure.Cosmos.Encryption.Custom
     using System;
     using System.Collections.Generic;
     using System.Linq;
+    using System.Text.Json.Serialization;
     using Newtonsoft.Json;
 
     /// <summary>
@@ -73,31 +74,36 @@ internal DataEncryptionKeyProperties(DataEncryptionKeyProperties source)
         /// </para>
         /// </remarks>
         [JsonProperty(PropertyName = "id")]
+        [JsonPropertyName("id")]
         public string Id { get; internal set; }
 
         /// <summary>
         /// Gets the Encryption algorithm that will be used along with this data encryption key to encrypt/decrypt data.
         /// </summary>
         [JsonProperty(PropertyName = "encryptionAlgorithm", NullValueHandling = NullValueHandling.Ignore)]
+        [JsonPropertyName("encryptionAlgorithm")]
         public string EncryptionAlgorithm { get; internal set; }
 
         /// <summary>
         /// Gets wrapped form of the data encryption key.
         /// </summary>
         [JsonProperty(PropertyName = "wrappedDataEncryptionKey", NullValueHandling = NullValueHandling.Ignore)]
+        [JsonPropertyName("wrappedDataEncryptionKey")]
         public byte[] WrappedDataEncryptionKey { get; internal set; }
 
         /// <summary>
         /// Gets metadata for the wrapping provider that can be used to unwrap the wrapped data encryption key.
         /// </summary>
         [JsonProperty(PropertyName = "keyWrapMetadata", NullValueHandling = NullValueHandling.Ignore)]
+        [JsonPropertyName("keyWrapMetadata")]
         public EncryptionKeyWrapMetadata EncryptionKeyWrapMetadata { get; internal set; }
 
         /// <summary>
         /// Gets the creation time of the resource from the Azure Cosmos DB service.
         /// </summary>
-        [JsonConverter(typeof(UnixDateTimeConverter))]
+        [Newtonsoft.Json.JsonConverter(typeof(UnixDateTimeConverter))]
         [JsonProperty(PropertyName = "createTime", NullValueHandling = NullValueHandling.Ignore)]
+        [JsonPropertyName("createTime")]
         public DateTime? CreatedTime { get; internal set; }
 
         /// <summary>
@@ -110,14 +116,16 @@ internal DataEncryptionKeyProperties(DataEncryptionKeyProperties source)
         /// ETags are used for concurrency checking when updating resources.
         /// </remarks>
         [JsonProperty(PropertyName = "_etag", NullValueHandling = NullValueHandling.Ignore)]
+        [JsonPropertyName("_etag")]
         public string ETag { get; internal set; }
 
         /// <summary>
         /// Gets the last modified time stamp associated with the resource from the Azure Cosmos DB service.
         /// </summary>
         /// <value>The last modified time stamp associated with the resource.</value>
-        [JsonConverter(typeof(UnixDateTimeConverter))]
+        [Newtonsoft.Json.JsonConverter(typeof(UnixDateTimeConverter))]
         [JsonProperty(PropertyName = "_ts", NullValueHandling = NullValueHandling.Ignore)]
+        [JsonPropertyName("_ts")]
         public DateTime? LastModified { get; internal set; }
 
         /// <summary>
@@ -129,6 +137,7 @@ internal DataEncryptionKeyProperties(DataEncryptionKeyProperties source)
         /// E.g. a self-link for a document could be dbs/db_resourceid/colls/coll_resourceid/documents/doc_resourceid
         /// </remarks>
         [JsonProperty(PropertyName = "_self", NullValueHandling = NullValueHandling.Ignore)]
+        [JsonPropertyName("_self")]
         public virtual string SelfLink { get; internal set; }
 
         /// <summary>
@@ -143,6 +152,7 @@ internal DataEncryptionKeyProperties(DataEncryptionKeyProperties source)
         /// These resource ids are used when building up SelfLinks, a static addressable Uri for each resource within a database account.
         /// </remarks>
         [JsonProperty(PropertyName = "_rid", NullValueHandling = NullValueHandling.Ignore)]
+        [JsonPropertyName("_rid")]
         internal string ResourceId { get; set; }
 
         /// <summary>

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

@@ -4,6 +4,8 @@
 
 namespace Microsoft.Azure.Cosmos.Encryption.Custom
 {
+    using System;
+    using System.Threading;
     using System.Threading.Tasks;
 
     /// <summary>
@@ -71,13 +73,26 @@ namespace Microsoft.Azure.Cosmos.Encryption.Custom
     /// ]]>
     /// </code>
     /// </example>
-    public abstract class DecryptableItem
+    public abstract class DecryptableItem : IDisposable
     {
         /// <summary>
         /// Decrypts and deserializes the content.
         /// </summary>
         /// <typeparam name="T">The type of item to be returned.</typeparam>
         /// <returns>The requested item and the decryption related context.</returns>
         public abstract Task<(T, DecryptionContext)> GetItemAsync<T>();
+
+        /// <summary>
+        /// Decrypts and deserializes the content.
+        /// </summary>
+        /// <param name="cancellationToken">Cancellation token.</param>
+        /// <typeparam name="T">The type of item to be returned.</typeparam>
+        /// <returns>The requested item and the decryption related context.</returns>
+        public abstract Task<(T, DecryptionContext)> GetItemAsync<T>(CancellationToken cancellationToken);
+
+        /// <summary>
+        /// Dispose unmanaged resources.
+        /// </summary>
+        public abstract void Dispose();
     }
 }

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

@@ -5,6 +5,7 @@
 namespace Microsoft.Azure.Cosmos.Encryption.Custom
 {
     using System;
+    using System.Threading;
     using System.Threading.Tasks;
     using Newtonsoft.Json.Linq;
 
@@ -27,7 +28,16 @@ public DecryptableItemCore(
             this.cosmosSerializer = cosmosSerializer ?? throw new ArgumentNullException(nameof(cosmosSerializer));
         }
 
-        public override async Task<(T, DecryptionContext)> GetItemAsync<T>()
+        public override void Dispose()
+        {
+        }
+
+        public override Task<(T, DecryptionContext)> GetItemAsync<T>()
+        {
+            return this.GetItemAsync<T>(CancellationToken.None);
+        }
+
+        public override async Task<(T, DecryptionContext)> GetItemAsync<T>(CancellationToken cancellationToken)
         {
             if (this.decryptableContent is not JObject document)
             {
@@ -40,7 +50,7 @@ public DecryptableItemCore(
                     document,
                     this.encryptor,
                     new CosmosDiagnosticsContext(),
-                    cancellationToken: default);
+                    cancellationToken: cancellationToken);
 
                 return (this.cosmosSerializer.FromStream<T>(EncryptionProcessor.BaseSerializer.ToStream(decryptedItem)), decryptionContext);
             }

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

@@ -4,13 +4,17 @@
 
 namespace Microsoft.Azure.Cosmos.Encryption.Custom
 {
+    using System;
     using System.IO;
+    using System.Threading;
+    using System.Threading.Tasks;
+    using Microsoft.Azure.Cosmos.Encryption.Custom.Transformation;
     using Newtonsoft.Json.Linq;
 
     /// <summary>
     /// Input type should implement this abstract class for lazy decryption and to retrieve the details in the write path.
     /// </summary>
-    public abstract class EncryptableItem
+    public abstract class EncryptableItem : IDisposable
     {
         /// <summary>
         /// Gets DecryptableItem
@@ -24,6 +28,15 @@ public abstract class EncryptableItem
         /// <returns>Input payload in stream format</returns>
         protected internal abstract Stream ToStream(CosmosSerializer serializer);
 
+        /// <summary>
+        /// Gets the input payload in stream format.
+        /// </summary>
+        /// <param name="serializer">Cosmos Serializer</param>
+        /// <param name="outputStream">Output stream</param>
+        /// <param name="cancellationToken">CancellationToken</param>
+        /// <returns>A <see cref="Task"/> representing the asynchronous operation.</returns>
+        protected internal abstract Task ToStreamAsync(CosmosSerializer serializer, Stream outputStream, CancellationToken cancellationToken);
+
         /// <summary>
         /// Populates the DecryptableItem that can be used getting the decryption result.
         /// </summary>
@@ -34,5 +47,27 @@ protected internal abstract void SetDecryptableItem(
             JToken decryptableContent,
             Encryptor encryptor,
             CosmosSerializer cosmosSerializer);
+
+#if ENCRYPTION_CUSTOM_PREVIEW && NET8_0_OR_GREATER
+        /// <summary>
+        /// Populates the DecryptableItem that can be used getting the decryption result.
+        /// </summary>
+        /// <param name="decryptableStream">The encrypted content stream which is yet to be decrypted.</param>
+        /// <param name="encryptor">Encryptor instance which will be used for decryption.</param>
+        /// <param name="jsonProcessor">Json processor for decryption.</param>
+        /// <param name="cosmosSerializer">Serializer instance which will be used for deserializing the content after decryption.</param>
+        /// <param name="streamManager">Stream manager providing output streams.</param>
+        protected internal abstract void SetDecryptableStream(
+            Stream decryptableStream,
+            Encryptor encryptor,
+            JsonProcessor jsonProcessor,
+            CosmosSerializer cosmosSerializer,
+            StreamManager streamManager);
+#endif
+
+        /// <summary>
+        /// Release unmananaged resources
+        /// </summary>
+        public abstract void Dispose();
     }
 }

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

@@ -6,6 +6,8 @@ namespace Microsoft.Azure.Cosmos.Encryption.Custom
 {
     using System;
     using System.IO;
+    using System.Threading;
+    using System.Threading.Tasks;
     using Newtonsoft.Json.Linq;
 
     /// <summary>
@@ -69,16 +71,45 @@ protected internal override void SetDecryptableItem(
                 cosmosSerializer);
         }
 
+        private void Dispose(bool disposing)
+        {
+            if (disposing)
+            {
+                this.StreamPayload?.Dispose();
+                this.DecryptableItem?.Dispose();
+            }
+        }
+
         /// <inheritdoc/>
-        public void Dispose()
+        public override void Dispose()
         {
-            this.StreamPayload.Dispose();
+            this.Dispose(true);
         }
 
         /// <inheritdoc/>
         protected internal override Stream ToStream(CosmosSerializer serializer)
         {
             return this.StreamPayload;
         }
+
+        /// <inheritdoc/>
+        /// <remarks>This solution is not performant with Newtonsoft.Json.</remarks>
+        protected internal override async Task ToStreamAsync(CosmosSerializer serializer, Stream outputStream, CancellationToken cancellationToken)
+        {
+#if NET8_0_OR_GREATER
+            await this.StreamPayload.CopyToAsync(outputStream, cancellationToken);
+#else
+            await this.StreamPayload.CopyToAsync(outputStream, 81920, cancellationToken);
+#endif
+        }
+
+#if NET8_0_OR_GREATER
+        /// <inheritdoc/>
+        /// <remarks>Direct stream based item is not supported with Newtonsoft.Json.</remarks>
+        protected internal override void SetDecryptableStream(Stream decryptableStream, Encryptor encryptor, JsonProcessor jsonProcessor, CosmosSerializer cosmosSerializer, StreamManager streamManager)
+        {
+            throw new NotImplementedException("Stream based item is only allowed for EncryptionContainerStream");
+        }
+#endif
     }
 }

Microsoft.Azure.Cosmos.Encryption.Custom/src/EncryptableItem{T}.cs

@@ -6,6 +6,8 @@ namespace Microsoft.Azure.Cosmos.Encryption.Custom
 {
     using System;
     using System.IO;
+    using System.Threading;
+    using System.Threading.Tasks;
     using Newtonsoft.Json.Linq;
 
     /// <summary>
@@ -85,5 +87,32 @@ protected internal override Stream ToStream(CosmosSerializer serializer)
         {
             return serializer.ToStream(this.Item);
         }
+
+        /// <inheritdoc/>
+        /// <remarks>This solution is not performant with Newtonsoft.Json.</remarks>
+        protected internal override async Task ToStreamAsync(CosmosSerializer serializer, Stream outputStream, CancellationToken cancellationToken)
+        {
+            Stream temp = serializer.ToStream(this.Item);
+#if NET8_0_OR_GREATER
+            await temp.CopyToAsync(outputStream, cancellationToken);
+#else
+            await temp.CopyToAsync(outputStream, 81920, cancellationToken);
+#endif
+        }
+
+#if NET8_0_OR_GREATER
+        /// <inheritdoc/>
+        /// <remarks>Direct stream based item is not supported with Newtonsoft.Json.</remarks>
+        protected internal override void SetDecryptableStream(Stream decryptableStream, Encryptor encryptor, JsonProcessor jsonProcessor, CosmosSerializer cosmosSerializer, StreamManager streamManager)
+        {
+            throw new NotImplementedException("Stream based item is only allowed for EncryptionContainerStream");
+        }
+#endif
+
+        /// <inheritdoc/>
+        /// <remarks>Does nothing with Newtonsoft based EncryptableItem.</remarks>
+        public override void Dispose()
+        {
+        }
     }
 }