Client Encryption : Adds integration with Cosmos SDK 3.15.0-preview (#…

…1956)

* Integrate with latest Cosmos SDK 3.15.0-preview

* Use customer serializer

Microsoft.Azure.Cosmos.Encryption/src/AssemblyInfo.cs

@@ -0,0 +1,8 @@
+//------------------------------------------------------------
+// Copyright (c) Microsoft Corporation.  All rights reserved.
+//------------------------------------------------------------
+
+using System.Runtime.CompilerServices;
+
+[assembly: InternalsVisibleTo("Microsoft.Azure.Cosmos.Encryption.Tests" + Microsoft.Azure.Cosmos.Encryption.AssemblyKeys.TestPublicKey)]
+[assembly: InternalsVisibleTo("Microsoft.Azure.Cosmos.Encryption.EmulatorTests" + Microsoft.Azure.Cosmos.Encryption.AssemblyKeys.TestPublicKey)]

Microsoft.Azure.Cosmos.Encryption/src/DecryptableFeedResponse.cs

@@ -49,11 +49,7 @@ internal static DecryptableFeedResponse<T> CreateResponse(
 
             using (responseMessage)
             {
-                // ReadFeed can return 304 in some scenarios (for example Change Feed)
-                if (responseMessage.StatusCode != HttpStatusCode.NotModified)
-                {
-                    responseMessage.EnsureSuccessStatusCode();
-                }
+                responseMessage.EnsureSuccessStatusCode();
 
                 return new DecryptableFeedResponse<T>(
                     responseMessage,

Microsoft.Azure.Cosmos.Encryption/src/EncryptionContainer.cs

@@ -800,75 +800,6 @@ public override Task<IReadOnlyList<FeedRange>> GetFeedRangesAsync(
             return this.container.GetFeedRangesAsync(cancellationToken);
         }
 
-        public override FeedIterator GetChangeFeedStreamIterator(
-            string continuationToken = null,
-            ChangeFeedRequestOptions changeFeedRequestOptions = null)
-        {
-            return new EncryptionFeedIterator(
-                this.container.GetChangeFeedStreamIterator(
-                    continuationToken,
-                    changeFeedRequestOptions),
-                this.Encryptor,
-                this.CosmosSerializer);
-        }
-
-        public override FeedIterator GetChangeFeedStreamIterator(
-            FeedRange feedRange,
-            ChangeFeedRequestOptions changeFeedRequestOptions = null)
-        {
-            return new EncryptionFeedIterator(
-                this.container.GetChangeFeedStreamIterator(
-                    feedRange,
-                    changeFeedRequestOptions),
-                this.Encryptor,
-                this.CosmosSerializer);
-        }
-
-        public override FeedIterator GetChangeFeedStreamIterator(
-            PartitionKey partitionKey,
-            ChangeFeedRequestOptions changeFeedRequestOptions = null)
-        {
-            return new EncryptionFeedIterator(
-                this.container.GetChangeFeedStreamIterator(
-                    partitionKey,
-                    changeFeedRequestOptions),
-                this.Encryptor,
-                this.CosmosSerializer);
-        }
-
-        public override FeedIterator<T> GetChangeFeedIterator<T>(
-            string continuationToken = null,
-            ChangeFeedRequestOptions changeFeedRequestOptions = null)
-        {
-            return new EncryptionFeedIterator<T>(
-                (EncryptionFeedIterator)this.GetChangeFeedStreamIterator(
-                    continuationToken,
-                    changeFeedRequestOptions),
-                this.ResponseFactory);
-        }
-
-        public override FeedIterator<T> GetChangeFeedIterator<T>(
-            FeedRange feedRange,
-            ChangeFeedRequestOptions changeFeedRequestOptions = null)
-        {
-            return new EncryptionFeedIterator<T>(
-                (EncryptionFeedIterator)this.GetChangeFeedStreamIterator(
-                    feedRange,
-                    changeFeedRequestOptions),
-                this.ResponseFactory);
-        }
-
-        public override FeedIterator<T> GetChangeFeedIterator<T>(
-            PartitionKey partitionKey,
-            ChangeFeedRequestOptions changeFeedRequestOptions = null)
-        {
-            return new EncryptionFeedIterator<T>(
-                (EncryptionFeedIterator)this.GetChangeFeedStreamIterator(
-                    partitionKey,
-                    changeFeedRequestOptions),
-                this.ResponseFactory);
-        }
-
         public override Task<IEnumerable<string>> GetPartitionKeyRangesAsync(
             FeedRange feedRange,
             CancellationToken cancellationToken = default)
@@ -906,5 +837,35 @@ public override FeedIterator<T> GetItemQueryIterator<T>(
                     requestOptions),
                 this.ResponseFactory);
         }
+
+        public override ChangeFeedEstimator GetChangeFeedEstimator(
+            string processorName,
+            Container leaseContainer)
+        {
+            return this.container.GetChangeFeedEstimator(processorName, leaseContainer);
+        }
+
+        public override FeedIterator GetChangeFeedStreamIterator(
+            ChangeFeedStartFrom changeFeedStartFrom,
+            ChangeFeedRequestOptions changeFeedRequestOptions = null)
+        {
+            return new EncryptionFeedIterator(
+                this.container.GetChangeFeedStreamIterator(
+                    changeFeedStartFrom,
+                    changeFeedRequestOptions),
+                this.Encryptor,
+                this.CosmosSerializer);
+        }
+
+        public override FeedIterator<T> GetChangeFeedIterator<T>(
+            ChangeFeedStartFrom changeFeedStartFrom,
+            ChangeFeedRequestOptions changeFeedRequestOptions = null)
+        {
+            return new EncryptionFeedIterator<T>(
+                (EncryptionFeedIterator)this.GetChangeFeedStreamIterator(
+                    changeFeedStartFrom,
+                    changeFeedRequestOptions),
+                this.ResponseFactory);
+        }
     }
 }

