prg: unrolled cSHAKE padding

draft-irtf-cfrg-vdaf.md

@@ -1661,33 +1661,46 @@ def expand_into_vec(Prg,
 ### PrgSha3 {#prg-sha3}
 
 This section describes PrgSha3, a PRG based on the Keccak permutation of SHA-3
-{{FIPS202}}. Keccak is used in the cSHAKE128 mode of operation {{SP800-185}}.
+{{FIPS202}}. Keccak is used in the SHAKE128 mode of operation {{FIPS202}}.
 This Prg is RECOMMENDED for all use cases within VDAFs.
 
 ~~~
 class PrgSha3(Prg):
-    """PRG based on SHA-3 (cSHAKE128)."""
+    """PRG based on SHA-3 (SHAKE128)."""
 
     # Associated parameters
     SEED_SIZE = 16
 
     def __init__(self, seed, dst, binder):
+        # Implementation note: it is required that the length of dst be
+        # at most 255 bytes. This comes from the use of cSHAKE128's padding
+        # scheme (see {{SP800-185}} for details).
         self.l = 0
         self.x = seed + binder
         self.s = dst
 
     def next(self, length: Unsigned) -> Bytes:
         self.l += length
 
-        # Function `cSHAKE128(x, l, n, s)` is as defined in
-        # [SP800-185, Section 3.3].
+        # Function `SHAKE128(x, l)` is as defined in
+        # [FIPS 202, Section 6.2].
         #
         # Implementation note: Rather than re-generate the output
         # stream each time `next()` is invoked, most implementations
         # of SHA-3 will expose an "absorb-then-squeeze" API that
         # allows stateful handling of the stream.
-        stream = cSHAKE128(self.x, self.l, b'', self.s)
+        dst_length_l = (len(self.s).bit_length() + 7)//8
+        dst_length = len(self.s).to_bytes(dst_length_l, byteorder = 'little')
+        dst_length_l = self.__to_bytes__(dst_length_l)
+        pre_pad = self.__to_bytes__(1) + self.__to_bytes__(168)
+        post_pad_l = (168 - ((len(self.s) + len(dst_length)+3) % 168)) % 168
+        post_pad = (0).to_bytes(post_pad_l, byteorder = 'little')
+        dom_sep = pre_pad + dst_length_l + dst_length + self.s + post_pad
+        stream = SHAKE128(dom_sep + self.x, self.l)
         return stream[-length:]
+
+    def __to_bytes__(self, i):
+        return i.to_bytes(1, byteorder = 'little)
 ~~~
 {: title="Definition of PRG PrgSha3."}
 
@@ -1720,9 +1733,19 @@ class PrgFixedKeyAes128(Prg):
         #
         # Implementation note: This step can be cached across PRG
         # evaluations with many different seeds.
-        self.fixed_key = cSHAKE128(binder, 16, b'', dst)
+        dst_length_l = (len(dst).bit_length + 7)//8
+        dst_length = len(dst).to_bytes(dst_length_l, byteorder = 'little')
+        dst_length_l = self.__to_bytes(dst_length_l) #one byte value
+        pre_pad = self.__to_bytes__(1) + self.__to_bytes__(168)
+        post_pad_l = (168 - ((len(dst) + len(dst_length)+3) % 168)) % 168
+        post_pad = (0).to_bytes(post_pad_l, byteorder = 'little')
+        dom_sep = pre_pad + dst_length_l + dst_length + dst + post_pad
+        self.fixed_key = SHAKE128(dom_sep + binder, 16)
         self.seed = seed
 
+    def __to_bytes__(self, i):
+        return i.to_bytes(1, byteorder = 'little)
+
     def next(self, length: Unsigned) -> Bytes:
         offset = self.length_consumed % 16
         new_length = self.length_consumed + length

poc/prg.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 from Cryptodome.Cipher import AES
-from Cryptodome.Hash import CMAC, cSHAKE128
+from Cryptodome.Hash import CMAC, SHAKE128
 from Cryptodome.Util import Counter
 
 from common import (TEST_VECTOR, VERSION, Bytes, Unsigned, concat, format_dst,
@@ -95,7 +95,7 @@ def next(self, length: Unsigned) -> Bytes:
 
 
 class PrgSha3(Prg):
-    """PRG based on SHA-3 (cSHAKE128)."""
+    """PRG based on SHA-3 (SHAKE128)."""
 
     # Associated parameters
     SEED_SIZE = 16
@@ -106,13 +106,24 @@ class PrgSha3(Prg):
     def __init__(self, seed, dst, binder):
         # `dst` is used as the customization string; `seed || binder` is
         # used as the main input string.
-        self.shake = cSHAKE128.new(custom=dst)
+        self.shake = SHAKE128.new()
+        dst_length_l = (len(dst).bit_length() + 7)//8
+        dst_length = len(dst).to_bytes(dst_length_l, byteorder='little')
+        dst_length_l = self.__to_bytes__(dst_length_l)  #one byte value
+        pre_pad = self.__to_bytes__(1) + self.__to_bytes__(168)
+        post_pad_l = (168 - ((len(dst) + len(dst_length)+3) % 168)) % 168
+        post_pad = (0).to_bytes(post_pad_l,byteorder='little')
+        dom_sep = pre_pad + dst_length_l + dst_length + dst + post_pad
+        self.shake.update(dom_sep)
         self.shake.update(seed)
         self.shake.update(binder)
 
     def next(self, length: Unsigned) -> Bytes:
         return self.shake.read(length)
 
+    def __to_bytes__(self, i):
+        return i.to_bytes(1, byteorder='little')
+
 
 class PrgFixedKeyAes128(Prg):
     """
@@ -136,7 +147,15 @@ def __init__(self, seed, dst, binder):
         #
         # Implementation note: This step can be cached across PRG
         # evaluations with many different seeds.
-        shake = cSHAKE128.new(custom=dst)
+        shake = SHAKE128.new()
+        dst_length_l = (len(dst).bit_length() + 7)//8
+        dst_length = len(dst).to_bytes(dst_length_l, byteorder='little')
+        dst_length_l = self.__to_bytes__(dst_length_l)  #one byte value
+        pre_pad = self.__to_bytes__(1) + self.__to_bytes__(168)
+        post_pad_l = (168 - ((len(dst) + len(dst_length)+3) % 168)) % 168
+        post_pad = (0).to_bytes(post_pad_l, byteorder='little')
+        dom_sep = pre_pad + dst_length_l + dst_length + dst + post_pad
+        shake.update(dom_sep)
         shake.update(binder)
         fixed_key = shake.read(16)
         self.cipher = AES.new(fixed_key, AES.MODE_ECB)
@@ -171,6 +190,9 @@ def hash_block(self, block):
         sigma_block = concat([hi, xor(hi, lo)])
         return xor(self.cipher.encrypt(sigma_block), sigma_block)
 
+    def __to_bytes__(self, i):
+        return i.to_bytes(1, byteorder='little')
+
 
 ##
 # TESTS