[Client encryption]: Switch to non-allocating MDE 2.0 api

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

@@ -0,0 +1,54 @@
+//------------------------------------------------------------
+// Copyright (c) Microsoft Corporation.  All rights reserved.
+//------------------------------------------------------------
+
+namespace Microsoft.Azure.Cosmos.Encryption.Custom
+{
+    using System;
+    using System.Buffers;
+    using System.Collections.Generic;
+
+#pragma warning disable SA1402 // File may only contain a single type
+    internal class ArrayPoolManager<T> : IDisposable
+#pragma warning restore SA1402 // File may only contain a single type
+    {
+        private List<T[]> rentedBuffers = new ();
+        private bool disposedValue;
+
+        public T[] Rent(int minimumLength)
+        {
+            T[] buffer = ArrayPool<T>.Shared.Rent(minimumLength);
+            this.rentedBuffers.Add(buffer);
+            return buffer;
+        }
+
+        protected virtual void Dispose(bool disposing)
+        {
+            if (!this.disposedValue)
+            {
+                if (disposing)
+                {
+                    foreach (T[] buffer in this.rentedBuffers)
+                    {
+                        ArrayPool<T>.Shared.Return(buffer, clearArray: true);
+                    }
+
+                    this.rentedBuffers = null;
+                }
+
+                this.disposedValue = true;
+            }
+        }
+
+        public void Dispose()
+        {
+            // Do not change this code. Put cleanup code in 'Dispose(bool disposing)' method
+            this.Dispose(disposing: true);
+            GC.SuppressFinalize(this);
+        }
+    }
+
+    internal class ArrayPoolManager : ArrayPoolManager<byte>
+    {
+    }
+}

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

