refactor

Author: DeagleGross

src/DataProtection/DataProtection/src/AuthenticatedEncryption/AuthenticatedEncryptorExtensions.cs

@@ -38,12 +38,13 @@ public static bool TryEncrypt(
         ReadOnlySpan<byte> plaintext,
         ReadOnlySpan<byte> additionalAuthenticatedData,
         Span<byte> destination,
-        uint preBufferSize,
-        uint postBufferSize,
+        int preBufferSize,
+        int postBufferSize,
         out int bytesWritten)
     {
-        var plaintextWithOffsets = plaintext.Slice((int)preBufferSize, plaintext.Length - (int)(preBufferSize + postBufferSize));
-        return encryptor.TryEncrypt(plaintextWithOffsets, additionalAuthenticatedData, destination, out bytesWritten);
+        var destinationBufferOffsets = destination.Slice(preBufferSize, destination.Length - (preBufferSize + postBufferSize));
+        // var plaintextWithOffsets = plaintext.Slice((int)preBufferSize, plaintext.Length - (int)(preBufferSize + postBufferSize));
+        return encryptor.TryEncrypt(plaintext, additionalAuthenticatedData, destinationBufferOffsets, out bytesWritten);
     }
 #endif
 

src/DataProtection/DataProtection/src/AuthenticatedEncryption/IOptimizedAuthenticatedEncryptor.cs

@@ -38,8 +38,4 @@ internal interface IOptimizedAuthenticatedEncryptor : IAuthenticatedEncryptor
     /// All cryptography-related exceptions should be homogenized to CryptographicException.
     /// </remarks>
     byte[] Encrypt(ArraySegment<byte> plaintext, ArraySegment<byte> additionalAuthenticatedData, uint preBufferSize, uint postBufferSize);
-
-#if NET10_0_OR_GREATER
-    int GetEncryptedSize(int plainTextLength, uint preBufferSize, uint postBufferSize);
-#endif
 }

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

@@ -300,10 +300,10 @@ private void DoCbcEncrypt(BCryptKeyHandle symmetricKeyHandle, byte* pbIV, byte*
     }
 
 #if NET10_0_OR_GREATER
-    public override int GetEncryptedSize(int plainTextLength, uint preBufferSize, uint postBufferSize)
+    public override int GetEncryptedSize(int plainTextLength)
     {
         uint paddedCiphertextLength = GetCbcEncryptedOutputSizeWithPadding((uint)plainTextLength);
-        return checked((int)(preBufferSize + KEY_MODIFIER_SIZE_IN_BYTES + _symmetricAlgorithmBlockSizeInBytes + paddedCiphertextLength + _hmacAlgorithmDigestLengthInBytes + postBufferSize));
+        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)

src/DataProtection/DataProtection/src/Cng/CngGcmAuthenticatedEncryptor.cs

@@ -232,11 +232,11 @@ private void DoGcmEncrypt(byte* pbKey, uint cbKey, byte* pbNonce, byte* pbPlaint
     }
 
 #if NET10_0_OR_GREATER
-    public override int GetEncryptedSize(int plainTextLength, uint preBufferSize, uint postBufferSize)
+    public override int GetEncryptedSize(int plainTextLength)
     {
         // A buffer to hold the key modifier, nonce, encrypted data, and tag.
         // In GCM, the encrypted output will be the same length as the plaintext input.
-        return checked((int)(preBufferSize + KEY_MODIFIER_SIZE_IN_BYTES + NONCE_SIZE_IN_BYTES + plainTextLength + TAG_SIZE_IN_BYTES + postBufferSize));
+        return checked((int)(KEY_MODIFIER_SIZE_IN_BYTES + NONCE_SIZE_IN_BYTES + plainTextLength + TAG_SIZE_IN_BYTES));
     }
 
     public override bool TryEncrypt(ReadOnlySpan<byte> plainText, ReadOnlySpan<byte> additionalAuthenticatedData, Span<byte> destination, out int bytesWritten)

src/DataProtection/DataProtection/src/Cng/Internal/CngAuthenticatedEncryptorBase.cs

