Files
mlkem_agent/src/agent.jazz

134 lines
3.3 KiB
Plaintext

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;
p = (u64)[sync_mem + (16 + i * 8)];
p += shared_mem;
return p;
}
inline fn extract_ipc_id(reg u64 sync_mem) -> reg u64 {
reg u64 id;
id = (u64)[sync_mem + 8];
return id;
}
inline fn generate_keypair(reg u64 shared_mem sync_mem private_mem, stack u64 key_id) -> stack u64 {
inline int i;
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);
(u64)[idptr] = key_id;
key_id += 1;
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;
// Futex varibles
reg u64 futex_op val timeout uaddr2 val3 woken_up;
// Mmap variables
reg u64 addr len prot flag off;
() = #spill(shared_mem, sync_mem);
private_mem_fd = private_mem_fd;
addr = 0;
len = 3200;
prot = 3; // PROT_READ | PROT_WRITE
flag = 1; // MAP_SHARED;
off = 0;
private_mem = #mmap(addr, len, prot, flag, private_mem_fd, off);
() = #spill(private_mem);
key_id = 0;
while (true) {
while {
() = #unspill(sync_mem);
futex_op = 0;
val = 0;
timeout = 0;
uaddr2 = 0;
val3 = 0;
woken_up = #futex(sync_mem, futex_op, val, timeout, uaddr2, val3);
} (woken_up != 0)
() = #unspill(shared_mem, sync_mem, private_mem);
reg u64 id;
id = extract_ipc_id(sync_mem);
if (id == 0) {
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) {
}
}
}
while {
() = #unspill(sync_mem);
futex_op = 1;
val = 1;
timeout = 0;
uaddr2 = 0;
val3 = 0;
woken_up = #futex(sync_mem, futex_op, val, timeout, uaddr2, val3);
} (woken_up != 1)
}
}