feat: patch libjade and use randomized functions

default.nix

@@ -26,12 +26,14 @@ let
       hash = "sha256-IWLMWExvdZ++V3a9zX0a9xjSDgOA98h3sSmLOj8hKPE=";
     };
 
-    buildPhase = ''
+    patches = [ 
+      ./patches/randombytes.patch
+      ./patches/returnaddresskind.patch
+    ];
+
+    installPhase = ''
       mkdir -p $out
-      cp -r $src/libjade $out
+      cp -r libjade $out
-      chmod -R u+w $out
-      sed -i '/#randombytes/c\' $out/libjade/crypto_kem/kyber_kyber768_avx2/kyber_kyber768_avx2.jazz
-      sed -i 's/#\[returnaddress="stack"\]/#\[returnaddress="mmx"\]/g' $out/libjade/crypto_kem/kyber_kyber768_avx2/kyber_kyber768_avx2.jazz
     '';
   };
 in

patches/randombytes.patch

@@ -0,0 +1,30 @@
+diff --git a/libjade/crypto_kem/kyber_kyber768_avx2/kyber_kyber768_avx2.jazz b/libjade/crypto_kem/kyber_kyber768_avx2/kyber_kyber768_avx2.jazz
+index 8323647..6ddd805 100644
+--- a/libjade/crypto_kem/kyber_kyber768_avx2/kyber_kyber768_avx2.jazz
++++ b/libjade/crypto_kem/kyber_kyber768_avx2/kyber_kyber768_avx2.jazz
+@@ -4769,7 +4769,11 @@ export fn jade_kem_kyber_kyber768_amd64_avx2_keypair(reg u64 public_key secret_k
+ 
+   public_key = public_key;
+   secret_key = secret_key;
+-  stack_coins = #randombytes(stack_coins);
++  reg u64 num_bytes flag;
++  while {
++    flag = 0;
++    stack_coins, num_bytes = #randombytes(stack_coins, flag);
++  } (num_bytes != 2*KYBER_SYMBYTES)
+   __crypto_kem_keypair_derand_jazz(public_key, secret_key, stack_coins);
+   ?{}, r = #set0();
+   return r;
+@@ -4797,7 +4801,11 @@ export fn jade_kem_kyber_kyber768_amd64_avx2_enc(reg u64 ciphertext shared_secre
+   ciphertext = ciphertext;
+   shared_secret = shared_secret;
+   public_key = public_key;
+-  stack_coins = #randombytes(stack_coins);
++  reg u64 num_bytes flag;
++  while {
++    flag = 0;
++    stack_coins, num_bytes = #randombytes(stack_coins, flag);
++  } (num_bytes != KYBER_SYMBYTES)
+   __crypto_kem_enc_derand_jazz(ciphertext, shared_secret, public_key, stack_coins);
+   ?{}, r = #set0();
+   return r;

patches/returnaddresskind.patch

@@ -0,0 +1,130 @@
+diff --git a/libjade/crypto_kem/kyber_kyber768_avx2/kyber_kyber768_avx2.jazz b/libjade/crypto_kem/kyber_kyber768_avx2/kyber_kyber768_avx2.jazz
+index 6ddd805..f4ed0e9 100644
+--- a/libjade/crypto_kem/kyber_kyber768_avx2/kyber_kyber768_avx2.jazz
++++ b/libjade/crypto_kem/kyber_kyber768_avx2/kyber_kyber768_avx2.jazz
+@@ -504,7 +504,7 @@ param int SHAKE256_RATE = 136;
+ param int SHA3_256_RATE = 136;
+ param int SHA3_512_RATE = 72;
+ 
+-#[returnaddress="stack"]
++#[returnaddress="mmx"]
+ fn _sha3_256(reg ptr u8[32] out, reg u64 in inlen) -> reg ptr u8[32]
+ {
+   reg u256[7] state;
+@@ -537,7 +537,7 @@ fn _sha3_256(reg ptr u8[32] out, reg u64 in inlen) -> reg ptr u8[32]
+ }
+ 
+ 
+-#[returnaddress="stack"]
++#[returnaddress="mmx"]
+ fn _sha3_256_32(reg ptr u8[32] out, reg ptr u8[KYBER_SYMBYTES] in) -> reg ptr u8[32]
+ {
+   reg u256[7] state;
+@@ -585,7 +585,7 @@ fn _sha3_256_32(reg ptr u8[32] out, reg ptr u8[KYBER_SYMBYTES] in) -> reg ptr u8
+ }
+ 
+ 
+-#[returnaddress="stack"]
++#[returnaddress="mmx"]
+ fn _shake256_64(reg u64 out outlen, reg const ptr u8[64] in)
+ {
+   reg u256[7] state;
+@@ -625,7 +625,7 @@ fn _shake256_64(reg u64 out outlen, reg const ptr u8[64] in)
+ }
+ 
+ 
+-#[returnaddress="stack"]
++#[returnaddress="mmx"]
+ fn _shake256_128_33(reg ptr u8[128] out, reg const ptr u8[33] in) -> stack u8[128]
+ {
+   reg u256[7] state;
+@@ -679,7 +679,7 @@ fn _shake256_128_33(reg ptr u8[128] out, reg const ptr u8[33] in) -> stack u8[12
+   return out;
+ }
+ 
+-#[returnaddress="stack"]
++#[returnaddress="mmx"]
+ fn _sha3_512_64(reg ptr u8[64] out, reg const ptr u8[64] in) -> stack u8[64]
+ {
+   reg u256[7] state;
+@@ -728,7 +728,7 @@ fn _sha3_512_64(reg ptr u8[64] out, reg const ptr u8[64] in) -> stack u8[64]
+   return out;
+ }
+ 
+-#[returnaddress="stack"]
++#[returnaddress="mmx"]
+ fn _sha3_512_32(reg ptr u8[64] out, reg const ptr u8[32] in) -> stack u8[64]
+ {
+   reg u256[7] state;
+@@ -1237,7 +1237,7 @@ inline fn __keccakf1600_4x_avx2(reg ptr u256[25] a) -> reg ptr u256[25]
+   return a;
+ }
+ 
+-#[returnaddress="stack"]
++#[returnaddress="mmx"]
+ fn _keccakf1600_4x_avx2(reg ptr u256[25] a) -> reg ptr u256[25]
+ {
+   a = __keccakf1600_4x_avx2(a);
+@@ -1256,7 +1256,7 @@ inline fn _keccakf1600_4x_avx2_(reg ptr u256[25] a) -> reg ptr u256[25]
+ 
+ u256 SHAKE_SEP = (4u64)[0x8000000000000000, 0x8000000000000000, 0x8000000000000000, 0x8000000000000000];
+ 
+-#[returnaddress="stack"]
++#[returnaddress="mmx"]
+ fn _shake128_absorb4x_34(reg ptr u256[25] s, reg ptr u8[34] m0 m1 m2 m3) -> reg ptr u256[25]
+ {
+ 	inline int i;
+@@ -1307,7 +1307,7 @@ fn _shake128_absorb4x_34(reg ptr u256[25] s, reg ptr u8[34] m0 m1 m2 m3) -> reg
+ }
+ 
+ 
+-#[returnaddress="stack"]
++#[returnaddress="mmx"]
+ fn _shake256_absorb4x_33(reg ptr u256[25] s, reg ptr u8[33] m0 m1 m2 m3) -> reg ptr u256[25]
+ {
+ 	inline int i;
+@@ -2459,7 +2459,7 @@ fn __poly_cbd_eta2(reg ptr u16[KYBER_N] rp, reg ptr u8[KYBER_ETA2*KYBER_N/4] buf
+ }
+ 
+ 
+-#[returnaddress="stack"]
++#[returnaddress="mmx"]
+ fn _poly_getnoise(reg ptr u16[KYBER_N] rp, reg ptr u8[KYBER_SYMBYTES] seed, reg u8 nonce) -> reg ptr u16[KYBER_N]
+ {
+   inline int i;
+@@ -2549,7 +2549,7 @@ fn __shake256_squeezenblocks4x(reg ptr u256[25] state, reg ptr u8[NOISE_NBLOCKS
+   return state, buf0, buf1, buf2, buf3;
+ }
+ 
+-#[returnaddress="stack"]
++#[returnaddress="mmx"]
+ fn _poly_getnoise_eta1_4x(
+   reg ptr u16[KYBER_N] r0 r1 r2 r3,
+   reg ptr u8[KYBER_SYMBYTES] seed,
+@@ -2593,7 +2593,7 @@ fn _poly_getnoise_eta1_4x(
+   return r0, r1, r2, r3;
+ }
+ 
+-#[returnaddress="stack"]
++#[returnaddress="mmx"]
+ fn _poly_getnoise_eta1122_4x(
+   reg ptr u16[KYBER_N] r0 r1 r2 r3,
+   reg ptr u8[KYBER_SYMBYTES] seed,
+@@ -4629,7 +4629,7 @@ fn __cmov(reg ptr u8[KYBER_SYMBYTES] dst, reg u64 src cnd) -> reg ptr u8[KYBER_S
+ // Note2: due to the integration of hakyber implementations into libjade, this file is no longer
+ //        required by kyber768/amd64/ref/
+ 
+-#[returnaddress="stack"]
++#[returnaddress="mmx"]
+ fn __crypto_kem_keypair_derand_jazz(reg u64 pkp, reg u64 skp, reg ptr u8[2*KYBER_SYMBYTES] coins)
+ {
+   stack u8[32] h_pk;
+@@ -4668,7 +4668,7 @@ fn __crypto_kem_keypair_derand_jazz(reg u64 pkp, reg u64 skp, reg ptr u8[2*KYBER
+   __fromstack32u8(skp, coins[32:32]);
+ }
+ 
+-#[returnaddress="stack"]
++#[returnaddress="mmx"]
+ fn __crypto_kem_enc_derand_jazz(reg u64 ctp, reg u64 shkp, reg u64 pkp, reg ptr u8[KYBER_SYMBYTES] coins)
+ {
+   stack u8[KYBER_SYMBYTES * 2] buf kr;

src/agent.jazz

@@ -20,26 +20,15 @@ inline fn extract_ipc_id(reg u64 sync_mem) -> reg u64 {
 }
 
 inline fn generate_keypair(reg u64 shared_mem sync_mem private_mem, stack u64 key_id) -> stack u64 {
-    inline int i;
+    reg u64 skptr idptr pkptr;
-    stack u8[2*KYBER_SYMBYTES] coins;
-    reg u64 skptr idptr pkptr flag num_bytes;
-
-    while {
-        flag = 0;
-        coins, num_bytes = #randombytes(coins, flag);
-    } (num_bytes != 2*KYBER_SYMBYTES)
 
     skptr = key_id * MLKEM_SK_SIZE;
     skptr += private_mem;
 
-    for i=0 to (2*KYBER_SYMBYTES)/8  {
-        (u64)[skptr + i * 8] = coins[u64 i];
-    }
-
     idptr = extract_nth_ptr(sync_mem, shared_mem, 0);
     pkptr = extract_nth_ptr(sync_mem, shared_mem, 1);
 
-    jade_kem_kyber_kyber768_amd64_avx2_keypair_derand(pkptr, skptr, skptr);
+    jade_kem_kyber_kyber768_amd64_avx2_keypair(pkptr, skptr);
 
     (u64)[idptr] = key_id;
 
@@ -49,24 +38,13 @@ inline fn generate_keypair(reg u64 shared_mem sync_mem private_mem, stack u64 ke
 }
 
 inline fn encapsulate(reg u64 shared_mem sync_mem) {
-    inline int i;
+    reg u64 ssptr ctptr pkptr;
-    reg u64 ssptr ctptr pkptr flag num_bytes;
-    stack u8[KYBER_SYMBYTES] coins;
-
-    while {
-        flag = 0;
-        coins, num_bytes = #randombytes(coins, flag);
-    } (num_bytes != KYBER_SYMBYTES)
 
     ctptr = extract_nth_ptr(sync_mem, shared_mem, 0);
     ssptr = extract_nth_ptr(sync_mem, shared_mem, 1);
     pkptr = extract_nth_ptr(sync_mem, shared_mem, 2);
 
-    for i=0 to KYBER_SYMBYTES/8 {
+    jade_kem_kyber_kyber768_amd64_avx2_enc(ctptr, ssptr, pkptr);
-        (u64)[ssptr + i * 8] = coins[u64 i];
-    }
-
-    jade_kem_kyber_kyber768_amd64_avx2_enc_derand(ctptr, ssptr, pkptr, ssptr);
 }
 
 inline fn decapsulate(reg u64 shared_mem sync_mem private_mem) {