@@ -85,8 +85,7 @@ public byte[] Encrypt(ArraySegment<byte> plaintext, ArraySegment<byte> additiona
     protected abstract byte[] EncryptImpl(byte* pbPlaintext, uint cbPlaintext, byte* pbAdditionalAuthenticatedData, uint cbAdditionalAuthenticatedData, uint cbPreBuffer, uint cbPostBuffer);
 
 #if NET10_0_OR_GREATER
-    public int GetEncryptedSize(int plainTextLength) => GetEncryptedSize(plainTextLength, 0, 0);
-    public abstract int GetEncryptedSize(int plainTextLength, uint preBufferSize, uint postBufferSize);
+    public abstract int GetEncryptedSize(int plainTextLength);
     public abstract bool TryEncrypt(ReadOnlySpan<byte> plainText, ReadOnlySpan<byte> additionalAuthenticatedData, Span<byte> destination, out int bytesWritten);
 #endif
 }

src/DataProtection/DataProtection/src/KeyManagement/KeyRingBasedDataProtector.cs

@@ -2,22 +2,24 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System;
-using System.Buffers.Binary;
 using System.Buffers;
+using System.Buffers.Binary;
+using System.Buffers.Text;
 using System.Collections.Generic;
 using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
 using System.IO;
 using System.Linq;
+using System.Reflection.PortableExecutable;
 using System.Runtime.CompilerServices;
 using System.Threading;
 using Microsoft.AspNetCore.Cryptography;
 using Microsoft.AspNetCore.DataProtection.AuthenticatedEncryption;
+using Microsoft.AspNetCore.DataProtection.Internal;
 using Microsoft.AspNetCore.DataProtection.KeyManagement.Internal;
 using Microsoft.AspNetCore.Shared;
 using Microsoft.Extensions.Logging;
-using System.Buffers.Text;
-using Microsoft.AspNetCore.DataProtection.Internal;
+using static System.Runtime.InteropServices.JavaScript.JSType;
 
 namespace Microsoft.AspNetCore.DataProtection.KeyManagement;
 
@@ -34,6 +36,8 @@ internal sealed unsafe class KeyRingBasedDataProtector : IDataProtector, IPersis
     private readonly IKeyRingProvider _keyRingProvider;
     private readonly ILogger? _logger;
 
+    private static readonly int _magicHeaderKeyIdSize = (sizeof(uint) + sizeof(Guid));
+
     public KeyRingBasedDataProtector(IKeyRingProvider keyRingProvider, ILogger? logger, string[]? originalPurposes, string newPurpose)
     {
         Debug.Assert(keyRingProvider != null);
@@ -98,7 +102,9 @@ public int GetProtectedSize(ReadOnlySpan<byte> plainText)
         var defaultEncryptor = currentKeyRing.DefaultAuthenticatedEncryptor;
         CryptoUtil.Assert(defaultEncryptor != null, "defaultEncryptorInstance != null");
 
-        return defaultEncryptor.GetEncryptedSize(plainText.Length);
+        // We allocate a 20-byte pre-buffer so that we can inject the magic header and key id into the return value.
+        // See Protect() / TryProtect() for details
+        return _magicHeaderKeyIdSize + defaultEncryptor.GetEncryptedSize(plainText.Length);
     }
 
     public bool TryProtect(ReadOnlySpan<byte> plaintext, Span<byte> destination, out int bytesWritten)
@@ -125,7 +131,7 @@ public bool TryProtect(ReadOnlySpan<byte> plaintext, Span<byte> destination, out
                 plaintext: plaintext,
                 additionalAuthenticatedData: aad,
                 destination: destination,
-                preBufferSize: (uint)(sizeof(uint) + sizeof(Guid)),
+                preBufferSize: _magicHeaderKeyIdSize,
                 postBufferSize: 0,
                 out bytesWritten);
 
