Skip to content

Commit 0d30a5b

Browse files
DeagleGrossDeagleGross
committed
use static TryHashData for specific hash implementation
1 parent ec70b6e commit 0d30a5b

File tree

2 files changed

+186
-43
lines changed

2 files changed

+186
-43
lines changed

‎src/DataProtection/DataProtection/src/Managed/ManagedAuthenticatedEncryptor.cs‎

Lines changed: 91 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -200,17 +200,34 @@ public byte[] Decrypt(ArraySegment<byte> protectedPayload, ArraySegment<byte> ad
200200
// The best optimization is to stackalloc. If the size is too big, we would want to rent from the pool,
201201
// but we can't due to the HashAlgorithm, ValidationAlgorithm and SymmetricAlgorithm requiring a byte[] instead of a Span<byte>
202202
// in the constructor / Key property.
203+
#if NET10_0_OR_GREATER
204+
byte[]? decryptedKdkLease = null;
205+
Span<byte> decryptedKdk = _keyDerivationKey.Length <= 128
206+
? stackalloc byte[_keyDerivationKey.Length]
207+
: (decryptedKdkLease = DataProtectionPool.Rent(_keyDerivationKey.Length)).AsSpan(0, _keyDerivationKey.Length);
208+
#else
203209
var decryptedKdk = new byte[_keyDerivationKey.Length];
210+
#endif
211+
204212
var decryptionSubkey = new byte[_symmetricAlgorithmSubkeyLengthInBytes];
205213
var validationSubkey = new byte[_validationAlgorithmSubkeyLengthInBytes];
206214

207-
fixed (byte* __unused__1 = decryptedKdk)
215+
fixed (byte* decryptedKdkUnsafe = decryptedKdk)
208216
fixed (byte* __unused__2 = decryptionSubkey)
209217
fixed (byte* __unused__3 = validationSubkey)
210218
{
211219
try
212220
{
213-
_keyDerivationKey.WriteSecretIntoBuffer(new ArraySegment<byte>(decryptedKdk));
221+
_keyDerivationKey.WriteSecretIntoBuffer(decryptedKdkUnsafe, decryptedKdk.Length);
222+
#if NET10_0_OR_GREATER
223+
ManagedSP800_108_CTR_HMACSHA512.DeriveKeysHMACSHA512(
224+
kdk: decryptedKdk,
225+
label: additionalAuthenticatedData,
226+
contextHeader: _contextHeader,
227+
contextData: keyModifier,
228+
operationSubKey: decryptionSubkey,
229+
validationSubKey: validationSubkey);
230+
#else
214231
ManagedSP800_108_CTR_HMACSHA512.DeriveKeys(
215232
kdk: decryptedKdk,
216233
label: additionalAuthenticatedData,
@@ -219,6 +236,7 @@ public byte[] Decrypt(ArraySegment<byte> protectedPayload, ArraySegment<byte> ad
219236
prfFactory: _kdkPrfFactory,
220237
operationSubKey: decryptionSubkey,
221238
validationSubKey: validationSubkey);
239+
#endif
222240

223241
// Step 3: Calculate the correct MAC for this payload.
224242
// correctHash := MAC(IV || ciphertext)
@@ -296,7 +314,19 @@ public byte[] Decrypt(ArraySegment<byte> protectedPayload, ArraySegment<byte> ad
296314
finally
297315
{
298316
// delete since these contain secret material
317+
#if NET10_0_OR_GREATER
318+
if (decryptedKdkLease is not null)
319+
{
320+
DataProtectionPool.Return(decryptedKdkLease, clearArray: true);
321+
}
322+
else
323+
{
324+
decryptedKdk.Clear();
325+
}
326+
#else
299327
Array.Clear(decryptedKdk, 0, decryptedKdk.Length);
328+
#endif
329+
300330
Array.Clear(decryptionSubkey, 0, decryptionSubkey.Length);
301331
Array.Clear(validationSubkey, 0, validationSubkey.Length);
302332
}
@@ -340,65 +370,83 @@ private byte[] EncryptImpl(
340370

341371
// Step 1: Decrypt the KDK, and use it to generate new encryption and HMAC keys.
342372
// We pin all unencrypted keys to limit their exposure via GC relocation.
343-
var decryptedKdk = new byte[_keyDerivationKey.Length];
373+
byte[]? decryptedKdkLease = null;
374+
Span<byte> decryptedKdk = _keyDerivationKey.Length <= 128
375+
? stackalloc byte[_keyDerivationKey.Length]
376+
: (decryptedKdkLease = DataProtectionPool.Rent(_keyDerivationKey.Length)).AsSpan(0, _keyDerivationKey.Length);
377+
344378
var encryptionSubkey = new byte[_symmetricAlgorithmSubkeyLengthInBytes];
345379
var validationSubkey = new byte[_validationAlgorithmSubkeyLengthInBytes];
346380

347-
fixed (byte* __unused__1 = decryptedKdk)
381+
fixed (byte* decryptedKdkUnsafe = decryptedKdk)
348382
fixed (byte* __unused__2 = encryptionSubkey)
349383
fixed (byte* __unused__3 = validationSubkey)
350384
{
351-
var keyModifier = ArrayPool<byte>.Shared.Rent(keyModifierLength);
352-
_genRandom.GenRandom(keyModifier);
385+
var keyModifier = DataProtectionPool.Rent(keyModifierLength);
353386

354-
_keyDerivationKey.WriteSecretIntoBuffer(new ArraySegment<byte>(decryptedKdk));
355-
ManagedSP800_108_CTR_HMACSHA512.DeriveKeys(
356-
kdk: decryptedKdk,
357-
label: additionalAuthenticatedData,
358-
contextHeader: _contextHeader,
359-
contextData: keyModifier,
360-
prfFactory: _kdkPrfFactory,
361-
operationSubKey: encryptionSubkey,
362-
validationSubKey: validationSubkey);
387+
try
388+
{
389+
_genRandom.GenRandom(keyModifier);
363390

364-
// idea of optimization here is firstly get all the types preset
365-
// for calculating length of the output array and allocating it.
366-
// then we are filling it with the data directly, without any additional copying
367-
using var symmetricAlgorithm = CreateSymmetricAlgorithm(key: encryptionSubkey);
368-
using var validationAlgorithm = CreateValidationAlgorithm(key: validationSubkey);
391+
_keyDerivationKey.WriteSecretIntoBuffer(decryptedKdkUnsafe, decryptedKdk.Length);
392+
ManagedSP800_108_CTR_HMACSHA512.DeriveKeysHMACSHA512(
393+
kdk: decryptedKdk,
394+
label: additionalAuthenticatedData,
395+
contextHeader: _contextHeader,
396+
contextData: keyModifier,
397+
operationSubKey: encryptionSubkey,
398+
validationSubKey: validationSubkey);
369399

370-
var cipherTextLength = symmetricAlgorithm.GetCiphertextLengthCbc(plainText.Length); // CBC because symmetricAlgorithm is created with CBC mode
371-
var macLength = _validationAlgorithmDigestLengthInBytes;
400+
// idea of optimization here is firstly get all the types preset
401+
// for calculating length of the output array and allocating it.
402+
// then we are filling it with the data directly, without any additional copying
403+
using var symmetricAlgorithm = CreateSymmetricAlgorithm(key: encryptionSubkey);
404+
using var validationAlgorithm = CreateValidationAlgorithm(key: validationSubkey);
372405

373-
// allocating an array of a specific required length
374-
var outputArray = new byte[keyModifierLength + ivLength + cipherTextLength + macLength];
375-
var outputSpan = outputArray.AsSpan();
406+
var cipherTextLength = symmetricAlgorithm.GetCiphertextLengthCbc(plainText.Length); // CBC because symmetricAlgorithm is created with CBC mode
407+
var macLength = _validationAlgorithmDigestLengthInBytes;
376408

377-
// Step 2: Copy the key modifier and the IV to the output stream since they'll act as a header.
378-
keyModifier.CopyTo(outputSpan.Slice(start: 0, length: keyModifierLength));
409+
// allocating an array of a specific required length
410+
var outputArray = new byte[keyModifierLength + ivLength + cipherTextLength + macLength];
411+
var outputSpan = outputArray.AsSpan();
379412

380-
// Step 3: Generate IV for this operation right into the result array (no allocation)
381-
_genRandom.GenRandom(outputSpan.Slice(start: keyModifierLength, length: ivLength));
382-
var iv = outputSpan.Slice(start: keyModifierLength, length: ivLength);
413+
// Step 2: Copy the key modifier and the IV to the output stream since they'll act as a header.
414+
keyModifier.CopyTo(outputSpan.Slice(start: 0, length: keyModifierLength));
383415

384-
// encrypting plaintext into the target array directly
385-
symmetricAlgorithm.EncryptCbc(plainText, iv, outputSpan.Slice(start: keyModifierLength + ivLength, length: cipherTextLength));
416+
// Step 3: Generate IV for this operation right into the result array (no allocation)
417+
_genRandom.GenRandom(outputSpan.Slice(start: keyModifierLength, length: ivLength));
418+
var iv = outputSpan.Slice(start: keyModifierLength, length: ivLength);
386419

387-
// At this point, outputStream := { keyModifier || IV || ciphertext }
420+
// encrypting plaintext into the target array directly
421+
symmetricAlgorithm.EncryptCbc(plainText, iv, outputSpan.Slice(start: keyModifierLength + ivLength, length: cipherTextLength));
388422

389-
// Step 4: Calculate the digest over the IV and ciphertext.
390-
// We don't need to calculate the digest over the key modifier since that
391-
// value has already been mixed into the KDF used to generate the MAC key.
423+
// At this point, outputStream := { keyModifier || IV || ciphertext }
392424

393-
var ivAndCipherTextSpan = outputSpan.Slice(start: keyModifierLength, length: ivLength + cipherTextLength);
394-
var macDestinationSpan = outputSpan.Slice(keyModifierLength + ivLength + cipherTextLength, macLength);
395-
validationAlgorithm.TryComputeHash(source: ivAndCipherTextSpan, destination: macDestinationSpan, bytesWritten: out _);
396-
// At this point, outputArray := { keyModifier || IV || ciphertext || MAC(IV || ciphertext) }
425+
// Step 4: Calculate the digest over the IV and ciphertext.
426+
// We don't need to calculate the digest over the key modifier since that
427+
// value has already been mixed into the KDF used to generate the MAC key.
397428

398-
// returning whatever was pooled back with clear (secret data to be cleaned)
399-
ArrayPool<byte>.Shared.Return(keyModifier, clearArray: true);
429+
var ivAndCipherTextSpan = outputSpan.Slice(start: keyModifierLength, length: ivLength + cipherTextLength);
430+
var macDestinationSpan = outputSpan.Slice(keyModifierLength + ivLength + cipherTextLength, macLength);
431+
validationAlgorithm.TryComputeHash(source: ivAndCipherTextSpan, destination: macDestinationSpan, bytesWritten: out _);
432+
// At this point, outputArray := { keyModifier || IV || ciphertext || MAC(IV || ciphertext) }
400433

401-
return outputArray;
434+
return outputArray;
435+
}
436+
finally
437+
{
438+
// returning whatever was pooled back with clear (secret data to be cleaned)
439+
DataProtectionPool.Return(keyModifier, clearArray: true);
440+
441+
if (decryptedKdkLease is not null)
442+
{
443+
DataProtectionPool.Return(keyModifier, clearArray: true);
444+
}
445+
else
446+
{
447+
decryptedKdk.Clear();
448+
}
449+
}
402450
}
403451
}
404452
#else

‎src/DataProtection/DataProtection/src/SP800_108/ManagedSP800_108_CTR_HMACSHA512.cs‎

Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,101 @@ public static void DeriveKeys(byte[] kdk, ArraySegment<byte> label, ArraySegment
5757
}
5858
}
5959

60+
#if NET10_0_OR_GREATER
61+
public static void DeriveKeysHMACSHA512(
62+
ReadOnlySpan<byte> kdk,
63+
ReadOnlySpan<byte> label,
64+
ReadOnlySpan<byte> contextHeader,
65+
ReadOnlySpan<byte> contextData,
66+
Span<byte> operationSubKey,
67+
Span<byte> validationSubKey)
68+
{
69+
var operationSubKeyIndex = 0;
70+
var validationSubKeyIndex = 0;
71+
var outputCount = operationSubKey.Length + validationSubKey.Length;
72+
73+
byte[]? prfOutput = null;
74+
75+
// See SP800-108, Sec. 5.1 for the format of the input to the PRF routine.
76+
var prfInputLength = checked(sizeof(uint) /* [i]_2 */ + label.Length + 1 /* 0x00 */ + (contextHeader.Length + contextData.Length) + sizeof(uint) /* [K]_2 */);
77+
78+
byte[]? prfInputLease = null;
79+
Span<byte> prfInput = prfInputLength <= 128
80+
? stackalloc byte[prfInputLength]
81+
: (prfInputLease = DataProtectionPool.Rent(prfInputLength)).AsSpan(0, prfInputLength);
82+
83+
try
84+
{
85+
// Copy [L]_2 to prfInput since it's stable over all iterations
86+
uint outputSizeInBits = (uint)checked((int)outputCount * 8);
87+
prfInput[prfInput.Length - 4] = (byte)(outputSizeInBits >> 24);
88+
prfInput[prfInput.Length - 3] = (byte)(outputSizeInBits >> 16);
89+
prfInput[prfInput.Length - 2] = (byte)(outputSizeInBits >> 8);
90+
prfInput[prfInput.Length - 1] = (byte)(outputSizeInBits);
91+
92+
// Copy label and context to prfInput since they're stable over all iterations
93+
label.CopyTo(prfInput.Slice(sizeof(uint)));
94+
contextHeader.CopyTo(prfInput.Slice(sizeof(uint) + label.Length + 1));
95+
contextData.CopyTo(prfInput.Slice(sizeof(uint) + label.Length + 1 + contextHeader.Length));
96+
97+
for (uint i = 1; outputCount > 0; i++)
98+
{
99+
// Copy [i]_2 to prfInput since it mutates with each iteration
100+
prfInput[0] = (byte)(i >> 24);
101+
prfInput[1] = (byte)(i >> 16);
102+
prfInput[2] = (byte)(i >> 8);
103+
prfInput[3] = (byte)(i);
104+
105+
// Run the PRF and copy the results to the output buffer
106+
// not using stackalloc here, because we are in a loop
107+
// and potentially can exhaust the stack memory: https://learn.microsoft.com/en-us/dotnet/fundamentals/code-analysis/quality-rules/ca2014
108+
prfOutput = DataProtectionPool.Rent(HMACSHA512.HashSizeInBytes);
109+
HMACSHA512.TryHashData(kdk, prfInput, prfOutput, out _);
110+
111+
CryptoUtil.Assert(HMACSHA512.HashSizeInBytes == prfOutput.Length, "prfOutputSizeInBytes == prfOutput.Length");
112+
var numBytesToCopyThisIteration = Math.Min(HMACSHA512.HashSizeInBytes, outputCount);
113+
114+
// we need to write into the operationSubkey
115+
// but it may be the case that we need to split the output
116+
// so lets count how many bytes we can write into the operationSubKey
117+
var bytesToWrite = Math.Min(numBytesToCopyThisIteration, operationSubKey.Length - operationSubKeyIndex);
118+
var leftOverBytes = numBytesToCopyThisIteration - bytesToWrite;
119+
if (operationSubKeyIndex < operationSubKey.Length) // meaning we need to write to operationSubKey
120+
{
121+
var destination = operationSubKey.Slice(operationSubKeyIndex, bytesToWrite);
122+
prfOutput.AsSpan(0, bytesToWrite).CopyTo(destination);
123+
operationSubKeyIndex += bytesToWrite;
124+
}
125+
if (operationSubKeyIndex == operationSubKey.Length && leftOverBytes != 0) // we have filled the operationSubKey. It's time for the validationSubKey
126+
{
127+
var destination = validationSubKey.Slice(validationSubKeyIndex, leftOverBytes);
128+
prfOutput.AsSpan(bytesToWrite, leftOverBytes).CopyTo(destination);
129+
validationSubKeyIndex += leftOverBytes;
130+
}
131+
132+
outputCount -= numBytesToCopyThisIteration;
133+
}
134+
}
135+
finally
136+
{
137+
if (prfOutput is not null)
138+
{
139+
DataProtectionPool.Return(prfOutput, clearArray: true); // contains key material, so delete it
140+
}
141+
142+
if (prfInputLease is not null)
143+
{
144+
DataProtectionPool.Return(prfInputLease, clearArray: true); // contains key material, so delete it
145+
}
146+
else
147+
{
148+
// to be extra careful - clear the stackalloc memory
149+
prfInput.Clear();
150+
}
151+
}
152+
}
153+
#endif
154+
60155
public static void DeriveKeys(
61156
byte[] kdk,
62157
ReadOnlySpan<byte> label,

0 commit comments

Comments
 (0)