Skip to content

Commit 41fd5e2

Browse files
HMBSbigeHMBSbige
committed
feat: Add support for overlapping output and info buffers in Hkdf.Expand
1 parent 888778b commit 41fd5e2

File tree

4 files changed

+77
-46
lines changed

4 files changed

+77
-46
lines changed

‎src/CryptoBase/KDF/Hkdf.cs‎

Lines changed: 59 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
using CryptoBase.Abstractions;
22
using CryptoBase.Digests;
33
using CryptoBase.Macs.Hmac;
4+
using System.Security.Cryptography;
45

56
namespace CryptoBase.KDF;
67

@@ -14,11 +15,6 @@ public static int Extract(DigestType type, ReadOnlySpan<byte> ikm, ReadOnlySpan<
1415
int hashLength = HashLength(type);
1516
ArgumentOutOfRangeException.ThrowIfLessThan(prk.Length, hashLength, nameof(prk));
1617

17-
if (prk.Length > hashLength)
18-
{
19-
prk = prk[..hashLength];
20-
}
21-
2218
ExtractInternal(type, ikm, salt, prk);
2319

2420
return hashLength;
@@ -38,6 +34,8 @@ public static void Expand(DigestType type, ReadOnlySpan<byte> prk, Span<byte> ou
3834

3935
ArgumentOutOfRangeException.ThrowIfZero(output.Length, nameof(output));
4036

37+
ArgumentOutOfRangeException.ThrowIfLessThan(prk.Length, hashLength, nameof(prk));
38+
4139
int maxOkmLength = 255 * hashLength;
4240
ArgumentOutOfRangeException.ThrowIfGreaterThan(output.Length, maxOkmLength, nameof(output));
4341

@@ -46,46 +44,68 @@ public static void Expand(DigestType type, ReadOnlySpan<byte> prk, Span<byte> ou
4644

4745
private static void ExpandInternal(DigestType type, int hashLength, ReadOnlySpan<byte> prk, Span<byte> output, ReadOnlySpan<byte> info)
4846
{
49-
ArgumentOutOfRangeException.ThrowIfLessThan(prk.Length, hashLength, nameof(prk));
50-
51-
if (output.Overlaps(info))
52-
{
53-
throw new InvalidOperationException(@"the info input overlaps with the output destination");
54-
}
55-
56-
Span<byte> counterSpan = stackalloc byte[1];
57-
ref byte counter = ref counterSpan[0];
47+
byte counter = 0;
48+
Span<byte> counterSpan = new(ref counter);
5849
Span<byte> t = Span<byte>.Empty;
5950
Span<byte> remainingOutput = output;
6051

61-
using IMac hmac = HmacUtils.Create(type, prk);
52+
const int maxStackInfoBuffer = 64;
53+
Span<byte> tempInfoBuffer = stackalloc byte[maxStackInfoBuffer];
54+
scoped ReadOnlySpan<byte> infoBuffer;
55+
byte[]? rentedTempInfoBuffer = null;
6256

63-
for (int i = 1; ; ++i)
57+
if (output.Overlaps(info))
6458
{
65-
hmac.Update(t);
66-
hmac.Update(info);
67-
counter = (byte)i;
68-
hmac.Update(counterSpan);
69-
70-
if (remainingOutput.Length >= hashLength)
59+
if (info.Length > maxStackInfoBuffer)
7160
{
72-
t = remainingOutput[..hashLength];
73-
remainingOutput = remainingOutput[hashLength..];
74-
hmac.GetMac(t);
61+
rentedTempInfoBuffer = ArrayPool<byte>.Shared.Rent(info.Length);
62+
tempInfoBuffer = rentedTempInfoBuffer;
7563
}
76-
else
64+
65+
tempInfoBuffer = tempInfoBuffer.Slice(0, info.Length);
66+
info.CopyTo(tempInfoBuffer);
67+
infoBuffer = tempInfoBuffer;
68+
}
69+
else
70+
{
71+
infoBuffer = info;
72+
}
73+
74+
using (IMac hmac = HmacUtils.Create(type, prk))
75+
{
76+
for (int i = 1; ; ++i)
7777
{
78-
if (remainingOutput.Length > 0)
78+
hmac.Update(t);
79+
hmac.Update(infoBuffer);
80+
counter = (byte)i;
81+
hmac.Update(counterSpan);
82+
83+
if (remainingOutput.Length >= hashLength)
7984
{
80-
// ReSharper disable once StackAllocInsideLoop
81-
Span<byte> lastChunk = stackalloc byte[hashLength];
82-
hmac.GetMac(lastChunk);
83-
lastChunk[..remainingOutput.Length].CopyTo(remainingOutput);
85+
t = remainingOutput.Slice(0, hashLength);
86+
remainingOutput = remainingOutput.Slice(hashLength);
87+
hmac.GetMac(t);
88+
}
89+
else
90+
{
91+
if (remainingOutput.Length > 0)
92+
{
93+
// ReSharper disable once StackAllocInsideLoop
94+
Span<byte> lastChunk = stackalloc byte[hashLength];
95+
hmac.GetMac(lastChunk);
96+
lastChunk.Slice(0, remainingOutput.Length).CopyTo(remainingOutput);
97+
}
98+
99+
break;
84100
}
85-
86-
break;
87101
}
88102
}
103+
104+
if (rentedTempInfoBuffer is not null)
105+
{
106+
CryptographicOperations.ZeroMemory(rentedTempInfoBuffer.AsSpan(0, info.Length));
107+
ArrayPool<byte>.Shared.Return(rentedTempInfoBuffer);
108+
}
89109
}
90110

91111
public static void DeriveKey(DigestType type, ReadOnlySpan<byte> ikm, Span<byte> output, ReadOnlySpan<byte> salt, ReadOnlySpan<byte> info)
@@ -97,10 +117,16 @@ public static void DeriveKey(DigestType type, ReadOnlySpan<byte> ikm, Span<byte>
97117
int maxOkmLength = 255 * hashLength;
98118
ArgumentOutOfRangeException.ThrowIfGreaterThan(output.Length, maxOkmLength, nameof(output));
99119

120+
DeriveKeyInternal(type, hashLength, ikm, output, salt, info);
121+
}
122+
123+
private static void DeriveKeyInternal(DigestType type, int hashLength, ReadOnlySpan<byte> ikm, Span<byte> output, ReadOnlySpan<byte> salt, ReadOnlySpan<byte> info)
124+
{
100125
Span<byte> prk = stackalloc byte[hashLength];
101126

102127
ExtractInternal(type, ikm, salt, prk);
103128
ExpandInternal(type, hashLength, prk, output, info);
129+
CryptographicOperations.ZeroMemory(prk);
104130
}
105131

106132
private static int HashLength(DigestType type)

‎src/CryptoBase/Macs/Hmac/HmacUtils.cs‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@ public static IMac Create(ReadOnlySpan<byte> key, HashAlgorithmName name)
1616
return new DefaultHmac(key, name);
1717
}
1818

19-
[MethodImpl(MethodImplOptions.AggressiveInlining)]
2019
public static IMac Create(DigestType type, ReadOnlySpan<byte> key)
2120
{
2221
return type switch

‎test/CryptoBase.Benchmark/HKDFBenchmark.cs‎

Lines changed: 11 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ namespace CryptoBase.Benchmark;
88
[MemoryDiagnoser]
99
public class HKDFBenchmark
1010
{
11-
[Params(1)]
11+
[Params(10)]
1212
public int Max { get; set; }
1313

1414
private byte[] _ikm = null!;
@@ -23,23 +23,25 @@ public void Setup()
2323
_info = RandomNumberGenerator.GetBytes(80);
2424
}
2525

26-
[Benchmark]
27-
public void Default()
26+
[Benchmark(Baseline = true)]
27+
public void NET()
2828
{
2929
Span<byte> output = stackalloc byte[82];
30-
for (var i = 0; i < Max; ++i)
30+
31+
for (int i = 0; i < Max; ++i)
3132
{
32-
Hkdf.DeriveKey(DigestType.Sha256, _ikm, output, _salt, _info);
33+
HKDF.DeriveKey(HashAlgorithmName.SHA256, _ikm, output, _salt, _info);
3334
}
3435
}
3536

36-
[Benchmark(Baseline = true)]
37-
public void NET()
37+
[Benchmark]
38+
public void Default()
3839
{
3940
Span<byte> output = stackalloc byte[82];
40-
for (var i = 0; i < Max; ++i)
41+
42+
for (int i = 0; i < Max; ++i)
4143
{
42-
HKDF.DeriveKey(HashAlgorithmName.SHA256, _ikm, output, _salt, _info);
44+
Hkdf.DeriveKey(DigestType.Sha256, _ikm, output, _salt, _info);
4345
}
4446
}
4547
}

‎test/CryptoBase.Tests/HkdfTest.cs‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -20,14 +20,18 @@ public class HkdfTest
2020
public void Test(DigestType type, string ikmHex, string saltHex, string infoHex, string prkHex, string okmHex)
2121
{
2222
Span<byte> prk = stackalloc byte[255];
23+
Span<byte> info = infoHex.FromHex();
2324

2425
int length = Hkdf.Extract(type, ikmHex.FromHex(), saltHex.FromHex(), prk);
25-
Assert.Equal(prkHex, prk[..length].ToHex());
26+
Assert.Equal(prkHex, prk.Slice(0, length).ToHex());
2627

2728
Span<byte> okm = stackalloc byte[okmHex.FromHex().Length];
28-
Hkdf.Expand(type, prkHex.FromHex(), okm, infoHex.FromHex());
29+
Hkdf.Expand(type, prkHex.FromHex(), okm, info);
2930
Assert.Equal(okmHex, okm.ToHex());
3031

31-
Hkdf.DeriveKey(type, ikmHex.FromHex(), okm, saltHex.FromHex(), infoHex.FromHex());
32+
Hkdf.DeriveKey(type, ikmHex.FromHex(), okm, saltHex.FromHex(), info);
33+
34+
info.CopyTo(okm);
35+
Hkdf.DeriveKey(type, ikmHex.FromHex(), okm, saltHex.FromHex(), okm.Slice(0, info.Length));
3236
}
3337
}

0 commit comments

Comments
 (0)