From 9a244b5c9cb1d12e22512c4483871e5c69834769 Mon Sep 17 00:00:00 2001
From: arobsn <87387688+arobsn@users.noreply.github.com>
Date: Wed, 25 Dec 2024 11:14:13 -0300
Subject: [PATCH] add `blake2b` function

---
 .changeset/few-chefs-sleep.md                 |  5 +++++
 packages/crypto/src/hashes.spec.ts            | 21 +++++++++++++++++--
 packages/crypto/src/hashes.ts                 | 20 +++++++++++++++---
 .../src/serializers/transactionSerializer.ts  | 17 +++++++--------
 4 files changed, 49 insertions(+), 14 deletions(-)
 create mode 100644 .changeset/few-chefs-sleep.md

diff --git a/.changeset/few-chefs-sleep.md b/.changeset/few-chefs-sleep.md
new file mode 100644
index 00000000..f2a8ade9
--- /dev/null
+++ b/.changeset/few-chefs-sleep.md
@@ -0,0 +1,5 @@
+---
+"@fleet-sdk/crypto": patch
+---
+
+Add `blake2b` function
diff --git a/packages/crypto/src/hashes.spec.ts b/packages/crypto/src/hashes.spec.ts
index e5f5f005..b2b47e47 100644
--- a/packages/crypto/src/hashes.spec.ts
+++ b/packages/crypto/src/hashes.spec.ts
@@ -1,14 +1,31 @@
 import { describe, expect, it } from "vitest";
 import { hex, utf8 } from "./coders";
-import { blake2b256, sha256 } from "./hashes";
+import { blake2b, blake2b256, sha256 } from "./hashes";
 
 describe("Hashes smoke tests", () => {
   it("Should hash message using BLAKE2b256", () => {
-    expect(blake2b256(utf8.decode("blake2b256"))).to.be.deep.equal(
+    const msg = utf8.decode("blake2b256");
+    expect(blake2b256(msg)).to.be.deep.equal(
       hex.decode("eb95e6932cedac15db722fcdb0cfd21437f94690339a716251fad2f89842ea8b")
     );
   });
 
+  it("Should hash message using BLAKE2b with parameters", () => {
+    const xpk =
+      "0488b21e000000000000000000b345a673afdeb85091c35d02083035f6e0ca284b1846223b23b566c4070a0cec02a3ad1969b60e85426791b75eccf038e6105c3afab8167af7eb6b73e709b81882";
+
+    expect(
+      hex.encode(
+        blake2b(utf8.decode(xpk), {
+          dkLen: 64,
+          personalization: "wallets checksum"
+        })
+      )
+    ).to.be.equal(
+      "5d33031ea3bbba9d3332559b1dafd8612683092f535273a4c15ffa103ffa3fc11f7b6992f5a034b3c8dd30f6f103b24e500c44ba4cff2e5c7f6e3e2eb124cd32"
+    );
+  });
+
   it("Should have the same result regardless input format", () => {
     const byte = Uint8Array.from([0xde, 0xad, 0xbe, 0xef]);
     const hex = "deadbeef";
diff --git a/packages/crypto/src/hashes.ts b/packages/crypto/src/hashes.ts
index cbe752c8..b96c7f94 100644
--- a/packages/crypto/src/hashes.ts
+++ b/packages/crypto/src/hashes.ts
@@ -1,14 +1,28 @@
-import { blake2b } from "@noble/hashes/blake2b";
+import { blake2b as _blake2b } from "@noble/hashes/blake2b";
 import { sha256 as _sha256 } from "@noble/hashes/sha256";
 import { hex } from "./coders";
 import type { ByteInput } from "./types";
 
+export type Blake2b256Options = {
+  key?: ByteInput;
+  salt?: ByteInput;
+  personalization?: ByteInput;
+};
+
+export type Blake2bOptions = Blake2b256Options & {
+  dkLen?: number;
+};
+
 export function ensureBytes(input: ByteInput): Uint8Array {
   return typeof input === "string" ? hex.decode(input) : input;
 }
 
-export function blake2b256(message: ByteInput): Uint8Array {
-  return blake2b(ensureBytes(message), { dkLen: 32 });
+export function blake2b(message: ByteInput, options?: Blake2bOptions): Uint8Array {
+  return _blake2b(ensureBytes(message), options);
+}
+
+export function blake2b256(message: ByteInput, options?: Blake2b256Options): Uint8Array {
+  return blake2b(ensureBytes(message), { dkLen: 32, ...options });
 }
 
 export function sha256(message: ByteInput): Uint8Array {
diff --git a/packages/serializer/src/serializers/transactionSerializer.ts b/packages/serializer/src/serializers/transactionSerializer.ts
index 515ce555..41eb49e9 100644
--- a/packages/serializer/src/serializers/transactionSerializer.ts
+++ b/packages/serializer/src/serializers/transactionSerializer.ts
@@ -1,7 +1,6 @@
 import {
   type Amount,
   type BoxCandidate,
-  type ContextExtension,
   type DataInput,
   isDefined,
   type UnsignedInput
@@ -46,24 +45,24 @@ function writeInput(writer: SigmaByteWriter, input: UnsignedInput): void {
   writeExtension(writer, input.extension);
 }
 
-function writeExtension(writer: SigmaByteWriter, extension: ContextExtension): void {
+function writeExtension(
+  writer: SigmaByteWriter,
+  extension: Record<string, string | undefined>
+): void {
   const keys = Object.keys(extension);
   let length = 0;
 
   for (const key of keys) {
-    const ext = extension[key as unknown as keyof ContextExtension];
-    if (isDefined(ext)) {
-      length++;
-    }
+    if (isDefined(extension[key])) length++;
   }
 
   writer.writeVLQ(length);
   if (length === 0) return;
 
   for (const key of keys) {
-    const ext = extension[key as unknown as keyof ContextExtension];
-    if (isDefined(ext)) {
-      writer.writeVLQ(Number(key)).writeHex(ext);
+    const val = extension[key];
+    if (isDefined(val)) {
+      writer.writeVLQ(Number(key)).writeHex(val);
     }
   }
 }