Microsoft.Azure.Cosmos.Encryption/src/EncryptionTransactionalBatch.cs

@@ -5,6 +5,7 @@
 namespace Microsoft.Azure.Cosmos.Encryption
 {
     using System;
+    using System.Collections.Generic;
     using System.IO;
     using System.Threading;
     using System.Threading.Tasks;
@@ -199,26 +200,60 @@ public override async Task<TransactionalBatchResponse> ExecuteAsync(
             using (diagnosticsContext.CreateScope("TransactionalBatch.ExecuteAsync"))
             {
                 TransactionalBatchResponse response = await this.transactionalBatch.ExecuteAsync(cancellationToken);
+                return await this.DecryptTransactionalBatchResponseAsync(
+                    response,
+                    diagnosticsContext,
+                    cancellationToken);
+            }
+        }
+
+        public override async Task<TransactionalBatchResponse> ExecuteAsync(
+            TransactionalBatchRequestOptions requestOptions,
+            CancellationToken cancellationToken = default)
+        {
+            CosmosDiagnosticsContext diagnosticsContext = CosmosDiagnosticsContext.Create(options: null);
+            using (diagnosticsContext.CreateScope("TransactionalBatch.ExecuteAsync.WithRequestOptions"))
+            {
+                TransactionalBatchResponse response = await this.transactionalBatch.ExecuteAsync(requestOptions, cancellationToken);
+                return await this.DecryptTransactionalBatchResponseAsync(
+                    response,
+                    diagnosticsContext,
+                    cancellationToken);
+            }
+        }
+
+        private async Task<TransactionalBatchResponse> DecryptTransactionalBatchResponseAsync(
+            TransactionalBatchResponse response,
+            CosmosDiagnosticsContext diagnosticsContext,
+            CancellationToken cancellationToken)
+        {
+            List<TransactionalBatchOperationResult> decryptedTransactionalBatchOperationResults = new List<TransactionalBatchOperationResult>();
 
-                if (response.IsSuccessStatusCode)
+            if (response.IsSuccessStatusCode)
+            {
+                for (int index = 0; index < response.Count; index++)
                 {
-                    for (int index = 0; index < response.Count; index++)
+                    TransactionalBatchOperationResult result = response[index];
+
+                    if (result.ResourceStream != null)
                     {
-                        TransactionalBatchOperationResult result = response[index];
-
-                        if (result.ResourceStream != null)
-                        {
-                            (result.ResourceStream, _) = await EncryptionProcessor.DecryptAsync(
-                                result.ResourceStream,
-                                this.encryptor,
-                                diagnosticsContext,
-                                cancellationToken);
-                        }
+                        (Stream decryptedStream, _) = await EncryptionProcessor.DecryptAsync(
+                            result.ResourceStream,
+                            this.encryptor,
+                            diagnosticsContext,
+                            cancellationToken);
+
+                        result = new EncryptionTransactionalBatchOperationResult(response[index], decryptedStream);
                     }
-                }
 
-                return response;
+                    decryptedTransactionalBatchOperationResults.Add(result);
+                }
             }
+
+            return new EncryptionTransactionalBatchResponse(
+                decryptedTransactionalBatchOperationResults,
+                response,
+                this.cosmosSerializer);
         }
     }
 }

