From f6f9625cf2bf0f4f51883951dd023de67fd7715f Mon Sep 17 00:00:00 2001
From: Infendro <infendro@noreply.infendro.com>
Date: Sun, 3 Aug 2025 00:12:34 +0200
Subject: [PATCH] enable passing hash function to HMAC PRNG

---
 .../kotlin/com/infendro/random/HMAC.kt        | 35 ++++++++++---------
 .../kotlin/com/infendro/random/PRNG.kt        | 31 ++++++++++++++--
 2 files changed, 47 insertions(+), 19 deletions(-)

diff --git a/src/commonMain/kotlin/com/infendro/random/HMAC.kt b/src/commonMain/kotlin/com/infendro/random/HMAC.kt
index 6d4209a..0707679 100644
--- a/src/commonMain/kotlin/com/infendro/random/HMAC.kt
+++ b/src/commonMain/kotlin/com/infendro/random/HMAC.kt
@@ -1,26 +1,26 @@
 package com.infendro.random
 
-import com.infendro.bytearray.toInt
-import com.infendro.hash.SHA256
-import com.infendro.mac.HMAC
-import kotlin.random.Random
+import com.infendro.bytearray.addAll
+import com.infendro.hash.HashFunction
 
-class HMAC internal constructor() : Random() {
+class HMAC internal constructor(
+    function: HashFunction,
+) : PRNG() {
     private val random = Default
-    private val hmac = HMAC(SHA256)
+    private val hmac = com.infendro.mac.HMAC(function)
     private var K: ByteArray
     private var V: ByteArray
 
     init {
-        K = ByteArray(32).also {
+        K = ByteArray(function.bytes).also {
             random.nextBytes(it)
         }
-        V = ByteArray(32).also {
+        V = ByteArray(function.bytes).also {
             random.nextBytes(it)
         }
     }
 
-    fun reseed(
+    override fun reseed(
         entropy: ByteArray,
     ) {
         K = hmac.hash(K, V + byteArrayOf(0x00) + entropy)
@@ -29,18 +29,19 @@ class HMAC internal constructor() : Random() {
         V = hmac.hash(K, V)
     }
 
-    override fun nextBits(
+    override fun generate(
         count: Int,
-    ): Int {
-        if (count !in 0..32) throw Exception()
-        if (count == 0) return 0
-
-        val bytes = hmac.hash(K, V)
-        val bits = bytes.take(4).toInt() ushr (32 - count)
+    ): ByteArray {
+        val bytes = buildList {
+            while (size < count) {
+                V = hmac.hash(K, V)
+                addAll(V)
+            }
+        }.take(count).toByteArray()
 
         K = hmac.hash(K, V + byteArrayOf(0x00))
         V = hmac.hash(K, V)
 
-        return bits
+        return bytes
     }
 }
diff --git a/src/commonMain/kotlin/com/infendro/random/PRNG.kt b/src/commonMain/kotlin/com/infendro/random/PRNG.kt
index 8df79d7..f176fa1 100644
--- a/src/commonMain/kotlin/com/infendro/random/PRNG.kt
+++ b/src/commonMain/kotlin/com/infendro/random/PRNG.kt
@@ -1,5 +1,32 @@
 package com.infendro.random
 
-object PRNG {
-    val HMAC = HMAC()
+import com.infendro.bytearray.toInt
+import com.infendro.hash.HashFunction
+import kotlin.random.Random
+
+abstract class PRNG : Random() {
+    abstract fun reseed(
+        entropy: ByteArray,
+    )
+
+    abstract fun generate(
+        count: Int,
+    ): ByteArray
+
+    override fun nextBits(
+        count: Int,
+    ): Int {
+        if (count !in 0..32) throw Exception()
+        if (count == 0) return 0
+
+        return generate(4).toInt() ushr (32 - count)
+    }
+
+    companion object {
+        fun HMAC(
+            function: HashFunction,
+        ): HMAC {
+            return com.infendro.random.HMAC(function)
+        }
+    }
 }