feat: Blake3 Stateful Hash Object implementation (#13)

* feat: Draft Blake3 Stateful Hash Object implementation

* chore: Draft theorem

* chore: size_of_extract theorem

* chore: mimic Rust version of the sponge implementation (#12)

* chore: mimic Rust version of the sponge implementation

* reset Sponge counter when ratcheting

---------

Co-authored-by: Arthur Paulino <[email protected]>

Author: storojs72

Blake3.lean

@@ -1,5 +1,10 @@
 /-! Bindings to the BLAKE3 hashing library. -/
 
+theorem ByteArray.size_of_extract {hash : ByteArray} (hbe : b ≤ e) (h : e ≤ hash.data.size) :
+    (hash.extract b e).size = e - b := by
+  simp [ByteArray.size, ByteArray.extract, ByteArray.copySlice, ByteArray.empty, ByteArray.mkEmpty]
+  rw [Nat.add_comm, Nat.sub_add_cancel hbe, Nat.min_eq_left h]
+
 namespace Blake3
 
 /-- BLAKE3 constants -/
@@ -53,40 +58,97 @@ opaque initDeriveKey (context : @& ByteArray) : Hasher
 @[extern "lean_blake3_hasher_update"]
 opaque update (hasher : Hasher) (input : @& ByteArray) : Hasher
 
+instance {length : USize} : Inhabited { r : ByteArray // r.size = length.toNat } where
+  default := ⟨
+    ⟨⟨List.replicate length.toNat 0⟩⟩,
+    by simp only [ByteArray.size, List.toArray_replicate, Array.size_mkArray]
+  ⟩
+
 /-- Finalize the hasher and write the output to an array of a given length. -/
 @[extern "lean_blake3_hasher_finalize"]
-opaque finalize : (hasher : Hasher) → (length : USize) → ByteArray
+opaque finalize : (hasher : Hasher) → (length : USize) →
+  {r : ByteArray // r.size = length.toNat}
+
+def finalizeWithLength (hasher : Hasher) (length : Nat)
+    (h : length % 2 ^ System.Platform.numBits = length := by native_decide) :
+    { r : ByteArray // r.size = length } :=
+  let ⟨hash, h'⟩ := hasher.finalize length.toUSize
+  have hash_size_eq_len : hash.size = length := by
+    cases System.Platform.numBits_eq with
+    | inl h32 => rw [h32] at h; simp [h', h32]; exact h
+    | inr h64 => rw [h64] at h; simp [h', h64]; exact h
+  ⟨hash, hash_size_eq_len⟩
 
 end Hasher
 
 /-- Hash a ByteArray -/
 def hash (input : ByteArray) : Blake3Hash :=
   let hasher := Hasher.init ()
   let hasher := hasher.update input
-  let output := hasher.finalize BLAKE3_OUT_LEN.toUSize
-  if h : output.size = BLAKE3_OUT_LEN then
-    ⟨output, h⟩
-  else
-    panic! "Incorrect output size"
+  hasher.finalizeWithLength BLAKE3_OUT_LEN
 
 /-- Hash a ByteArray using keyed initializer -/
 def hashKeyed (input : @& ByteArray) (key : @& Blake3Key) : Blake3Hash :=
   let hasher := Hasher.initKeyed key
   let hasher := hasher.update input
-  let output := hasher.finalize BLAKE3_OUT_LEN.toUSize
-  if h : output.size = BLAKE3_OUT_LEN then
-    ⟨output, h⟩
-  else
-    panic! "Incorrect output size"
+  hasher.finalizeWithLength BLAKE3_OUT_LEN
 
 /-- Hash a ByteArray using initializer parameterized by some context -/
 def hashDeriveKey (input context : @& ByteArray) : Blake3Hash :=
   let hasher := Hasher.initDeriveKey context
   let hasher := hasher.update input
-  let output := hasher.finalize BLAKE3_OUT_LEN.toUSize
-  if h : output.size = BLAKE3_OUT_LEN then
-    ⟨output, h⟩
-  else
-    panic! "Incorrect output size"
+  hasher.finalizeWithLength BLAKE3_OUT_LEN
+
+structure Sponge where
+  hasher : Hasher
+  counter : Nat
+
+namespace Sponge
+
+abbrev ABSORB_MAX_BYTES := UInt32.size - 1
+abbrev DEFAULT_REKEYING_STAGE := UInt16.size - 1
+
+def init (label : String) (_h : ¬label.isEmpty := by decide) : Sponge :=
+  ⟨Hasher.initDeriveKey label.toUTF8, 0⟩
+
+def ratchet (sponge : Sponge) : Sponge :=
+  let key := sponge.hasher.finalizeWithLength BLAKE3_KEY_LEN
+  { sponge with hasher := Hasher.initKeyed key, counter := 0 }
+
+def absorb (sponge : Sponge) (bytes : ByteArray)
+    (_h : bytes.size < ABSORB_MAX_BYTES := by norm_cast) : Sponge :=
+  let highCounter := sponge.counter >= DEFAULT_REKEYING_STAGE
+  let sponge := if highCounter then sponge.ratchet else sponge
+  ⟨sponge.hasher.update bytes, sponge.counter + 1⟩
+
+def squeeze (sponge : Sponge) (length : USize)
+    (h_len_bound : 2 * BLAKE3_OUT_LEN + length.toNat < 2 ^ System.Platform.numBits :=
+      by native_decide) :
+    { r : ByteArray // r.size = length.toNat } :=
+  let ⟨tmp, h⟩ := sponge.hasher.finalize (2 * BLAKE3_OUT_LEN.toUSize + length)
+  let b := (2 * BLAKE3_OUT_LEN)
+  let e := (2 * BLAKE3_OUT_LEN + length.toNat)
+  let y := tmp.extract b e
+  have sub_e_b_eq_length : e - b = length.toNat := by
+    simp only [b, e]
+    rw [Nat.add_comm _ length.toNat, Nat.add_sub_cancel]
+  have le_b_e : b ≤ e := by simp only [Nat.le_add_right, e, b]
+  have h_e_bound : e ≤ tmp.size := by
+    simp [e, h]
+    cases System.Platform.numBits_eq with
+    | inl h32 =>
+      have hbound : (2 * BLAKE3_OUT_LEN) + length.toNat < 2 ^ 32 := by
+        rwa [h32, Nat.pow_succ] at h_len_bound
+      rw [h32, BLAKE3_OUT_LEN]
+      simp only [Nat.mod_eq_of_lt hbound, Nat.le_refl]
+    | inr h64 =>
+      have hbound : (2 * BLAKE3_OUT_LEN) + length.toNat < 2 ^ 64 := by
+        rwa [h64, Nat.pow_succ] at h_len_bound
+      rw [h64, BLAKE3_OUT_LEN]
+      simp only [Nat.mod_eq_of_lt hbound, Nat.le_refl]
+  have size_of_extract := ByteArray.size_of_extract le_b_e h_e_bound
+  ⟨y, by simp only [y, size_of_extract, sub_e_b_eq_length]⟩
+
+end Sponge
 
 end Blake3

Blake3Test.lean

@@ -1,7 +1,7 @@
 import Blake3
 
 abbrev input : ByteArray := ⟨#[0]⟩
-abbrev expected_output_regular_hashing : ByteArray := ⟨#[
+abbrev expectedOutputRegularHashing : ByteArray := ⟨#[
    45,  58, 222, 223, 241,  27,  97, 241,
    76, 136, 110,  53, 175, 160,  54, 115,
   109, 205, 135, 167,  77,  39, 181, 193,
@@ -15,7 +15,7 @@ abbrev key : Blake3.Blake3Key := .ofBytes ⟨#[
    213,  14, 159, 189, 82, 166,  91, 107,
     33,  78,  26, 226, 89,  65, 188, 92
 ]⟩
-abbrev expected_output_keyed_hashing: ByteArray := ⟨#[
+abbrev expectedOutputKeyedHashing: ByteArray := ⟨#[
    145, 187, 220, 234, 206, 139, 205, 138,
    220, 103,  35,  65, 199,  96, 210,  18,
    145, 201, 131, 254,  79, 208, 229, 157,
@@ -24,16 +24,37 @@ abbrev expected_output_keyed_hashing: ByteArray := ⟨#[
 
 -- Context (with "bad" randomness)
 abbrev context : ByteArray := ⟨#[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]⟩
-abbrev expected_output_derive_key_hashing: ByteArray := ⟨#[
+abbrev expectedOutputDeriveKeyHashing: ByteArray := ⟨#[
     34,  57,  22,  43, 164, 211,  65, 131,
     90,  55,  55,  92,  68,  90,  63, 136,
      3, 235,  52, 117, 143, 246, 158, 224,
-   178, 170, 234, 211, 148,  21,  97, 183]⟩
+   178, 170, 234, 211, 148,  21,  97, 183
+]⟩
+
+abbrev expectedOutputSponge : ByteArray := ⟨#[
+    21, 150, 121,  38, 177,  21,  33, 148, 213, 127,
+   208, 231, 200,  18,  46, 115, 244, 188, 245, 188,
+   176,  29,  43, 123, 135, 148, 132, 218,  45, 244,
+   127, 103, 178,  82, 145,  67,  43, 106, 204, 137,
+   154,  99, 175,   9, 196, 102, 126,  72,  27,  86
+]⟩
+
+def hashingFails :=
+  (Blake3.hash input).val.data != expectedOutputRegularHashing.data ||
+  (Blake3.hashKeyed input key).val.data != expectedOutputKeyedHashing.data ||
+  (Blake3.hashDeriveKey input context).val.data != expectedOutputDeriveKeyHashing.data
+
+def spongeFails :=
+  let sponge := Blake3.Sponge.init "ix 2025-01-01 16:18:03 content-addressing v1"
+  let sponge := sponge.absorb ⟨#[1]⟩
+  let sponge := sponge.absorb ⟨#[2]⟩
+  let sponge := sponge.absorb ⟨#[3]⟩
+  let sponge := sponge.absorb ⟨#[4, 5]⟩
+  let output := sponge.squeeze 50
+  output.val.data != expectedOutputSponge.data
 
 def main : IO UInt32 := do
   println! s!"BLAKE3 version: {Blake3.version}"
-  if (Blake3.hash input).val.data != expected_output_regular_hashing.data ||
-     (Blake3.hashKeyed input key).val.data != expected_output_keyed_hashing.data ||
-     (Blake3.hashDeriveKey input context).val.data != expected_output_derive_key_hashing.data
+  if hashingFails || spongeFails
     then IO.eprintln "BLAKE3 test failed";    return 1
     else IO.println  "BLAKE3 test succeeded"; return 0