Microsoft.Azure.Cosmos.Encryption/src/EncryptionTransactionalBatchOperationResult.cs

@@ -0,0 +1,32 @@
+//------------------------------------------------------------
+// Copyright (c) Microsoft Corporation.  All rights reserved.
+//------------------------------------------------------------
+
+namespace Microsoft.Azure.Cosmos.Encryption
+{
+    using System;
+    using System.IO;
+    using System.Net;
+
+    internal sealed class EncryptionTransactionalBatchOperationResult : TransactionalBatchOperationResult
+    {
+        private readonly Stream encryptionResourceStream;
+        private readonly TransactionalBatchOperationResult response;
+
+        public EncryptionTransactionalBatchOperationResult(TransactionalBatchOperationResult response, Stream encryptionResourceStream)
+        {
+            this.response = response;
+            this.encryptionResourceStream = encryptionResourceStream;
+        }
+
+        public override Stream ResourceStream => this.encryptionResourceStream;
+
+        public override HttpStatusCode StatusCode => this.response.StatusCode;
+
+        public override bool IsSuccessStatusCode => this.response.IsSuccessStatusCode;
+
+        public override string ETag => this.response.ETag;
+
+        public override TimeSpan RetryAfter => this.response.RetryAfter;
+    }
+}

Microsoft.Azure.Cosmos.Encryption/src/EncryptionTransactionalBatchOperationResult{T}.cs

@@ -0,0 +1,24 @@
+//------------------------------------------------------------
+// Copyright (c) Microsoft Corporation.  All rights reserved.
+//------------------------------------------------------------
+
+namespace Microsoft.Azure.Cosmos.Encryption
+{
+    internal sealed class EncryptionTransactionalBatchOperationResult<T> : TransactionalBatchOperationResult<T>
+    {
+        /// <summary>
+        /// Initializes a new instance of the <see cref="TransactionalBatchOperationResult{T}"/> class.
+        /// </summary>
+        /// <param name="result">BatchOperationResult with stream resource.</param>
+        /// <param name="resource">Deserialized resource.</param>
+        internal EncryptionTransactionalBatchOperationResult(T resource)
+        {
+            this.Resource = resource;
+        }
+
+        /// <summary>
+        /// Gets or sets the content of the resource.
+        /// </summary>
+        public override T Resource { get; set; }
+    }
+}

Microsoft.Azure.Cosmos.Encryption/src/EncryptionTransactionalBatchResponse.cs

