From 0d83487dd098f9498d907f06535e645cc8bcf3de Mon Sep 17 00:00:00 2001
From: Frank Denis <github@pureftpd.org>
Date: Thu, 29 Dec 2022 11:29:50 +0100
Subject: [PATCH] hkdf: add prk_length and extractInit()

The HKDF extract function uses HMAC under the hood, but requiring
applications to directly use HMAC functions reduces clarity and
feels like the wrong abstraction.

So, in order to get the PRK length, add a `prk_length` constant
that applications can use directly.

Also add an `extractInit()` function for cases where the keying
material has to be provided as multiple chunks.
---
 lib/std/crypto/hkdf.zig | 45 +++++++++++++++++++++++++++++++----------
 1 file changed, 34 insertions(+), 11 deletions(-)

diff --git a/lib/std/crypto/hkdf.zig b/lib/std/crypto/hkdf.zig
index 637cf42297..7102ffe780 100644
--- a/lib/std/crypto/hkdf.zig
+++ b/lib/std/crypto/hkdf.zig
@@ -13,39 +13,56 @@ pub const HkdfSha512 = Hkdf(hmac.sha2.HmacSha512);
 /// derives one or more uniform keys from it.
 pub fn Hkdf(comptime Hmac: type) type {
     return struct {
+        /// Length of a master key, in bytes.
+        pub const prk_length = Hmac.mac_length;
+
         /// Return a master key from a salt and initial keying material.
-        pub fn extract(salt: []const u8, ikm: []const u8) [Hmac.mac_length]u8 {
-            var prk: [Hmac.mac_length]u8 = undefined;
+        pub fn extract(salt: []const u8, ikm: []const u8) [prk_length]u8 {
+            var prk: [prk_length]u8 = undefined;
             Hmac.create(&prk, ikm, salt);
             return prk;
         }
 
+        /// Initialize the creation of a master key from a salt
+        /// and keying material that can be added later, possibly in chunks.
+        /// Example:
+        /// ```
+        /// var prk: [hkdf.prk_length]u8 = undefined;
+        /// var hkdf = HkdfSha256.extractInit(salt);
+        /// hkdf.update(ikm1);
+        /// hkdf.update(ikm2);
+        /// hkdf.final(&prk);
+        /// ```
+        pub fn extractInit(salt: []const u8) Hmac {
+            return Hmac.init(salt);
+        }
+
         /// Derive a subkey from a master key `prk` and a subkey description `ctx`.
-        pub fn expand(out: []u8, ctx: []const u8, prk: [Hmac.mac_length]u8) void {
-            assert(out.len <= Hmac.mac_length * 255); // output size is too large for the Hkdf construction
+        pub fn expand(out: []u8, ctx: []const u8, prk: [prk_length]u8) void {
+            assert(out.len <= prk_length * 255); // output size is too large for the Hkdf construction
             var i: usize = 0;
             var counter = [1]u8{1};
-            while (i + Hmac.mac_length <= out.len) : (i += Hmac.mac_length) {
+            while (i + prk_length <= out.len) : (i += prk_length) {
                 var st = Hmac.init(&prk);
                 if (i != 0) {
-                    st.update(out[i - Hmac.mac_length ..][0..Hmac.mac_length]);
+                    st.update(out[i - prk_length ..][0..prk_length]);
                 }
                 st.update(ctx);
                 st.update(&counter);
-                st.final(out[i..][0..Hmac.mac_length]);
+                st.final(out[i..][0..prk_length]);
                 counter[0] +%= 1;
                 assert(counter[0] != 1);
             }
-            const left = out.len % Hmac.mac_length;
+            const left = out.len % prk_length;
             if (left > 0) {
                 var st = Hmac.init(&prk);
                 if (i != 0) {
-                    st.update(out[i - Hmac.mac_length ..][0..Hmac.mac_length]);
+                    st.update(out[i - prk_length ..][0..prk_length]);
                 }
                 st.update(ctx);
                 st.update(&counter);
-                var tmp: [Hmac.mac_length]u8 = undefined;
-                st.final(tmp[0..Hmac.mac_length]);
+                var tmp: [prk_length]u8 = undefined;
+                st.final(tmp[0..prk_length]);
                 mem.copy(u8, out[i..][0..left], tmp[0..left]);
             }
         }
@@ -64,4 +81,10 @@ test "Hkdf" {
     var out: [42]u8 = undefined;
     kdf.expand(&out, &context, prk);
     try htest.assertEqual("3cb25f25faacd57a90434f64d0362f2a2d2d0a90cf1a5a4c5db02d56ecc4c5bf34007208d5b887185865", &out);
+
+    var hkdf = kdf.extractInit(&salt);
+    hkdf.update(&ikm);
+    var prk2: [kdf.prk_length]u8 = undefined;
+    hkdf.final(&prk2);
+    try htest.assertEqual("077709362c2e32df0ddc3f0dc47bba6390b6c73bb50f9c3122ec844ad7c2b3e5", &prk2);
 }