@@ -13,6 +13,15 @@ namespace Microsoft.Azure.Cosmos.Encryption.Custom
     /// </summary>
     public abstract class Encryptor
     {
+        /// <summary>
+        /// Retrieve Data Encryption Key.
+        /// </summary>
+        /// <param name="dataEncryptionKeyId">Identifier of the data encryption key.</param>
+        /// <param name="encryptionAlgorithm">Identifier of the encryption algorithm.</param>
+        /// <param name="cancellationToken">Token for cancellation.</param>
+        /// <returns>Data Encryption Key</returns>
+        public abstract Task<DataEncryptionKey> GetEncryptionKeyAsync(string dataEncryptionKeyId, string encryptionAlgorithm, CancellationToken cancellationToken = default);
+
         /// <summary>
         /// Encrypts the plainText using the key and algorithm provided.
         /// </summary>
@@ -27,15 +36,6 @@ public abstract Task<byte[]> EncryptAsync(
             string encryptionAlgorithm,
             CancellationToken cancellationToken = default);
 
-        /// <summary>
-        /// Retrieve Data Encryption Key.
-        /// </summary>
-        /// <param name="dataEncryptionKeyId">Identifier of the data encryption key.</param>
-        /// <param name="encryptionAlgorithm">Identifier of the encryption algorithm.</param>
-        /// <param name="cancellationToken">Token for cancellation.</param>
-        /// <returns>Data Encryption Key</returns>
-        public abstract Task<DataEncryptionKey> GetEncryptionKeyAsync(string dataEncryptionKeyId, string encryptionAlgorithm, CancellationToken cancellationToken = default);
-
         /// <summary>
         /// Decrypts the cipherText using the key and algorithm provided.
         /// </summary>

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

@@ -0,0 +1,161 @@
+//------------------------------------------------------------
+// Copyright (c) Microsoft Corporation.  All rights reserved.
+//------------------------------------------------------------
+
+namespace Microsoft.Azure.Cosmos.Encryption.Custom
+{
+    using System;
+    using System.Diagnostics.Contracts;
+    using System.IO;
+
+    /// <summary>
+    /// Adjusted implementation of .Net StringReader reading from a Memory<char> instead of a string.
+    /// </summary>
+    internal class MemoryTextReader : TextReader
+    {
+        private Memory<char> chars;
+        private int length;
+        private int pos;
+        private bool closed;
+
+        public MemoryTextReader(Memory<char> chars)
+        {
+            this.chars = chars;
+            this.length = chars.Length;
+        }
+
+        public override void Close()
+        {
+            this.Dispose(true);
+        }
+
+        protected override void Dispose(bool disposing)
+        {
+            this.chars = null;
+            this.pos = 0;
+            this.length = 0;
+            this.closed = true;
+            base.Dispose(disposing);
+        }
+
+        [Pure]
+        public override int Peek()
+        {
+            if (this.closed)
+            {
+                throw new InvalidOperationException("Reader is closed");
+            }
+
+            if (this.pos == this.length)
+            {
+                return -1;
+            }
+
+            return this.chars.Span[this.pos];
+        }
+
+        public override int Read()
+        {
+            if (this.closed)
+            {
+                throw new InvalidOperationException("Reader is closed");
+            }
+
+            if (this.pos == this.length)
+            {
+                return -1;
+            }
+
+            return this.chars.Span[this.pos++];
+        }
+
+        public override int Read(char[] buffer, int index, int count)
+        {
+            if (buffer == null)
+            {
+                throw new ArgumentNullException(nameof(buffer));
+            }
+
+            if (index < 0)
+            {
+                throw new ArgumentOutOfRangeException(nameof(index));
+            }
+
+            if (count < 0)
+            {
+                throw new ArgumentOutOfRangeException(nameof(count));
+            }
+
+            if (buffer.Length - index < count)
+            {
+                throw new ArgumentOutOfRangeException();
+            }
+
+            if (this.closed)
+            {
+                throw new InvalidOperationException("Reader is closed");
+            }
+
+            int n = this.length - this.pos;
+            if (n > 0)
+            {
+                if (n > count)
+                {
+                    n = count;
+                }
+
+                this.chars.Span.Slice(this.pos, n).CopyTo(buffer.AsSpan(index, n));
+                this.pos += n;
+            }
+
+            return n;
+        }
+
+        public override string ReadToEnd()
+        {
+            if (this.closed)
+            {
+                throw new InvalidOperationException("Reader is closed");
+            }
+
+            this.pos = this.length;
+            return new string(this.chars.Slice(this.pos, this.length - this.pos).ToArray());
+        }
+
+        public override string ReadLine()
+        {
+            if (this.closed)
+            {
+                throw new InvalidOperationException("Reader is closed");
+            }
+
+            int i = this.pos;
+            while (i < this.length)
+            {
+                char ch = this.chars.Span[i];
+                if (ch == '\r' || ch == '\n')
+                {
+                    string result = new (this.chars.Slice(this.pos, i - this.pos).ToArray());
+                    this.pos = i + 1;
+                    if (ch == '\r' && this.pos < this.length && this.chars.Span[this.pos] == '\n')
+                    {
+                        this.pos++;
+                    }
+
+                    return result;
+                }
+
+                i++;
+            }
+
+            if (i > this.pos)
+            {
+                string result = new (this.chars.Slice(this.pos, i - this.pos).ToArray());
+                this.pos = i;
+                return result;
+            }
+
+            return null;
+        }
+    }
+}

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

@@ -37,7 +37,7 @@
   <ItemGroup>
      <PackageReference Include="Azure.Core" Version="1.38.0" />
      <PackageReference Include="Azure.Identity" Version="1.11.4" />
-     <PackageReference Include="Microsoft.Data.Encryption.Cryptography" Version="1.2.0" />
+     <PackageReference Include="Microsoft.Data.Encryption.Cryptography" Version="2.0.0-pre007" />
      <PackageReference Include="System.Threading.Tasks.Extensions" Version="4.5.4" />     
   </ItemGroup>
 

Microsoft.Azure.Cosmos.Encryption.Custom/tests/EmulatorTests/MdeCustomEncryptionTests.cs

@@ -370,7 +370,7 @@ public async Task ValidateCachingOfProtectedDataEncryptionKey()
                 await CreateItemAsync(encryptionContainer, dekId, TestDoc.PathsToEncrypt);
 
             testEncryptionKeyStoreProvider.UnWrapKeyCallsCount.TryGetValue(masterKeyUri1.ToString(), out unwrapcount);
-            Assert.AreEqual(32, unwrapcount);
+            Assert.AreEqual(4, unwrapcount);
 
             // 2 hours default
             testEncryptionKeyStoreProvider = new TestEncryptionKeyStoreProvider();

Microsoft.Azure.Cosmos.Encryption.Custom/tests/Microsoft.Azure.Cosmos.Encryption.Custom.Performance.Tests/Readme.md

@@ -9,11 +9,11 @@ Job=MediumRun  Toolchain=InProcessEmitToolchain  IterationCount=15
 LaunchCount=2  WarmupCount=10  
 
 ```
-|  Method | DocumentSizeInKb |        Mean |      Error |     StdDev |     Gen0 |     Gen1 |    Gen2 |  Allocated |
-|-------- |----------------- |------------:|-----------:|-----------:|---------:|---------:|--------:|-----------:|
-| **Encrypt** |                **1** |    **37.15 μs** |   **0.683 μs** |   **1.002 μs** |   **3.8452** |   **0.9766** |       **-** |   **47.28 KB** |
-| Decrypt |                1 |    45.29 μs |   0.757 μs |   1.062 μs |   4.3945 |   1.0986 |       - |   54.17 KB |
-| **Encrypt** |               **10** |   **111.83 μs** |   **1.252 μs** |   **1.874 μs** |  **15.1367** |   **3.0518** |       **-** |  **186.07 KB** |
-| Decrypt |               10 |   151.46 μs |   2.259 μs |   3.311 μs |  19.5313 |   2.1973 |       - |  239.94 KB |
-| **Encrypt** |              **100** | **1,567.24 μs** | **153.944 μs** | **230.416 μs** | **152.3438** | **109.3750** | **76.1719** | **1773.95 KB** |
-| Decrypt |              100 | 2,088.77 μs | 232.084 μs | 347.372 μs | 160.1563 | 113.2813 | 76.1719 | 2042.61 KB |
+|  Method | DocumentSizeInKb |        Mean |      Error |     StdDev |      Median |     Gen0 |     Gen1 |     Gen2 |  Allocated |
+|-------- |----------------- |------------:|-----------:|-----------:|------------:|---------:|---------:|---------:|-----------:|
+| **Encrypt** |                **1** |    **28.40 μs** |   **0.428 μs** |   **0.640 μs** |    **28.40 μs** |   **3.3569** |   **0.8240** |        **-** |   **41.15 KB** |
+| Decrypt |                1 |    33.19 μs |   0.532 μs |   0.779 μs |    33.54 μs |   3.2349 |   0.7935 |        - |    39.7 KB |
+| **Encrypt** |               **10** |   **105.95 μs** |   **2.230 μs** |   **3.337 μs** |   **106.49 μs** |  **13.7939** |   **0.6104** |        **-** |  **169.78 KB** |
+| Decrypt |               10 |   113.47 μs |   1.716 μs |   2.569 μs |   111.81 μs |  12.5732 |   1.2207 |        - |  154.62 KB |
+| **Encrypt** |              **100** | **1,486.58 μs** | **389.596 μs** | **583.129 μs** | **1,487.32 μs** | **216.7969** | **177.7344** | **142.5781** |  **1655.2 KB** |
+| Decrypt |              100 | 1,404.48 μs | 137.824 μs | 206.288 μs | 1,409.23 μs | 144.5313 | 107.4219 |  87.8906 | 1248.31 KB |

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

@@ -41,18 +41,22 @@ public static void ClassInitialize(TestContext testContext)
                 .Returns((int plainTextLength) => plainTextLength);
             DekMock.Setup(m => m.EncryptData(It.IsAny<byte[]>(), It.IsAny<int>(), It.IsAny<int>(), It.IsAny<byte[]>(), It.IsAny<int>()))
                 .Returns((byte[] plainText, int plainTextOffset, int plainTextLength, byte[] output, int outputOffset) => TestCommon.EncryptData(plainText, plainTextOffset, plainTextLength, output, outputOffset));
+            DekMock.Setup(m => m.DecryptData(It.IsAny<byte[]>(), It.IsAny<int>(), It.IsAny<int>(), It.IsAny<byte[]>(), It.IsAny<int>()))
+                .Returns((byte[] plainText, int plainTextOffset, int plainTextLength, byte[] output, int outputOffset) => TestCommon.DecryptData(plainText, plainTextOffset, plainTextLength, output, outputOffset));
+            DekMock.Setup(m => m.GetDecryptByteCount(It.IsAny<int>()))
+                .Returns((int cipherTextLength) => cipherTextLength);
 
 
-            MdeEncryptionProcessorTests.mockEncryptor = new Mock<Encryptor>();
-            MdeEncryptionProcessorTests.mockEncryptor.Setup(m => m.GetEncryptionKeyAsync(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<CancellationToken>()))
+            mockEncryptor = new Mock<Encryptor>();
+            mockEncryptor.Setup(m => m.GetEncryptionKeyAsync(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<CancellationToken>()))
                 .ReturnsAsync((string dekId, string algorithm, CancellationToken token) =>
                     dekId == MdeEncryptionProcessorTests.dekId ? DekMock.Object : throw new InvalidOperationException("DEK not found."));
 
-            MdeEncryptionProcessorTests.mockEncryptor.Setup(m => m.EncryptAsync(It.IsAny<byte[]>(), It.IsAny<string>(), It.IsAny<string>(), It.IsAny<CancellationToken>()))
+            mockEncryptor.Setup(m => m.EncryptAsync(It.IsAny<byte[]>(), It.IsAny<string>(), It.IsAny<string>(), It.IsAny<CancellationToken>()))
                 .ReturnsAsync((byte[] plainText, string dekId, string algo, CancellationToken t) =>
                     dekId == MdeEncryptionProcessorTests.dekId ? TestCommon.EncryptData(plainText) : throw new InvalidOperationException("DEK not found."));
 
-            MdeEncryptionProcessorTests.mockEncryptor.Setup(m => m.DecryptAsync(It.IsAny<byte[]>(), It.IsAny<string>(), It.IsAny<string>(), It.IsAny<CancellationToken>()))
+            mockEncryptor.Setup(m => m.DecryptAsync(It.IsAny<byte[]>(), It.IsAny<string>(), It.IsAny<string>(), It.IsAny<CancellationToken>()))
                 .ReturnsAsync((byte[] cipherText, string dekId, string algo, CancellationToken t) =>
                     dekId == MdeEncryptionProcessorTests.dekId ? TestCommon.DecryptData(cipherText) : throw new InvalidOperationException("Null DEK was returned."));
         }