Implement RFC-5869 HMAC-based Extract-and-Expand Key Derivation Function (HKDF)

The HKDF is initialized ("extracted") from a (non-secret) salt and a secret key.
From this, any number of secret keys can be derived ("expanded") deterministically.

When multiple keys are to be derived from the same initial keying/salting material,
each separate key should use a distinct "context". This ensures that there exists
a domain separation between the keys. Using the same context as another key on a
HKDF initialized with the same salt+key results in the exact same derived key
material as that key.

This implementation only offers HMAC-SHA256-based key derivation.

Tested with all HMAC-SHA256 test vectors in RFC-5869, with added edge case tests.

Analogous to BouncyCastle's `HKDFBytesGenerator`, but with a simpler API that
tries to be very explicit in its operation, as well as fully thread safe due to
not storing intermediate calculations in member fields.

2 files changed, 435 insertions, 0 deletions

diff --git a/security-utils/src/main/java/com/yahoo/security/HKDF.java b/security-utils/src/main/java/com/yahoo/security/HKDF.java
new file mode 100644
index 00000000000..3aff89d71c2
--- /dev/null
+++ b/security-utils/src/main/java/com/yahoo/security/HKDF.java
@@ -0,0 +1,221 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.security;
+
+import javax.crypto.Mac;
+import javax.crypto.spec.SecretKeySpec;
+import java.nio.ByteBuffer;
+import java.security.InvalidKeyException;
+import java.security.NoSuchAlgorithmException;
+import java.util.Objects;
+
+/**
+ * Implementation of RFC-5869 HMAC-based Extract-and-Expand Key Derivation Function (HKDF).
+ *
+ * <p>The HKDF is initialized ("extracted") from a (non-secret) salt and a secret key.
+ * From this, any number of secret keys can be derived ("expanded") deterministically.</p>
+ *
+ * <p>When multiple keys are to be derived from the same initial keying/salting material,
+ * each separate key should use a distinct "context" in the {@link #expand(int, byte[])}
+ * call. This ensures that there exists a domain separation between the keys.
+ * Using the same context as another key on a HKDF initialized with the same salt+key
+ * results in the exact same derived key material as that key.</p>
+ *
+ * <p>This implementation only offers HMAC-SHA256-based key derivation.</p>
+ *
+ * @see <a href="https://tools.ietf.org/html/rfc5869">RFC-5869</a>
+ * @see <a href="https://en.wikipedia.org/wiki/HKDF">HKDF on Wikipedia</a>
+ *
+ * @author vekterli
+ */
+public class HKDF {
+
+    private static final int    HASH_LEN        = 32; // Fixed output size of HMAC-SHA256. Corresponds to HashLen in the spec
+    private static final byte[] EMPTY_BYTES     = new byte[0];
+    private static final byte[] ALL_ZEROS_SALT  = new byte[HASH_LEN];
+    public static final int     MAX_OUTPUT_SIZE = 255 * HASH_LEN;
+
+    private final byte[] pseudoRandomKey; // Corresponds to "PRK" in spec
+
+    private HKDF(byte[] pseudoRandomKey) {
+        this.pseudoRandomKey = pseudoRandomKey;
+    }
+
+    private static Mac createHmacSha256() {
+        try {
+            return Mac.getInstance("HmacSHA256");
+        } catch (NoSuchAlgorithmException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    private static SecretKeySpec hmacKeyFrom(byte[] rawKey) {
+        return new SecretKeySpec(rawKey, "HmacSHA256");
+    }
+
+    private static Mac createKeyedHmacSha256(byte[] rawKey) {
+        var hmac = createHmacSha256();
+        try {
+            hmac.init(hmacKeyFrom(rawKey));
+        } catch (InvalidKeyException e) {
+            throw new RuntimeException(e);
+        }
+        return hmac;
+    }
+
+    private static void validateExtractionParams(byte[] salt, byte[] ikm) {
+        Objects.requireNonNull(salt);
+        Objects.requireNonNull(ikm);
+        if (ikm.length == 0) {
+            throw new IllegalArgumentException("HKDF extraction IKM array can not be empty");
+        }
+        if (salt.length == 0) {
+            throw new IllegalArgumentException("HKDF extraction salt array can not be empty");
+        }
+    }
+
+    /**
+     * Creates and returns a new HKDF instance extracted from the given salt and key.
+     *
+     * <p>Both the salt and input key value may be of arbitrary size, but it is recommended
+     * to have both be at least 16 bytes in size.</p>
+     *
+     * @param salt a non-secret salt value. Should ideally be high entropy and functionally
+     *             "as if random". May not be empty, use {@link #unsaltedExtractedFrom(byte[])}
+     *             if unsalted extraction is desired (though this is not recommended).
+     * @param ikm secret initial Input Keying Material value.
+     * @return a new HDFK instance ready for deriving keys based on the salt and IKM.
+     */
+    public static HKDF extractedFrom(byte[] salt, byte[] ikm) {
+        validateExtractionParams(salt, ikm);
+        /*
+          RFC-5869, Step 2.2, Extract:
+
+          HKDF-Extract(salt, IKM) -> PRK
+
+          Options:
+             Hash     a hash function; HashLen denotes the length of the
+                      hash function output in octets
+
+          Inputs:
+             salt     optional salt value (a non-secret random value);
+                      if not provided, it is set to a string of HashLen zeros.
+             IKM      input keying material
+
+          Output:
+             PRK      a pseudorandom key (of HashLen octets)
+
+          The output PRK is calculated as follows:
+
+          PRK = HMAC-Hash(salt, IKM)
+        */
+        var mac = createKeyedHmacSha256(salt); // Note: HDFK is initially keyed on the salt, _not_ on ikm!
+        mac.update(ikm);
+        return new HKDF(/*PRK = */ mac.doFinal());
+    }
+
+    /**
+     * Creates and returns a new <em>unsalted</em> HKDF instance extracted from the given key.
+     *
+     * <p>Prefer using the salted {@link #extractedFrom(byte[], byte[])} method if possible.</p>
+     *
+     * @param ikm secret initial Input Keying Material value.
+     * @return a new HDFK instance ready for deriving keys based on the IKM and an all-zero salt.
+     */
+    public static HKDF unsaltedExtractedFrom(byte[] ikm) {
+        return extractedFrom(ALL_ZEROS_SALT, ikm);
+    }
+
+    /**
+     * Derives a key with a given number of bytes for a particular context. The returned
+     * key is always deterministic for a given unique context and a HKDF initialized with
+     * a specific salt+IKM pair.
+     *
+     * <p>Thread safety: multiple threads can safely call <code>expand()</code> simultaneously
+     * on the same HKDF object.</p>
+     *
+     * @param wantedBytes Positive number of output bytes. Must be less than or equal to {@link #MAX_OUTPUT_SIZE}
+     * @param context Context for key derivation. Derivation is deterministic for a given context.
+     *                Note: this maps to the "info" field in RFC-5869.
+     * @return A byte buffer of size wantedBytes filled with derived key material
+     */
+    public byte[] expand(int wantedBytes, byte[] context) {
+        Objects.requireNonNull(context);
+        verifyWantedBytesWithinBounds(wantedBytes);
+        return expandImpl(wantedBytes, context);
+    }
+
+    /**
+     * Derives a key with a given number of bytes. The returned key is always deterministic
+     * for a HKDF initialized with a specific salt+IKM pair.
+     *
+     * <p>If more than one key is to be derived, use {@link #expand(int, byte[])}</p>
+     *
+     * <p>Thread safety: multiple threads can safely call <code>expand()</code> simultaneously
+     * on the same HKDF object.</p>
+     *
+     * @param wantedBytes Positive number of output bytes. Must be less than or equal to {@link #MAX_OUTPUT_SIZE}
+     * @return A byte buffer of size wantedBytes filled with derived key material
+     */
+    public byte[] expand(int wantedBytes) {
+        return expand(wantedBytes, EMPTY_BYTES);
+    }
+
+    private void verifyWantedBytesWithinBounds(int wantedBytes) {
+        if (wantedBytes <= 0) {
+            throw new IllegalArgumentException("Requested negative or zero number of HKDF output bytes");
+        }
+        if (wantedBytes > MAX_OUTPUT_SIZE) {
+            throw new IllegalArgumentException("Too many requested HKDF output bytes (max %d, got %d)"
+                                               .formatted(MAX_OUTPUT_SIZE, wantedBytes));
+        }
+    }
+
+    private byte[] expandImpl(int wantedBytes, byte[] context) {
+        /*
+           RFC-5869, Step 2.3, Expand:
+
+           HKDF-Expand(PRK, info, L) -> OKM
+
+           Inputs:
+              PRK      a pseudorandom key of at least HashLen octets
+                       (usually, the output from the extract step)
+              info     optional context and application specific information
+                       (can be a zero-length string)
+              L        length of output keying material in octets
+                       (<= 255*HashLen)
+
+           Output:
+              OKM      output keying material (of L octets)
+
+           The output OKM is calculated as follows:
+
+           N = ceil(L/HashLen)
+           T = T(1) | T(2) | T(3) | ... | T(N)
+           OKM = first L octets of T
+
+           where:
+           T(0) = empty string (zero length)
+           T(1) = HMAC-Hash(PRK, T(0) | info | 0x01)
+           T(2) = HMAC-Hash(PRK, T(1) | info | 0x02)
+           T(3) = HMAC-Hash(PRK, T(2) | info | 0x03)
+           ...
+        */
+        var prkHmac = createKeyedHmacSha256(pseudoRandomKey);
+        int blocks  = (wantedBytes / HASH_LEN) + ((wantedBytes % HASH_LEN) != 0 ? 1 : 0); // N
+        var buffer  = ByteBuffer.allocate(blocks * HASH_LEN); // T
+        byte[] lastBlock = EMPTY_BYTES; // initially T(0)
+        for (int i = 0; i < blocks; ++i) {
+            prkHmac.update(lastBlock);
+            prkHmac.update(context);
+            prkHmac.update((byte)(i + 1)); // Number of blocks shall never exceed 255
+            // HMAC instance can be reused across doFinal() calls; resets back to initially keyed state.
+            lastBlock = prkHmac.doFinal();
+            buffer.put(lastBlock);
+        }
+        buffer.flip();
+        byte[] outputKeyingMaterial = new byte[wantedBytes]; // OKM
+        buffer.get(outputKeyingMaterial);
+        return outputKeyingMaterial;
+    }
+
+}
diff --git a/security-utils/src/test/java/com/yahoo/security/HKDFTest.java b/security-utils/src/test/java/com/yahoo/security/HKDFTest.java
new file mode 100644
index 00000000000..c2d8b99e9a7
--- /dev/null
+++ b/security-utils/src/test/java/com/yahoo/security/HKDFTest.java
@@ -0,0 +1,214 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.security;
+
+import org.bouncycastle.util.encoders.Hex;
+import org.junit.jupiter.api.Test;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+/**
+ * HKDF tests that ensure that the output of our own implementation matches the test
+ * vectors given in <a href="https://tools.ietf.org/html/rfc5869">RFC-5869</a>.
+ *
+ * We don't expose the internal PRK (pseudo-random key) value of the HKDF itself,
+ * so we don't test it explicitly. The actual OKM (output keying material) inherently
+ * depends on it, so its correctness is verified transitively.
+ *
+ * @author vekterli
+ */
+public class HKDFTest {
+
+    private static byte[] fromHex(String hex) {
+        return Hex.decode(hex);
+    }
+
+    private static String toHex(byte[] bytes) {
+        return Hex.toHexString(bytes);
+    }
+
+    /*
+       A.1.  Test Case 1
+
+       Basic test case with SHA-256
+
+       Hash = SHA-256
+       IKM  = 0x0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b (22 octets)
+       salt = 0x000102030405060708090a0b0c (13 octets)
+       info = 0xf0f1f2f3f4f5f6f7f8f9 (10 octets)
+       L    = 42
+
+       PRK  = 0x077709362c2e32df0ddc3f0dc47bba63
+              90b6c73bb50f9c3122ec844ad7c2b3e5 (32 octets)
+       OKM  = 0x3cb25f25faacd57a90434f64d0362f2a
+              2d2d0a90cf1a5a4c5db02d56ecc4c5bf
+              34007208d5b887185865 (42 octets)
+     */
+    @Test
+    void rfc_5869_test_vector_case_1() {
+        var ikm  = fromHex("0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b");
+        var salt = fromHex("000102030405060708090a0b0c");
+        var info = fromHex("f0f1f2f3f4f5f6f7f8f9");
+
+        var hkdf = HKDF.extractedFrom(salt, ikm);
+        var okm = hkdf.expand(42, info);
+        assertEquals(toHex(okm),
+                     "3cb25f25faacd57a90434f64d0362f2a" +
+                     "2d2d0a90cf1a5a4c5db02d56ecc4c5bf" +
+                     "34007208d5b887185865");
+    }
+
+    @Test
+    void rfc_5869_test_vector_case_1_block_boundary_edge_cases() {
+        var ikm  = fromHex("0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b");
+        var salt = fromHex("000102030405060708090a0b0c");
+        var info = fromHex("f0f1f2f3f4f5f6f7f8f9");
+
+        var hkdf = HKDF.extractedFrom(salt, ikm);
+        var okm = hkdf.expand(31, info); // One less than block size
+        assertEquals(toHex(okm),
+                     "3cb25f25faacd57a90434f64d0362f2a" +
+                     "2d2d0a90cf1a5a4c5db02d56ecc4c5");
+
+        okm = hkdf.expand(32, info); // Exactly equal to block size
+        assertEquals(toHex(okm),
+                     "3cb25f25faacd57a90434f64d0362f2a" +
+                     "2d2d0a90cf1a5a4c5db02d56ecc4c5bf");
+
+        okm = hkdf.expand(33, info); // One more than block size
+        assertEquals(toHex(okm),
+                     "3cb25f25faacd57a90434f64d0362f2a" +
+                     "2d2d0a90cf1a5a4c5db02d56ecc4c5bf" +
+                     "34");
+    }
+
+    /*
+       A.2.  Test Case 2
+
+       Test with SHA-256 and longer inputs/outputs
+
+       Hash = SHA-256
+       IKM  = 0x000102030405060708090a0b0c0d0e0f
+              101112131415161718191a1b1c1d1e1f
+              202122232425262728292a2b2c2d2e2f
+              303132333435363738393a3b3c3d3e3f
+              404142434445464748494a4b4c4d4e4f (80 octets)
+       salt = 0x606162636465666768696a6b6c6d6e6f
+              707172737475767778797a7b7c7d7e7f
+              808182838485868788898a8b8c8d8e8f
+              909192939495969798999a9b9c9d9e9f
+              a0a1a2a3a4a5a6a7a8a9aaabacadaeaf (80 octets)
+       info = 0xb0b1b2b3b4b5b6b7b8b9babbbcbdbebf
+              c0c1c2c3c4c5c6c7c8c9cacbcccdcecf
+              d0d1d2d3d4d5d6d7d8d9dadbdcdddedf
+              e0e1e2e3e4e5e6e7e8e9eaebecedeeef
+              f0f1f2f3f4f5f6f7f8f9fafbfcfdfeff (80 octets)
+       L    = 82
+
+       PRK  = 0x06a6b88c5853361a06104c9ceb35b45c
+              ef760014904671014a193f40c15fc244 (32 octets)
+       OKM  = 0xb11e398dc80327a1c8e7f78c596a4934
+              4f012eda2d4efad8a050cc4c19afa97c
+              59045a99cac7827271cb41c65e590e09
+              da3275600c2f09b8367793a9aca3db71
+              cc30c58179ec3e87c14c01d5c1f3434f
+              1d87 (82 octets)
+     */
+    @Test
+    void rfc_5869_test_vector_case_2() {
+        var ikm  = fromHex("000102030405060708090a0b0c0d0e0f" +
+                           "101112131415161718191a1b1c1d1e1f" +
+                           "202122232425262728292a2b2c2d2e2f" +
+                           "303132333435363738393a3b3c3d3e3f" +
+                           "404142434445464748494a4b4c4d4e4f");
+        var salt = fromHex("606162636465666768696a6b6c6d6e6f" +
+                           "707172737475767778797a7b7c7d7e7f" +
+                           "808182838485868788898a8b8c8d8e8f" +
+                           "909192939495969798999a9b9c9d9e9f" +
+                           "a0a1a2a3a4a5a6a7a8a9aaabacadaeaf");
+        var info = fromHex("b0b1b2b3b4b5b6b7b8b9babbbcbdbebf" +
+                           "c0c1c2c3c4c5c6c7c8c9cacbcccdcecf" +
+                           "d0d1d2d3d4d5d6d7d8d9dadbdcdddedf" +
+                           "e0e1e2e3e4e5e6e7e8e9eaebecedeeef" +
+                           "f0f1f2f3f4f5f6f7f8f9fafbfcfdfeff");
+
+        var hkdf = HKDF.extractedFrom(salt, ikm);
+        var okm = hkdf.expand(82, info);
+        assertEquals(toHex(okm),
+                     "b11e398dc80327a1c8e7f78c596a4934" +
+                     "4f012eda2d4efad8a050cc4c19afa97c" +
+                     "59045a99cac7827271cb41c65e590e09" +
+                     "da3275600c2f09b8367793a9aca3db71" +
+                     "cc30c58179ec3e87c14c01d5c1f3434f" +
+                     "1d87");
+    }
+
+    /*
+       A.3.  Test Case 3
+
+       Test with SHA-256 and zero-length salt/info
+
+       Hash = SHA-256
+       IKM  = 0x0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b (22 octets)
+       salt = (0 octets)
+       info = (0 octets)
+       L    = 42
+
+       PRK  = 0x19ef24a32c717b167f33a91d6f648bdf
+              96596776afdb6377ac434c1c293ccb04 (32 octets)
+       OKM  = 0x8da4e775a563c18f715f802a063c5a31
+              b8a11f5c5ee1879ec3454e5f3c738d2d
+              9d201395faa4b61a96c8 (42 octets)
+     */
+    @Test
+    void rfc_5869_test_vector_case_3() {
+        var ikm  = fromHex("0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b");
+        var info = new byte[0];
+
+        // We don't allow empty salt to the salted factory function, so this is equivalent.
+        var hkdf = HKDF.unsaltedExtractedFrom(ikm);
+        var okm = hkdf.expand(42, info);
+        var expectedKey = "8da4e775a563c18f715f802a063c5a31" +
+                          "b8a11f5c5ee1879ec3454e5f3c738d2d" +
+                          "9d201395faa4b61a96c8";
+        assertEquals(toHex(okm), expectedKey);
+
+        // expand() without explicit context should return as if an empty context array was passed
+        okm = hkdf.expand(42);
+        assertEquals(toHex(okm), expectedKey);
+    }
+
+    @Test
+    void requested_key_size_is_bounded_and_checked() {
+        var ikm  = fromHex("0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b");
+        var salt = fromHex("000102030405060708090a0b0c");
+        var hkdf = HKDF.extractedFrom(salt, ikm);
+
+        assertThrows(IllegalArgumentException.class, () -> hkdf.expand(-1)); // Can't request negative output size
+
+        assertThrows(IllegalArgumentException.class, () -> hkdf.expand(0)); // Need at least 1 key byte
+
+        assertThrows(IllegalArgumentException.class, () -> hkdf.expand(HKDF.MAX_OUTPUT_SIZE + 1)); // 1 too large
+
+        // However, exactly max should work (though a strange size to request in practice)
+        var okm = hkdf.expand(HKDF.MAX_OUTPUT_SIZE);
+        assertEquals(okm.length, HKDF.MAX_OUTPUT_SIZE);
+    }
+
+    @Test
+    void missing_salt_to_salted_factory_function_throws_exception() {
+        var ikm  = fromHex("0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b");
+        assertThrows(NullPointerException.class, () -> HKDF.extractedFrom(null, ikm));
+        assertThrows(IllegalArgumentException.class, () -> HKDF.extractedFrom(new byte[0], ikm));
+    }
+
+    @Test
+    void ikm_can_not_be_null_or_empty() {
+        var salt = fromHex("000102030405060708090a0b0c");
+        assertThrows(NullPointerException.class, () -> HKDF.extractedFrom(salt, null));
+        assertThrows(IllegalArgumentException.class, () -> HKDF.extractedFrom(salt, new byte[0]));
+        assertThrows(NullPointerException.class, () -> HKDF.unsaltedExtractedFrom(null));
+        assertThrows(IllegalArgumentException.class, () -> HKDF.unsaltedExtractedFrom(new byte[0]));
+    }
+
+}