@@ -136,6 +142,7 @@ public bool TryProtect(ReadOnlySpan<byte> plaintext, Span<byte> destination, out
             BinaryPrimitives.WriteUInt32BigEndian(destination.Slice(0, sizeof(uint)), MAGIC_HEADER_V0);
             var writeKeyIdResult = defaultKeyId.TryWriteBytes(destination.Slice(sizeof(uint), sizeof(Guid)));
             Debug.Assert(writeKeyIdResult, "Failed to write Guid to destination.");
+            bytesWritten += _magicHeaderKeyIdSize;
 
             // At this point, destination := { magicHeader || keyId || encryptorSpecificProtectedPayload }
             // And we're done!

src/DataProtection/DataProtection/src/Managed/AesGcmAuthenticatedEncryptor.cs

@@ -142,13 +142,10 @@ public byte[] Decrypt(ArraySegment<byte> ciphertext, ArraySegment<byte> addition
     }
 
     public int GetEncryptedSize(int plainTextLength)
-        => GetEncryptedSize(plainTextLength, 0, 0);
-
-    public int GetEncryptedSize(int plainTextLength, uint preBufferSize, uint postBufferSize)
     {
         // A buffer to hold the key modifier, nonce, encrypted data, and tag.
         // In GCM, the encrypted output will be the same length as the plaintext input.
-        return checked((int)(preBufferSize + KEY_MODIFIER_SIZE_IN_BYTES + NONCE_SIZE_IN_BYTES + plainTextLength + TAG_SIZE_IN_BYTES + postBufferSize));
+        return checked((int)(KEY_MODIFIER_SIZE_IN_BYTES + NONCE_SIZE_IN_BYTES + plainTextLength + TAG_SIZE_IN_BYTES));
     }
 
     public bool TryEncrypt(ReadOnlySpan<byte> plaintext, ReadOnlySpan<byte> additionalAuthenticatedData, Span<byte> destination, out int bytesWritten)

src/DataProtection/DataProtection/test/Microsoft.AspNetCore.DataProtection.Tests/KeyManagement/KeyRingBasedDataProtectorTests.cs

@@ -623,7 +623,44 @@ public void CreateProtector_ChainsPurposes()
     [Fact]
     public void GetProtectedSize_TryProtect_CorrectlyEstimatesDataLength()
     {
-        // TODO!
+        // Arrange
+        byte[] plaintext = new byte[] { 0x10, 0x20, 0x30, 0x40, 0x50 };
+        var encryptorFactory = new AuthenticatedEncryptorFactory(NullLoggerFactory.Instance);
+        Key key = new Key(Guid.NewGuid(), DateTimeOffset.Now, DateTimeOffset.Now, DateTimeOffset.Now, new AuthenticatedEncryptorConfiguration().CreateNewDescriptor(), new[] { encryptorFactory });
+        var keyRing = new KeyRing(key, new[] { key });
+        var mockKeyRingProvider = new Mock<IKeyRingProvider>();
+        mockKeyRingProvider.Setup(o => o.GetCurrentKeyRing()).Returns(keyRing);
+
+        var protector = new KeyRingBasedDataProtector(
+            keyRingProvider: mockKeyRingProvider.Object,
+            logger: GetLogger(),
+            originalPurposes: null,
+            newPurpose: "purpose");
+
+        // Act - get estimated size
+        int estimatedSize = protector.GetProtectedSize(plaintext);
+        
+        // Act - allocate buffer and try protect
+        byte[] destination = new byte[estimatedSize];
+        bool success = protector.TryProtect(plaintext, destination, out int bytesWritten);
+
+        // Assert
+        Assert.True(success, "TryProtect should succeed with estimated buffer size");
+        Assert.Equal(estimatedSize, bytesWritten);
+        Assert.True(bytesWritten > 0, "Should write some bytes");
+        Assert.True(bytesWritten >= plaintext.Length, "Protected data should be at least as large as plaintext");
+
+        // Verify the protected data can be unprotected to get original plaintext
+        byte[] actualDestination = new byte[bytesWritten];
+        Array.Copy(destination, actualDestination, bytesWritten);
+        byte[] unprotectedData = protector.Unprotect(actualDestination);
+        Assert.Equal(plaintext, unprotectedData);
+
+        // Test with buffer that's too small
+        byte[] smallBuffer = new byte[estimatedSize - 1];
+        bool smallBufferSuccess = protector.TryProtect(plaintext, smallBuffer, out int smallBufferBytesWritten);
+        Assert.False(smallBufferSuccess, "TryProtect should fail with buffer that's too small");
+        Assert.Equal(0, smallBufferBytesWritten);
     }
 
     private static byte[] BuildAadFromPurposeStrings(Guid keyId, params string[] purposes)