From c8018c55439faa6c626af6fbf4cb3b0e49010d02 Mon Sep 17 00:00:00 2001
From: Christoffer Lerno <christoffer@aegik.com>
Date: Mon, 30 Sep 2024 00:07:49 +0200
Subject: [PATCH] Added generic PBKDF2 implementation.

---
 lib/std/core/mem.c3           |  5 ++++
 lib/std/hash/hmac.c3          | 54 ++++++++++++++++++++++++++++++++++-
 lib/std/hash/md5.c3           |  1 +
 lib/std/hash/sha1.c3          |  1 +
 releasenotes.md               |  1 +
 test/unit/stdlib/hash/sha1.c3 | 31 ++++++++++++++++++++
 6 files changed, 92 insertions(+), 1 deletion(-)

diff --git a/lib/std/core/mem.c3 b/lib/std/core/mem.c3
index 24264e4af..3ffbd6d85 100644
--- a/lib/std/core/mem.c3
+++ b/lib/std/core/mem.c3
@@ -281,6 +281,11 @@ fn bool ptr_is_aligned(void* ptr, usz alignment) @inline
 	return (uptr)ptr & ((uptr)alignment - 1) == 0;
 }
 
+macro void zero_volatile(char[] data)
+{
+	$$memset(data.ptr, (char)0, data.len, true, (usz)1);
+}
+
 macro void clear(void* dst, usz len, usz $dst_align = 0, bool $is_volatile = false, bool $inlined = false)
 {
 	$$memset(dst, (char)0, len, $is_volatile, $dst_align);
diff --git a/lib/std/hash/hmac.c3 b/lib/std/hash/hmac.c3
index 1bfaca771..540e2368c 100644
--- a/lib/std/hash/hmac.c3
+++ b/lib/std/hash/hmac.c3
@@ -1,4 +1,5 @@
 module std::hash::hmac(<HashAlg, HASH_BYTES, BLOCK_BYTES>);
+import std::crypto;
 
 struct Hmac
 {
@@ -13,6 +14,34 @@ fn char[HASH_BYTES] hash(char[] key, char[] message)
 	return hmac.final();
 }
 
+/**
+ * @require output.len > 0 "Output must be greater than zero"
+ * @require output.len < int.max / HASH_BYTES "Output is too large"
+ **/
+fn void pbkdf2(char[] pw, char[] salt, uint iterations, char[] output)
+{
+	usz l = output.len / HASH_BYTES;
+	usz r = output.len % HASH_BYTES;
+
+	Hmac hmac;
+	hmac.init(pw);
+
+	char[] dst_curr = output;
+	for (usz i = 1; i <= l; i++)
+	{
+		@derive(&hmac, salt, iterations, i, dst_curr[:HASH_BYTES]);
+		dst_curr = dst_curr[HASH_BYTES..];
+	}
+
+	if (r > 0)
+	{
+		char[HASH_BYTES] tmp;
+		@derive(&hmac, salt, iterations, l + 1, &tmp);
+		dst_curr[..] = tmp[:dst_curr.len];
+		mem::zero_volatile(&tmp);
+	}
+}
+
 fn void Hmac.init(&self, char[] key)
 {
 	char[BLOCK_BYTES] buffer;
@@ -37,7 +66,7 @@ fn void Hmac.init(&self, char[] key)
 	self.b.init();
 	self.b.update(&buffer);
 
-	buffer = {};
+	mem::zero_volatile(&buffer);
 }
 
 fn void Hmac.update(&self, char[] data)
@@ -53,3 +82,26 @@ fn char[HASH_BYTES] Hmac.final(&self)
 
 const IPAD @private = 0x36;
 const OPAD @private = 0x5C;
+
+macro @derive(Hmac *hmac_start, char[] salt, uint iterations, usz index, char[] out)
+{
+	assert(out.len == HASH_BYTES);
+	char[HASH_BYTES] tmp @noinit;
+	defer mem::zero_volatile(&tmp);
+	Hmac hmac  = *hmac_start;
+	hmac.update(salt);
+	UIntBE be = { (uint)index };
+	hmac.update(&&bitcast(be, char[4]));
+	tmp = hmac.final();
+	out[..] = tmp;
+	for (int it = 1; it < iterations; it++)
+	{
+		hmac = *hmac_start;
+		hmac.update(&tmp);
+		tmp = hmac.final();
+		foreach (i, v : tmp)
+		{
+			out[i] ^= v;
+		}
+	}
+}
\ No newline at end of file
diff --git a/lib/std/hash/md5.c3 b/lib/std/hash/md5.c3
index b2749c058..5f870629e 100644
--- a/lib/std/hash/md5.c3
+++ b/lib/std/hash/md5.c3
@@ -15,6 +15,7 @@ struct Md5
 
 def HmacMd5 = Hmac(<Md5, HASH_BYTES, BLOCK_BYTES>);
 def hmac = hmac::hash(<Md5, HASH_BYTES, BLOCK_BYTES>);
+def pbkdf2 = hmac::pbkdf2(<Md5, HASH_BYTES, BLOCK_BYTES>);
 
 fn char[HASH_BYTES] hash(char[] data)
 {
diff --git a/lib/std/hash/sha1.c3 b/lib/std/hash/sha1.c3
index 51c15e0ea..57dd1917d 100644
--- a/lib/std/hash/sha1.c3
+++ b/lib/std/hash/sha1.c3
@@ -20,6 +20,7 @@ struct Sha1
 
 def HmacSha1 = Hmac(<Sha1, HASH_BYTES, BLOCK_BYTES>);
 def hmac = hmac::hash(<Sha1, HASH_BYTES, BLOCK_BYTES>);
+def pbkdf2 = hmac::pbkdf2(<Sha1, HASH_BYTES, BLOCK_BYTES>);
 
 fn char[HASH_BYTES] hash(char[] data)
 {
diff --git a/releasenotes.md b/releasenotes.md
index 23ac000db..d6cfa5a66 100644
--- a/releasenotes.md
+++ b/releasenotes.md
@@ -77,6 +77,7 @@
 - Additional timezone related functions for `datetime`.
 - Added MD5 and crypto::safe_compare.
 - Added generic HMAC.
+- Added generic PBKDF2 implementation.
 
 ## 0.6.2 Change list
 
diff --git a/test/unit/stdlib/hash/sha1.c3 b/test/unit/stdlib/hash/sha1.c3
index 1364f92c3..81e736bfc 100644
--- a/test/unit/stdlib/hash/sha1.c3
+++ b/test/unit/stdlib/hash/sha1.c3
@@ -17,6 +17,37 @@ fn void test_sha1_longer()
 	assert(sha.final() == x"84983E44 1C3BD26E BAAE4AA1 F95129E5 E54670F1");
 }
 
+fn void test_pbkdf2()
+{
+	char[] pw = "password";
+	char[] s = "salt";
+	char[20] out;
+	sha1::pbkdf2(pw, s, 1, &out);
+	assert(out == x'0c60c80f961f0e71f3a9b524af6012062fe037a6');
+	sha1::pbkdf2(pw, s, 2, &out);
+	assert(out == x'ea6c014dc72d6f8ccd1ed92ace1d41f0d8de8957');
+	sha1::pbkdf2(pw, s, 4096, &out);
+	assert(out == x'4b007901b765489abead49d926f721d065a429c1');
+}
+
+fn void test_pbkdf2_2()
+{
+	char[] pw = "passwordPASSWORDpassword";
+	char[] s = "saltSALTsaltSALTsaltSALTsaltSALTsalt";
+	char[25] out;
+	sha1::pbkdf2(pw, s, 4096, &out);
+	assert(out == x'3d2eec4fe41c849b80c8d83662c0e44a8b291a964cf2f07038');
+}
+
+fn void test_pbkdf2_3()
+{
+	char[] pw = "pass\0word";
+	char[] salt = "sa\0lt";
+	char[16] out;
+	sha1::pbkdf2(pw, salt, 4096, &out);
+	assert(out == x'56fa6aa75548099dcc37d7f03425e0c3');
+}
+
 fn void test_sha1_million_a()
 {
 	Sha1 sha;