feat: "Spanify" DataProtector Protect [part 1]

src/DataProtection/Abstractions/src/IDataProtector.cs

@@ -1,6 +1,7 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using System;
 using System.Diagnostics.CodeAnalysis;
 
 namespace Microsoft.AspNetCore.DataProtection;

src/DataProtection/Abstractions/src/ISpanDataProtector.cs

@@ -0,0 +1,33 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Text;
+using System.Threading.Tasks;
+
+namespace Microsoft.AspNetCore.DataProtection;
+
+/// <summary>
+/// An interface that can provide data protection services.
+/// Is an optimized version of <see cref="IDataProtector"/>.
+/// </summary>
+public interface ISpanDataProtector : IDataProtector
+{
+    /// <summary>
+    /// Determines the size of the protected data in order to then use <see cref="TryProtect(ReadOnlySpan{byte}, Span{byte}, out int)"/>."/>.
+    /// </summary>
+    /// <param name="plainText">The plain text that will be encrypted later</param>
+    /// <returns>The size of the protected data.</returns>
+    int GetProtectedSize(ReadOnlySpan<byte> plainText);
+
+    /// <summary>
+    /// Attempts to encrypt and tamper-proof a piece of data.
+    /// </summary>
+    /// <param name="plainText">The input to encrypt.</param>
+    /// <param name="destination">The ciphertext blob, including authentication tag.</param>
+    /// <param name="bytesWritten">When this method returns, the total number of bytes written into destination</param>
+    /// <returns>true if destination is long enough to receive the encrypted data; otherwise, false.</returns>
+    bool TryProtect(ReadOnlySpan<byte> plainText, Span<byte> destination, out int bytesWritten);
+}

src/DataProtection/Abstractions/src/Microsoft.AspNetCore.DataProtection.Abstractions.csproj

@@ -22,6 +22,10 @@ Microsoft.AspNetCore.DataProtection.IDataProtector</Description>
     <Compile Include="$(SharedSourceRoot)CallerArgument\CallerArgumentExpressionAttribute.cs" LinkBase="Shared" />
   </ItemGroup>
 
+  <ItemGroup Condition="'$(TargetFramework)' != '$(DefaultNetCoreTargetFramework)'">
+    <Reference Include="System.Memory" />
+  </ItemGroup>
+
   <ItemGroup>
     <InternalsVisibleTo Include="Microsoft.AspNetCore.DataProtection.Abstractions.Tests" />
   </ItemGroup>

src/DataProtection/Abstractions/src/PublicAPI.Unshipped.txt

@@ -1 +1,4 @@
 #nullable enable
+Microsoft.AspNetCore.DataProtection.ISpanDataProtector
+Microsoft.AspNetCore.DataProtection.ISpanDataProtector.GetProtectedSize(System.ReadOnlySpan<byte> plainText, out int cipherTextLength) -> int
+Microsoft.AspNetCore.DataProtection.ISpanDataProtector.TryProtect(System.ReadOnlySpan<byte> plainText, System.Span<byte> destination, out int bytesWritten) -> bool

src/DataProtection/DataProtection/src/AuthenticatedEncryption/ISpanAuthenticatedEncryptor.cs

@@ -0,0 +1,37 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Text;
+using System.Threading.Tasks;
+
+namespace Microsoft.AspNetCore.DataProtection.AuthenticatedEncryption;
+
+/// <summary>
+/// Provides an authenticated encryption and decryption routine via a span-based API.
+/// </summary>
+public interface ISpanAuthenticatedEncryptor : IAuthenticatedEncryptor
+{
+    /// <summary>
+    /// Returns the size of the encrypted data for a given plaintext length.
+    /// </summary>
+    /// <param name="plainTextLength">Length of the plain text that will be encrypted later</param>
+    /// <returns>The length of the encrypted data</returns>
+    int GetEncryptedSize(int plainTextLength);
+
+    /// <summary>
+    /// Attempts to encrypt and tamper-proof a piece of data.
+    /// </summary>
+    /// <param name="plaintext">The input to encrypt.</param>
+    /// <param name="additionalAuthenticatedData">
+    /// A piece of data which will not be included in
+    /// the returned ciphertext but which will still be covered by the authentication tag.
+    /// This input may be zero bytes in length. The same AAD must be specified in the corresponding decryption call.
+    /// </param>
+    /// <param name="destination">The ciphertext blob, including authentication tag.</param>
+    /// <param name="bytesWritten">When this method returns, the total number of bytes written into destination</param>
+    /// <returns>true if destination is long enough to receive the encrypted data; otherwise, false.</returns>
+    bool TryEncrypt(ReadOnlySpan<byte> plaintext, ReadOnlySpan<byte> additionalAuthenticatedData, Span<byte> destination, out int bytesWritten);
+}