@@ -0,0 +1,79 @@
+//------------------------------------------------------------
+// Copyright (c) Microsoft Corporation.  All rights reserved.
+//------------------------------------------------------------
+
+namespace Microsoft.Azure.Cosmos.Encryption
+{
+    using System;
+    using System.Collections.Generic;
+    using System.Net;
+
+    internal sealed class EncryptionTransactionalBatchResponse : TransactionalBatchResponse
+    {
+        private readonly IReadOnlyList<TransactionalBatchOperationResult> results;
+        private readonly TransactionalBatchResponse response;
+        private readonly CosmosSerializer cosmosSerializer;
+        private bool isDisposed = false;
+
+        public EncryptionTransactionalBatchResponse(
+            IReadOnlyList<TransactionalBatchOperationResult> results,
+            TransactionalBatchResponse response,
+            CosmosSerializer cosmosSerializer)
+        {
+            this.results = results;
+            this.response = response;
+            this.cosmosSerializer = cosmosSerializer;
+        }
+
+        public override TransactionalBatchOperationResult this[int index] => this.results[index];
+
+        public override TransactionalBatchOperationResult<T> GetOperationResultAtIndex<T>(int index)
+        {
+            TransactionalBatchOperationResult result = this.results[index];
+
+            T resource = default;
+            if (result.ResourceStream != null)
+            {
+                resource = this.cosmosSerializer.FromStream<T>(result.ResourceStream);
+            }
+
+            return new EncryptionTransactionalBatchOperationResult<T>(resource);
+        }
+
+        public override IEnumerator<TransactionalBatchOperationResult> GetEnumerator()
+        {
+            return this.results.GetEnumerator();
+        }
+
+        public override Headers Headers => this.response.Headers;
+
+        public override string ActivityId => this.response.ActivityId;
+
+        public override double RequestCharge => this.response.RequestCharge;
+
+        public override TimeSpan? RetryAfter => this.response.RetryAfter;
+
+        public override HttpStatusCode StatusCode => this.response.StatusCode;
+
+        public override string ErrorMessage => this.response.ErrorMessage;
+
+        public override bool IsSuccessStatusCode => this.response.IsSuccessStatusCode;
+
+        public override int Count => this.results?.Count ?? 0;
+
+        public override CosmosDiagnostics Diagnostics => this.response.Diagnostics;
+
+        protected override void Dispose(bool disposing)
+        {
+            if (disposing && !this.isDisposed)
+            {
+                this.isDisposed = true;
+
+                if (this.response != null)
+                {
+                    this.response.Dispose();
+                }
+            }
+        }
+    }
+}

Microsoft.Azure.Cosmos.Encryption/src/Microsoft.Azure.Cosmos.Encryption.csproj

@@ -26,7 +26,7 @@
 
   <ItemGroup>
      <PackageReference Include="Azure.Security.KeyVault.Keys" Version="4.0.3" />
-     <PackageReference Include="Microsoft.Azure.Cosmos" Version="3.9.1-preview" />
+     <PackageReference Include="Microsoft.Azure.Cosmos" Version="3.15.0-preview" />
      <PackageReference Include="Azure.Core" Version="1.3.0" />
      <PackageReference Include="Azure.Identity" Version="1.1.1" />
      <PackageReference Include="System.Threading.Tasks.Extensions" Version="4.5.4" />     

Microsoft.Azure.Cosmos.Encryption/src/Mirrored/AssemblyKeys.cs

@@ -0,0 +1,12 @@
+//------------------------------------------------------------
+// Copyright (c) Microsoft Corporation.  All rights reserved.
+//------------------------------------------------------------
+
+namespace Microsoft.Azure.Cosmos.Encryption
+{
+    internal static class AssemblyKeys
+    {
+        /// <summary>TestPublicKey is an unsupported strong key for testing and internal use only</summary>
+        internal const string TestPublicKey = ", PublicKey=0024000004800000940000000602000000240000525341310004000001000100197c25d0a04f73cb271e8181dba1c0c713df8deebb25864541a66670500f34896d280484b45fe1ff6c29f2ee7aa175d8bcbd0c83cc23901a894a86996030f6292ce6eda6e6f3e6c74b3c5a3ded4903c951e6747e6102969503360f7781bf8bf015058eb89b7621798ccc85aaca036ff1bc1556bb7f62de15908484886aa8bbae";
+    }
+}