diff --git a/src/agent.jazz b/src/agent.jazz index 15b6b70..b6e8dd7 100644 --- a/src/agent.jazz +++ b/src/agent.jazz @@ -1,5 +1,7 @@ from Jade require "crypto_kem/kyber_kyber768_avx2/kyber_kyber768_avx2.jazz" +param int MLKEM_SK_SIZE = 2400; + inline fn extract_nth_ptr(reg u64 sync_mem shared_mem, inline int i) -> reg u64 { reg u64 p; @@ -27,10 +29,10 @@ inline fn generate_keypair(reg u64 shared_mem sync_mem private_mem, stack u64 ke coins, num_bytes = #randombytes(coins, flag); } (num_bytes != 2*KYBER_SYMBYTES) - skptr = key_id * 2400; + skptr = key_id * MLKEM_SK_SIZE; skptr += private_mem; - for i=0 to 8 { + for i=0 to (2*KYBER_SYMBYTES)/8 { (u64)[skptr + i * 8] = coins[u64 i]; } @@ -46,6 +48,27 @@ inline fn generate_keypair(reg u64 shared_mem sync_mem private_mem, stack u64 ke return key_id; } +inline fn encapsulate(reg u64 shared_mem sync_mem private_mem) { + inline int i; + 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) + + for i=0 to KYBER_SYMBYTES/8 { + (u64)[ssptr + i * 8] = coins[u64 i]; + } + + 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); + + jade_kem_kyber_kyber768_amd64_avx2_enc_derand(ctptr, ssptr, pkptr, ssptr); +} + export fn agent_start(reg u64 shared_mem sync_mem private_mem_fd) { stack u64 key_id; reg u64 private_mem; @@ -87,6 +110,7 @@ export fn agent_start(reg u64 shared_mem sync_mem private_mem_fd) { key_id = generate_keypair(shared_mem, sync_mem, private_mem, key_id); } else { if (id == 1) { + encapsulate(shared_mem, sync_mem, private_mem); } else { if (id == 2) { }