src/DataProtection/DataProtection/src/Cng/CbcAuthenticatedEncryptor.cs

@@ -2,6 +2,7 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System;
+using System.Runtime.CompilerServices;
 using Microsoft.AspNetCore.Cryptography;
 using Microsoft.AspNetCore.Cryptography.Cng;
 using Microsoft.AspNetCore.Cryptography.SafeHandles;
@@ -299,92 +300,165 @@ private void DoCbcEncrypt(BCryptKeyHandle symmetricKeyHandle, byte* pbIV, byte*
         CryptoUtil.Assert(dwEncryptedBytes == cbOutput, "dwEncryptedBytes == cbOutput");
     }
 
-    protected override byte[] EncryptImpl(byte* pbPlaintext, uint cbPlaintext, byte* pbAdditionalAuthenticatedData, uint cbAdditionalAuthenticatedData, uint cbPreBuffer, uint cbPostBuffer)
+    public override int GetEncryptedSize(int plainTextLength)
     {
-        // This buffer will be used to hold the symmetric encryption and HMAC subkeys
-        // used in the generation of this payload.
-        var cbTempSubkeys = checked(_symmetricAlgorithmSubkeyLengthInBytes + _hmacAlgorithmSubkeyLengthInBytes);
-        byte* pbTempSubkeys = stackalloc byte[checked((int)cbTempSubkeys)];
+        uint paddedCiphertextLength = GetCbcEncryptedOutputSizeWithPadding((uint)plainTextLength);
+        return checked((int)(KEY_MODIFIER_SIZE_IN_BYTES + _symmetricAlgorithmBlockSizeInBytes + paddedCiphertextLength + _hmacAlgorithmDigestLengthInBytes));
+    }
+
+    public override bool TryEncrypt(ReadOnlySpan<byte> plaintext, ReadOnlySpan<byte> additionalAuthenticatedData, Span<byte> destination, out int bytesWritten)
+    {
+        bytesWritten = 0;
 
         try
         {
-            // Randomly generate the key modifier and IV.
-            var cbKeyModifierAndIV = checked(KEY_MODIFIER_SIZE_IN_BYTES + _symmetricAlgorithmBlockSizeInBytes);
-            byte* pbKeyModifierAndIV = stackalloc byte[checked((int)cbKeyModifierAndIV)];
-            _genRandom.GenRandom(pbKeyModifierAndIV, cbKeyModifierAndIV);
+            // This buffer will be used to hold the symmetric encryption and HMAC subkeys
+            // used in the generation of this payload.
+            var cbTempSubkeys = checked(_symmetricAlgorithmSubkeyLengthInBytes + _hmacAlgorithmSubkeyLengthInBytes);
+            byte* pbTempSubkeys = stackalloc byte[checked((int)cbTempSubkeys)];
 
-            // Calculate offsets
-            byte* pbKeyModifier = pbKeyModifierAndIV;
-            byte* pbIV = &pbKeyModifierAndIV[KEY_MODIFIER_SIZE_IN_BYTES];
+            try
+            {
+                // Randomly generate the key modifier and IV.
+                var cbKeyModifierAndIV = checked(KEY_MODIFIER_SIZE_IN_BYTES + _symmetricAlgorithmBlockSizeInBytes);
+                byte* pbKeyModifierAndIV = stackalloc byte[checked((int)cbKeyModifierAndIV)];
+                _genRandom.GenRandom(pbKeyModifierAndIV, cbKeyModifierAndIV);
 
-            // Use the KDF to generate a new symmetric encryption and HMAC subkey
-            _sp800_108_ctr_hmac_provider.DeriveKeyWithContextHeader(
-                pbLabel: pbAdditionalAuthenticatedData,
-                cbLabel: cbAdditionalAuthenticatedData,
-                contextHeader: _contextHeader,
-                pbContext: pbKeyModifier,
-                cbContext: KEY_MODIFIER_SIZE_IN_BYTES,
-                pbDerivedKey: pbTempSubkeys,
-                cbDerivedKey: cbTempSubkeys);
+                // Calculate offsets
+                byte* pbKeyModifier = pbKeyModifierAndIV;
+                byte* pbIV = &pbKeyModifierAndIV[KEY_MODIFIER_SIZE_IN_BYTES];
 
-            // Calculate offsets
-            byte* pbSymmetricEncryptionSubkey = pbTempSubkeys;
-            byte* pbHmacSubkey = &pbTempSubkeys[_symmetricAlgorithmSubkeyLengthInBytes];
+                // Use the KDF to generate a new symmetric encryption and HMAC subkey
+                fixed (byte* pbAdditionalAuthenticatedData = additionalAuthenticatedData)
+                {
+                    _sp800_108_ctr_hmac_provider.DeriveKeyWithContextHeader(
+                        pbLabel: pbAdditionalAuthenticatedData,
+                        cbLabel: (uint)additionalAuthenticatedData.Length,
+                        contextHeader: _contextHeader,
+                        pbContext: pbKeyModifier,
+                        cbContext: KEY_MODIFIER_SIZE_IN_BYTES,
+                        pbDerivedKey: pbTempSubkeys,
+                        cbDerivedKey: cbTempSubkeys);
+                }
 
-            using (var symmetricKeyHandle = _symmetricAlgorithmHandle.GenerateSymmetricKey(pbSymmetricEncryptionSubkey, _symmetricAlgorithmSubkeyLengthInBytes))
-            {
-                // We can't assume PKCS#7 padding (maybe the underlying provider is really using CTS),
-                // so we need to query the padded output size before we can allocate the return value array.
-                var cbOutputCiphertext = GetCbcEncryptedOutputSizeWithPadding(symmetricKeyHandle, pbPlaintext, cbPlaintext);
+                // Calculate offsets
+                byte* pbSymmetricEncryptionSubkey = pbTempSubkeys;
+                byte* pbHmacSubkey = &pbTempSubkeys[_symmetricAlgorithmSubkeyLengthInBytes];
 
-                // Allocate return value array and start copying some data
-                var retVal = new byte[checked(cbPreBuffer + KEY_MODIFIER_SIZE_IN_BYTES + _symmetricAlgorithmBlockSizeInBytes + cbOutputCiphertext + _hmacAlgorithmDigestLengthInBytes + cbPostBuffer)];
-                fixed (byte* pbRetVal = retVal)
+                using (var symmetricKeyHandle = _symmetricAlgorithmHandle.GenerateSymmetricKey(pbSymmetricEncryptionSubkey, _symmetricAlgorithmSubkeyLengthInBytes))
                 {
-                    // Calculate offsets
-                    byte* pbOutputKeyModifier = &pbRetVal[cbPreBuffer];
-                    byte* pbOutputIV = &pbOutputKeyModifier[KEY_MODIFIER_SIZE_IN_BYTES];
-                    byte* pbOutputCiphertext = &pbOutputIV[_symmetricAlgorithmBlockSizeInBytes];
-                    byte* pbOutputHmac = &pbOutputCiphertext[cbOutputCiphertext];
-
-                    UnsafeBufferUtil.BlockCopy(from: pbKeyModifierAndIV, to: pbOutputKeyModifier, byteCount: cbKeyModifierAndIV);
-
-                    // retVal will eventually contain { preBuffer | keyModifier | iv | encryptedData | HMAC(iv | encryptedData) | postBuffer }
-                    // At this point, retVal := { preBuffer | keyModifier | iv | _____ | _____ | postBuffer }
-
-                    DoCbcEncrypt(
-                        symmetricKeyHandle: symmetricKeyHandle,
-                        pbIV: pbIV,
-                        pbInput: pbPlaintext,
-                        cbInput: cbPlaintext,
-                        pbOutput: pbOutputCiphertext,
-                        cbOutput: cbOutputCiphertext);
-
-                    // At this point, retVal := { preBuffer | keyModifier | iv | encryptedData | _____ | postBuffer }
-
-                    // Compute the HMAC over the IV and the ciphertext (prevents IV tampering).
-                    // The HMAC is already implicitly computed over the key modifier since the key
-                    // modifier is used as input to the KDF.
-                    using (var hashHandle = _hmacAlgorithmHandle.CreateHmac(pbHmacSubkey, _hmacAlgorithmSubkeyLengthInBytes))
+                    // Get the padded output size
+                    byte dummy;
+                    fixed (byte* pbPlaintextArray = plaintext)
                     {
-                        hashHandle.HashData(
-                            pbInput: pbOutputIV,
-                            cbInput: checked(_symmetricAlgorithmBlockSizeInBytes + cbOutputCiphertext),
-                            pbHashDigest: pbOutputHmac,
-                            cbHashDigest: _hmacAlgorithmDigestLengthInBytes);
+                        var pbPlaintext = (pbPlaintextArray != null) ? pbPlaintextArray : &dummy;
+                        var cbOutputCiphertext = GetCbcEncryptedOutputSizeWithPadding(symmetricKeyHandle, pbPlaintext, (uint)plaintext.Length);
+
+                        fixed (byte* pbDestination = destination)
+                        {
+                            // Calculate offsets in destination
+                            byte* pbOutputKeyModifier = pbDestination;
+                            byte* pbOutputIV = &pbOutputKeyModifier[KEY_MODIFIER_SIZE_IN_BYTES];
+                            byte* pbOutputCiphertext = &pbOutputIV[_symmetricAlgorithmBlockSizeInBytes];
+                            byte* pbOutputHmac = &pbOutputCiphertext[cbOutputCiphertext];
+
+                            // Copy key modifier and IV to destination
+                            Unsafe.CopyBlock(pbOutputKeyModifier, pbKeyModifierAndIV, cbKeyModifierAndIV);
+                            bytesWritten += checked((int)cbKeyModifierAndIV);
+
+                            // Perform CBC encryption directly into destination
+                            DoCbcEncrypt(
+                                symmetricKeyHandle: symmetricKeyHandle,
+                                pbIV: pbIV,
+                                pbInput: pbPlaintext,
+                                cbInput: (uint)plaintext.Length,
+                                pbOutput: pbOutputCiphertext,
+                                cbOutput: cbOutputCiphertext);
+                            bytesWritten += checked((int)cbOutputCiphertext);
+
+                            // Compute the HMAC over the IV and the ciphertext
+                            using (var hashHandle = _hmacAlgorithmHandle.CreateHmac(pbHmacSubkey, _hmacAlgorithmSubkeyLengthInBytes))
+                            {
+                                hashHandle.HashData(
+                                    pbInput: pbOutputIV,
+                                    cbInput: checked(_symmetricAlgorithmBlockSizeInBytes + cbOutputCiphertext),
+                                    pbHashDigest: pbOutputHmac,
+                                    cbHashDigest: _hmacAlgorithmDigestLengthInBytes);
+                            }
+                            bytesWritten += checked((int)_hmacAlgorithmDigestLengthInBytes);
+
+                            return true;
+                        }
                     }
-
-                    // At this point, retVal := { preBuffer | keyModifier | iv | encryptedData | HMAC(iv | encryptedData) | postBuffer }
-                    // And we're done!
-                    return retVal;
                 }
             }
+            finally
+            {
+                // Buffer contains sensitive material; delete it.
+                UnsafeBufferUtil.SecureZeroMemory(pbTempSubkeys, cbTempSubkeys);
+            }
         }
-        finally
+        catch (Exception ex) when (ex.RequiresHomogenization())
         {
-            // Buffer contains sensitive material; delete it.
-            UnsafeBufferUtil.SecureZeroMemory(pbTempSubkeys, cbTempSubkeys);
+            // Homogenize all exceptions to CryptographicException.
+            throw Error.CryptCommon_GenericError(ex);
+        }
+    }
+
+    public override byte[] Encrypt(ArraySegment<byte> plaintext, ArraySegment<byte> additionalAuthenticatedData, uint preBufferSize, uint postBufferSize)
+    {
+        plaintext.Validate();
+        additionalAuthenticatedData.Validate();
+
+        var size = GetEncryptedSize(plaintext.Count);
+        var ciphertext = new byte[preBufferSize + size + postBufferSize];
+        var destination = ciphertext.AsSpan((int)preBufferSize, size);
+
+        if (!TryEncrypt(
+            plaintext: plaintext,
+            additionalAuthenticatedData: additionalAuthenticatedData,
+            destination: destination,
+            out var bytesWritten))
+        {
+            throw Error.CryptCommon_GenericError(new ArgumentException("Not enough space in destination array"));
         }
+
+        CryptoUtil.Assert(bytesWritten == size, "bytesWritten == size");
+        return ciphertext;
+    }
+
+    /// <summary>
+    /// Should be used only for expected encrypt/decrypt size calculation,
+    /// use the other overload <see cref="GetCbcEncryptedOutputSizeWithPadding(BCryptKeyHandle, byte*, uint)"/>
+    /// for the actual encryption algorithm
+    /// </summary>
+    private uint GetCbcEncryptedOutputSizeWithPadding(uint cbInput)
+    {
+        // Create a temporary key with dummy data for size calculation only
+        // The actual key material doesn't matter for size calculation
+        byte* pbDummyKey = stackalloc byte[checked((int)_symmetricAlgorithmSubkeyLengthInBytes)];
+        // Leave pbDummyKey uninitialized (all zeros) - BCrypt doesn't care for size queries
+
+        using var tempKeyHandle = _symmetricAlgorithmHandle.GenerateSymmetricKey(pbDummyKey, _symmetricAlgorithmSubkeyLengthInBytes);
+
+        // Use uninitialized IV and input data - only the lengths matter
+        byte* pbDummyIV = stackalloc byte[checked((int)_symmetricAlgorithmBlockSizeInBytes)];
+        byte* pbDummyInput = stackalloc byte[checked((int)cbInput)];
+
+        var ntstatus = UnsafeNativeMethods.BCryptEncrypt(
+            hKey: tempKeyHandle,
+            pbInput: pbDummyInput,
+            cbInput: cbInput,
+            pPaddingInfo: null,
+            pbIV: pbDummyIV,
+            cbIV: _symmetricAlgorithmBlockSizeInBytes,
+            pbOutput: null,  // NULL output = size query only
+            cbOutput: 0,
+            pcbResult: out var dwResult,
+            dwFlags: BCryptEncryptFlags.BCRYPT_BLOCK_PADDING);
+        UnsafeNativeMethods.ThrowExceptionForBCryptStatus(ntstatus);
+
+        return dwResult;
     }
 
     private uint GetCbcEncryptedOutputSizeWithPadding(BCryptKeyHandle symmetricKeyHandle, byte* pbInput, uint cbInput)
@@ -394,7 +468,6 @@ private uint GetCbcEncryptedOutputSizeWithPadding(BCryptKeyHandle symmetricKeyHa
 
         // Calling BCryptEncrypt with a null output pointer will cause it to return the total number
         // of bytes required for the output buffer.
-        uint dwResult;
         var ntstatus = UnsafeNativeMethods.BCryptEncrypt(
             hKey: symmetricKeyHandle,
             pbInput: pbInput,
@@ -404,7 +477,7 @@ private uint GetCbcEncryptedOutputSizeWithPadding(BCryptKeyHandle symmetricKeyHa
             cbIV: _symmetricAlgorithmBlockSizeInBytes,
             pbOutput: null,
             cbOutput: 0,
-            pcbResult: out dwResult,
+            pcbResult: out var dwResult,
             dwFlags: BCryptEncryptFlags.BCRYPT_BLOCK_PADDING);
         UnsafeNativeMethods.ThrowExceptionForBCryptStatus(ntstatus);