♻️ refractory hypertree, 🐛 fix sha2_35

colinxu2020 · Jul 8, 2024 · d613a18 · d613a18
1 parent 0cbaa32
commit d613a18
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 40 deletions.
diff --git a/slhdsa/lowlevel/hypertree.py b/slhdsa/lowlevel/hypertree.py
@@ -1,38 +1,45 @@
+from dataclasses import dataclass
+
 from slhdsa.lowlevel.xmss import XMSS
 from slhdsa.lowlevel.addresses import Address
 from slhdsa.lowlevel.parameters import Parameter
 from slhdsa.lowlevel.wots import WOTSParameter
 
 
-def sign(msg: bytes, sk_seed: bytes, pk_seed: bytes, tree_idx: int, leaf_idx: int, par: Parameter) -> bytes:
-    address = Address(0, tree_idx, 0)
-    tree = XMSS(par)
-    ht_sign = tmp_sign = tree.sign(msg, sk_seed, leaf_idx, pk_seed, address)
-    root = tree.public_key_from_sign(leaf_idx, tmp_sign, msg, pk_seed, address)
-    for j in range(1, par.d):
-        leaf_idx = tree_idx % (2 ** par.h_m)
-        tree_idx >>= par.h_m
-        address.layer = j
-        address.tree = tree_idx
-        tmp_sign = tree.sign(root, sk_seed, leaf_idx, pk_seed, address)
-        ht_sign += tmp_sign
-        if j < par.d - 1:
-            root = tree.public_key_from_sign(leaf_idx, tmp_sign, root, pk_seed, address)
-    return ht_sign
-
-
-def verify(msg: bytes, ht_sign: bytes, pk_seed: bytes, tree_idx: int, leaf_idx: int, pk_root: bytes, par: Parameter) -> bool:
-    address = Address(0, tree_idx, 0)
-    wots_par = WOTSParameter(par)
-    tmp_sign = ht_sign[:(par.h_m + wots_par.len) * par.n]
-    tree = XMSS(par)
-    node = tree.public_key_from_sign(leaf_idx, tmp_sign, msg, pk_seed, address)
-
-    for j in range(1, par.d):
-        leaf_idx = tree_idx % (2 ** par.h_m)
-        tree_idx >>= par.h_m
-        address.layer = j
-        address.tree = tree_idx
-        tmp_sign = ht_sign[(par.h_m + wots_par.len) * par.n * j:(par.h_m + wots_par.len) * par.n * (j + 1)]
-        node = tree.public_key_from_sign(leaf_idx, tmp_sign, node, pk_seed, address)
-    return node == pk_root
+
+@dataclass
+class HyperTree:
+    parameter: Parameter
+
+    def sign(self, msg: bytes, sk_seed: bytes, pk_seed: bytes, tree_idx: int, leaf_idx: int) -> bytes:
+        address = Address(0, tree_idx, 0)
+        tree = XMSS(self.parameter)
+        ht_sign = tmp_sign = tree.sign(msg, sk_seed, leaf_idx, pk_seed, address)
+        root = tree.public_key_from_sign(leaf_idx, tmp_sign, msg, pk_seed, address)
+        for j in range(1, self.parameter.d):
+            leaf_idx = tree_idx % (2 ** self.parameter.h_m)
+            tree_idx >>= self.parameter.h_m
+            address.layer = j
+            address.tree = tree_idx
+            tmp_sign = tree.sign(root, sk_seed, leaf_idx, pk_seed, address)
+            ht_sign += tmp_sign
+            if j < self.parameter.d - 1:
+                root = tree.public_key_from_sign(leaf_idx, tmp_sign, root, pk_seed, address)
+        return ht_sign
+
+
+    def verify(self, msg: bytes, ht_sign: bytes, pk_seed: bytes, tree_idx: int, leaf_idx: int, pk_root: bytes) -> bool:
+        address = Address(0, tree_idx, 0)
+        wots_par = WOTSParameter(self.parameter)
+        tmp_sign = ht_sign[:(self.parameter.h_m + wots_par.len) * self.parameter.n]
+        tree = XMSS(self.parameter)
+        node = tree.public_key_from_sign(leaf_idx, tmp_sign, msg, pk_seed, address)
+
+        for j in range(1, self.parameter.d):
+            leaf_idx = tree_idx % (2 ** self.parameter.h_m)
+            tree_idx >>= self.parameter.h_m
+            address.layer = j
+            address.tree = tree_idx
+            tmp_sign = ht_sign[(self.parameter.h_m + wots_par.len) * self.parameter.n * j:(self.parameter.h_m + wots_par.len) * self.parameter.n * (j + 1)]
+            node = tree.public_key_from_sign(leaf_idx, tmp_sign, node, pk_seed, address)
+        return node == pk_root
diff --git a/slhdsa/lowlevel/parameters.py b/slhdsa/lowlevel/parameters.py
@@ -31,7 +31,7 @@ def h_msg(r: bytes, pk_seed: bytes, pk_root: bytes, msg: bytes) -> bytes:
         return shake_256(r + pk_seed + pk_root + msg).digest(m)
 
     def prf(pk_seed: bytes, sk_seed: bytes, address: Address) -> bytes:
