Skip to content

Commit 87d118b

Browse files
DeagleGrossDeagleGross
committed
prfInput and prfOutput rented
1 parent 919d003 commit 87d118b

File tree

1 file changed

+78
-45
lines changed

1 file changed

+78
-45
lines changed

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

Lines changed: 78 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
using System.Buffers;
66
using System.Security.Cryptography;
77
using Microsoft.AspNetCore.Cryptography;
8+
using Microsoft.AspNetCore.DataProtection.Internal;
89
using Microsoft.AspNetCore.DataProtection.Managed;
910

1011
namespace Microsoft.AspNetCore.DataProtection.SP800_108;
@@ -71,67 +72,99 @@ public static void DeriveKeys(
7172

7273
using (var prf = prfFactory(kdk))
7374
{
74-
// See SP800-108, Sec. 5.1 for the format of the input to the PRF routine.
75-
var prfInput = new byte[checked(sizeof(uint) /* [i]_2 */ + label.Length + 1 /* 0x00 */ + (contextHeader.Length + contextData.Length) + sizeof(uint) /* [K]_2 */)];
76-
//var prfInputLength = checked(sizeof(uint) /* [i]_2 */ + label.Length + 1 /* 0x00 */ + (contextHeader.Length + contextData.Length) + sizeof(uint) /* [K]_2 */);
77-
//var prfInput = ArrayPool<byte>.Shared.Rent(prfInputLength);
75+
byte[]? prfOutput = null;
7876

79-
// Copy [L]_2 to prfInput since it's stable over all iterations
80-
uint outputSizeInBits = (uint)checked((int)outputCount * 8);
81-
prfInput[prfInput.Length - 4] = (byte)(outputSizeInBits >> 24);
82-
prfInput[prfInput.Length - 3] = (byte)(outputSizeInBits >> 16);
83-
prfInput[prfInput.Length - 2] = (byte)(outputSizeInBits >> 8);
84-
prfInput[prfInput.Length - 1] = (byte)(outputSizeInBits);
77+
// See SP800-108, Sec. 5.1 for the format of the input to the PRF routine.
78+
var prfInputLength = checked(sizeof(uint) /* [i]_2 */ + label.Length + 1 /* 0x00 */ + (contextHeader.Length + contextData.Length) + sizeof(uint) /* [K]_2 */);
8579

86-
// Copy label and context to prfInput since they're stable over all iterations
87-
label.CopyTo(prfInput.AsSpan(sizeof(uint)));
88-
contextHeader.CopyTo(prfInput.AsSpan(sizeof(uint) + label.Length + 1));
89-
contextData.CopyTo(prfInput.AsSpan(sizeof(uint) + label.Length + 1 + contextHeader.Length));
80+
#if NET10_0_OR_GREATER
81+
byte[]? prfInputLease = null;
82+
Span<byte> prfInput = prfInputLength <= 128
83+
? stackalloc byte[prfInputLength]
84+
: (prfInputLease = DataProtectionPool.Rent(prfInputLength)).AsSpan(0, prfInputLength);
85+
#else
86+
var prfInputArray = new byte[prfInputLength];
87+
var prfInput = prfInputArray.AsSpan();
88+
#endif
9089

91-
var prfOutputSizeInBytes = prf.GetDigestSizeInBytes();
92-
for (uint i = 1; outputCount > 0; i++)
90+
try
9391
{
94-
// Copy [i]_2 to prfInput since it mutates with each iteration
95-
prfInput[0] = (byte)(i >> 24);
96-
prfInput[1] = (byte)(i >> 16);
97-
prfInput[2] = (byte)(i >> 8);
98-
prfInput[3] = (byte)(i);
92+
// Copy [L]_2 to prfInput since it's stable over all iterations
93+
uint outputSizeInBits = (uint)checked((int)outputCount * 8);
94+
prfInput[prfInput.Length - 4] = (byte)(outputSizeInBits >> 24);
95+
prfInput[prfInput.Length - 3] = (byte)(outputSizeInBits >> 16);
96+
prfInput[prfInput.Length - 2] = (byte)(outputSizeInBits >> 8);
97+
prfInput[prfInput.Length - 1] = (byte)(outputSizeInBits);
98+
99+
// Copy label and context to prfInput since they're stable over all iterations
100+
label.CopyTo(prfInput.Slice(sizeof(uint)));
101+
contextHeader.CopyTo(prfInput.Slice(sizeof(uint) + label.Length + 1));
102+
contextData.CopyTo(prfInput.Slice(sizeof(uint) + label.Length + 1 + contextHeader.Length));
103+
104+
var prfOutputSizeInBytes = prf.GetDigestSizeInBytes();
105+
for (uint i = 1; outputCount > 0; i++)
106+
{
107+
// Copy [i]_2 to prfInput since it mutates with each iteration
108+
prfInput[0] = (byte)(i >> 24);
109+
prfInput[1] = (byte)(i >> 16);
110+
prfInput[2] = (byte)(i >> 8);
111+
prfInput[3] = (byte)(i);
99112

100-
// Run the PRF and copy the results to the output buffer
113+
// Run the PRF and copy the results to the output buffer
101114
#if NET10_0_OR_GREATER
102-
var prfOutput = ArrayPool<byte>.Shared.Rent(prfOutputSizeInBytes);
103-
prf.TryComputeHash(prfInput, prfOutput, out _);
115+
// not using stackalloc here, because we are in a loop
116+
// and potentially can exhaust the stack memory: https://learn.microsoft.com/en-us/dotnet/fundamentals/code-analysis/quality-rules/ca2014
117+
prfOutput = DataProtectionPool.Rent(prfOutputSizeInBytes);
118+
prf.TryComputeHash(prfInput, prfOutput, out _);
104119
#else
105-
var prfOutput = prf.ComputeHash(prfInput);
120+
prfOutput = prf.ComputeHash(prfInputArray);
106121
#endif
107122

108-
CryptoUtil.Assert(prfOutputSizeInBytes == prfOutput.Length, "prfOutputSizeInBytes == prfOutput.Length");
109-
var numBytesToCopyThisIteration = Math.Min(prfOutputSizeInBytes, outputCount);
110-
111-
// we need to write into the operationSubkey
112-
// but it may be the case that we need to split the output
113-
// so lets count how many bytes we can write into the operationSubKey
114-
var bytesToWrite = Math.Min(numBytesToCopyThisIteration, operationSubKey.Length - operationSubKeyIndex);
115-
var leftOverBytes = numBytesToCopyThisIteration - bytesToWrite;
116-
if (operationSubKeyIndex < operationSubKey.Length) // meaning we need to write to operationSubKey
123+
CryptoUtil.Assert(prfOutputSizeInBytes == prfOutput.Length, "prfOutputSizeInBytes == prfOutput.Length");
124+
var numBytesToCopyThisIteration = Math.Min(prfOutputSizeInBytes, outputCount);
125+
126+
// we need to write into the operationSubkey
127+
// but it may be the case that we need to split the output
128+
// so lets count how many bytes we can write into the operationSubKey
129+
var bytesToWrite = Math.Min(numBytesToCopyThisIteration, operationSubKey.Length - operationSubKeyIndex);
130+
var leftOverBytes = numBytesToCopyThisIteration - bytesToWrite;
131+
if (operationSubKeyIndex < operationSubKey.Length) // meaning we need to write to operationSubKey
132+
{
133+
var destination = operationSubKey.Slice(operationSubKeyIndex, bytesToWrite);
134+
prfOutput.AsSpan(0, bytesToWrite).CopyTo(destination);
135+
operationSubKeyIndex += bytesToWrite;
136+
}
137+
if (operationSubKeyIndex == operationSubKey.Length && leftOverBytes != 0) // we have filled the operationSubKey. It's time for the validationSubKey
138+
{
139+
var destination = validationSubKey.Slice(validationSubKeyIndex, leftOverBytes);
140+
prfOutput.AsSpan(bytesToWrite, leftOverBytes).CopyTo(destination);
141+
validationSubKeyIndex += leftOverBytes;
142+
}
143+
144+
outputCount -= numBytesToCopyThisIteration;
145+
}
146+
}
147+
finally
148+
{
149+
#if NET10_0_OR_GREATER
150+
if (prfOutput is not null)
117151
{
118-
var destination = operationSubKey.Slice(operationSubKeyIndex, bytesToWrite);
119-
prfOutput.AsSpan(0, bytesToWrite).CopyTo(destination);
120-
operationSubKeyIndex += bytesToWrite;
152+
DataProtectionPool.Return(prfOutput, clearArray: true); // contains key material, so delete it
121153
}
122-
if (operationSubKeyIndex == operationSubKey.Length && leftOverBytes != 0) // we have filled the operationSubKey. It's time for the validationSubKey
154+
155+
if (prfInputLease is not null)
123156
{
124-
var destination = validationSubKey.Slice(validationSubKeyIndex, leftOverBytes);
125-
prfOutput.AsSpan(bytesToWrite, leftOverBytes).CopyTo(destination);
126-
validationSubKeyIndex += leftOverBytes;
157+
DataProtectionPool.Return(prfInputLease, clearArray: true); // contains key material, so delete it
158+
}
159+
else
160+
{
161+
// to be extra careful - clear the stackalloc memory
162+
prfInput.Clear();
127163
}
128-
129-
#if NET10_0_OR_GREATER
130-
ArrayPool<byte>.Shared.Return(prfOutput, clearArray: true); // contains key material, so delete it
131164
#else
165+
Array.Clear(prfInputArray, 0, prfInputArray.Length); // contains key material, so delete it
132166
Array.Clear(prfOutput, 0, prfOutput.Length); // contains key material, so delete it
133167
#endif
134-
outputCount -= numBytesToCopyThisIteration;
135168
}
136169
}
137170
}

0 commit comments

Comments
 (0)