From 3a90ce74a5f50ea19e4f170e256cd8cc9030171c Mon Sep 17 00:00:00 2001
From: Yawning Angel <yawning@schwanenlied.me>
Date: Thu, 14 Nov 2024 18:48:20 +0900
Subject: [PATCH] core/crypto/blake2: Add the ability to easily alter digest
 size

---
 core/crypto/blake2b/blake2b.odin | 9 ++++++---
 core/crypto/blake2s/blake2s.odin | 9 ++++++---
 2 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/core/crypto/blake2b/blake2b.odin b/core/crypto/blake2b/blake2b.odin
index 74396b103e6..16b19e349cf 100644
--- a/core/crypto/blake2b/blake2b.odin
+++ b/core/crypto/blake2b/blake2b.odin
@@ -18,7 +18,7 @@ package blake2b
 import "../_blake2"
 
 // DIGEST_SIZE is the BLAKE2b digest size in bytes.
-DIGEST_SIZE :: 64
+DIGEST_SIZE :: _blake2.BLAKE2B_SIZE
 
 // BLOCK_SIZE is the BLAKE2b block size in bytes.
 BLOCK_SIZE :: _blake2.BLAKE2B_BLOCK_SIZE
@@ -27,9 +27,12 @@ BLOCK_SIZE :: _blake2.BLAKE2B_BLOCK_SIZE
 Context :: _blake2.Blake2b_Context
 
 // init initializes a Context with the default BLAKE2b config.
-init :: proc(ctx: ^Context) {
+init :: proc(ctx: ^Context, digest_size := DIGEST_SIZE) {
+	if digest_size > 255 {
+		panic("blake2b: invalid digest size")
+	}
 	cfg: _blake2.Blake2_Config
-	cfg.size = _blake2.BLAKE2B_SIZE
+	cfg.size = u8(digest_size)
 	_blake2.init(ctx, &cfg)
 }
 
diff --git a/core/crypto/blake2s/blake2s.odin b/core/crypto/blake2s/blake2s.odin
index 339ddf0272c..f7d3abe03ba 100644
--- a/core/crypto/blake2s/blake2s.odin
+++ b/core/crypto/blake2s/blake2s.odin
@@ -18,7 +18,7 @@ package blake2s
 import "../_blake2"
 
 // DIGEST_SIZE is the BLAKE2s digest size in bytes.
-DIGEST_SIZE :: 32
+DIGEST_SIZE :: _blake2.BLAKE2S_SIZE
 
 // BLOCK_SIZE is the BLAKE2s block size in bytes.
 BLOCK_SIZE :: _blake2.BLAKE2S_BLOCK_SIZE
@@ -27,9 +27,12 @@ BLOCK_SIZE :: _blake2.BLAKE2S_BLOCK_SIZE
 Context :: _blake2.Blake2s_Context
 
 // init initializes a Context with the default BLAKE2s config.
-init :: proc(ctx: ^Context) {
+init :: proc(ctx: ^Context, digest_size := DIGEST_SIZE) {
+	if digest_size > 255 {
+		panic("blake2s: invalid digest size")
+	}
 	cfg: _blake2.Blake2_Config
-	cfg.size = _blake2.BLAKE2S_SIZE
+	cfg.size = u8(digest_size)
 	_blake2.init(ctx, &cfg)
 }