-        return shake_256(pk_seed + address.to_bytes() + sk_seed).digest(n)
+        return shake_256(pk_seed +  address.to_bytes() + sk_seed).digest(n)
 
     def prf_msg(sk_prf: bytes, opt_rand: bytes, msg: bytes) -> bytes:
         return shake_256(sk_prf + opt_rand + msg).digest(n)
@@ -110,7 +110,7 @@ def h_msg(r: bytes, pk_seed: bytes, pk_root: bytes, msg: bytes) -> bytes:
         return mgf1_sha512(r + pk_seed + sha512(r + pk_seed + pk_root + msg).digest(), m)
 
     def prf(pk_seed: bytes, sk_seed: bytes, address: Address) -> bytes:
-        return trunc(sha256(pk_seed + b"\x00" * (64 - n) + compact_address(address.to_bytes() + sk_seed)).digest(), n)
+        return trunc(sha256(pk_seed + b"\x00" * (64 - n) + compact_address(address.to_bytes()) + sk_seed).digest(), n)
 
     def prf_msg(sk_prf: bytes, opt_rand: bytes, msg: bytes) -> bytes:
         return trunc(hmac_digest(sk_prf, opt_rand + msg, "sha512"), n)

diff --git a/slhdsa/lowlevel/slhdsa.py b/slhdsa/lowlevel/slhdsa.py
@@ -5,7 +5,7 @@
 from slhdsa.lowlevel.xmss import XMSS
 from slhdsa.lowlevel._utils import ceil_div
 from slhdsa.lowlevel.fors import FORS
-from slhdsa.lowlevel.hypertree import sign as ht_sign, verify as ht_verify
+from slhdsa.lowlevel.hypertree import HyperTree
 from slhdsa.lowlevel.wots import WOTSParameter
 
 
@@ -44,8 +44,8 @@ def sign(msg: bytes, secret_key: tuple[bytes, ...], par: Parameter, randomize: b
     fors_sign = fors.sign(md, sk_seed, pk_seed, address)
     sig += fors_sign
     fors_pk = fors.publickey_from_sign(fors_sign, md, pk_seed, address)
-    ht_sign_ = ht_sign(fors_pk, sk_seed, pk_seed, tree_idx, leaf_idx, par)
-    sig += ht_sign_
+    ht_sign = HyperTree(par).sign(fors_pk, sk_seed, pk_seed, tree_idx, leaf_idx)
+    sig += ht_sign
     return sig
 
 
@@ -69,6 +69,5 @@ def verify(msg: bytes, sig: bytes, public_key: tuple[bytes, ...], par: Parameter
     leaf_id %= 2 ** (par.h // par.d)
     address.tree = tree_id
     address.keypair = leaf_id
-    fors = FORS(par)
-    fors_pk = fors.publickey_from_sign(fors_sign, md, pk_seed, address)
-    return ht_verify(fors_pk, ht_sign_, pk_seed, tree_id, leaf_id, pk_root, par)
+    fors_pk = FORS(par).publickey_from_sign(fors_sign, md, pk_seed, address)
+    return HyperTree(par).verify(fors_pk, ht_sign_, pk_seed, tree_id, leaf_id, pk_root)