some refactoring

default.nix

@@ -4,8 +4,8 @@ let
   jasmin-src = fetchFromGitHub {
     owner = "jasmin-lang";
     repo = "jasmin";
-    rev = "4d42c212b924ad7553c2eab49d337fd128ad629b";
-    hash = "sha256-Ve9Eezpvi1wynlqQxO4KDSqw13MDXQmR8NjOS3PHWzg=";
+    rev = "c2adabafb8df0a60e7cf4e56041d3be3aca387a3";
+    hash = "sha256-ySw4eVbuV7suO8EjyaMOTqeNsPq1D6vsorCZlXffTPU=";
   };
 
   jasmin-drv = callPackage "${jasmin-src}/default.nix" { inherit pkgs; };

src/agent.jazz

@@ -1,5 +1,9 @@
 from Jade require "crypto_scalarmult/curve25519_mulx/curve25519_mulx.jazz"
 
+param int SK_SIZE = 32;
+param int SHARED_MEM_SIZE = 1024;
+param int PRIVATE_MEM_SIZE = 3200;
+
 inline fn extract_nth_ptr(reg u64 sync_mem shared_mem, inline int i) -> reg u64 {
     reg u64 p;
 
@@ -17,41 +21,45 @@ inline fn extract_ipc_id(reg u64 sync_mem) -> reg u64 {
     return id;
 }
 
-inline fn generate_keypair(reg u64 shared_mem sync_mem private_mem, stack u64 key_id) -> stack u64 {
+inline fn generate_keypair(reg u64 shared_mem sync_mem private_mem, stack u64 key_id) -> (stack u64, reg u64) {
     inline int i;
-    stack u8[32] key;
-    reg u64 addr idptr pkptr flag num_bytes;
+    stack u8[SK_SIZE] key;
+    reg u64 addr idptr pkptr flag num_bytes return_value;
 
-    flag = 0;
-    key, num_bytes = #randombytes(key, flag);
+    return_value = 0;
+
+    while {
+        flag = 0;
+        key, num_bytes = #randombytes(key, flag);
+    } (num_bytes != SK_SIZE)
     key[0] &= 248;
     key[31] &= 127;
     key[31] |= 64;
 
-    addr = key_id * 32;
+    addr = key_id * SK_SIZE;
     addr += private_mem;
 
-    for i=0 to 4 {
+    for i=0 to SK_SIZE/8 {
         (u64)[addr + i * 8] = key[u64 i];
     }
 
-    idptr = (u64)[sync_mem + 16];
-    idptr += shared_mem;
+    idptr = extract_nth_ptr(sync_mem, shared_mem, 0);
 
     (u64)[idptr] = key_id;
 
-    pkptr = (u64)[sync_mem + 24];
-    pkptr += shared_mem;
+    pkptr = extract_nth_ptr(sync_mem, shared_mem, 1);
 
     jade_scalarmult_curve25519_amd64_mulx_base(pkptr, addr);
 
     key_id += 1; 
 
-    return key_id;
+    return_value = 1;
+    return key_id, return_value;
 }
 
-inline fn calculate_pubkey(reg u64 shared_mem sync_mem private_mem) {
-    reg u64 skptr outptr;
+inline fn calculate_pubkey(reg u64 shared_mem sync_mem private_mem) -> reg u64 {
+    reg u64 skptr outptr return_value;
+    return_value = 0;
 
     skptr = extract_nth_ptr(sync_mem, shared_mem, 0);
     outptr = extract_nth_ptr(sync_mem, shared_mem, 1);
@@ -61,12 +69,17 @@ inline fn calculate_pubkey(reg u64 shared_mem sync_mem private_mem) {
     skptr += private_mem;
 
     jade_scalarmult_curve25519_amd64_mulx_base(outptr, skptr);
+
+    return_value = 1;
+    return return_value;
 }
 
-inline fn x25519(reg u64 shared_mem sync_mem private_mem) {
-    reg u64 outptr idptr pkptr keyptr;
+inline fn x25519(reg u64 shared_mem sync_mem private_mem) -> reg u64 {
+    reg u64 outptr idptr pkptr keyptr return_value;
     reg u64 keyid;
 
+    return_value = 0;
+
     outptr = extract_nth_ptr(sync_mem, shared_mem, 0);
     idptr = extract_nth_ptr(sync_mem, shared_mem, 1);
     pkptr = extract_nth_ptr(sync_mem, shared_mem, 2);
@@ -76,10 +89,13 @@ inline fn x25519(reg u64 shared_mem sync_mem private_mem) {
     keyptr = private_mem + keyid;
 
     jade_scalarmult_curve25519_amd64_mulx(outptr, keyptr, pkptr);
+
+    return_value = 1;
+    return return_value;
 }
 
 export fn agent_start(reg u64 shared_mem sync_mem private_mem_fd) {
-    stack u64 key_id;
+    stack u64 key_id return_value;
     reg u64 private_mem;
     // Futex varibles
     reg u64 futex_op val timeout uaddr2 val3 woken_up;
@@ -90,7 +106,7 @@ export fn agent_start(reg u64 shared_mem sync_mem private_mem_fd) {
 
     private_mem_fd = private_mem_fd;
     addr = 0;
-    len = 3200;
+    len = PRIVATE_MEM_SIZE;
     prot = 3; // PROT_READ | PROT_WRITE
     flag = 1; // MAP_SHARED;
     off = 0;
@@ -100,15 +116,17 @@ export fn agent_start(reg u64 shared_mem sync_mem private_mem_fd) {
 
     key_id = 0;
 
+    #no_termination_check
      while (true) {
-        () = #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);
+        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);
         
@@ -116,24 +134,31 @@ export fn agent_start(reg u64 shared_mem sync_mem private_mem_fd) {
         id = extract_ipc_id(sync_mem);
 
         if (id == 0) {
-            key_id = generate_keypair(shared_mem, sync_mem, private_mem, key_id);
+            if (key_id < 100) {
+                key_id, return_value = generate_keypair(shared_mem, sync_mem, private_mem, key_id);
+            }
         } else {
             if (id == 1) {
-                calculate_pubkey(shared_mem, sync_mem, private_mem);
+                return_value = calculate_pubkey(shared_mem, sync_mem, private_mem);
             } else {
                 if (id == 2) {
-                    x25519(shared_mem, sync_mem, private_mem);
+                    return_value = x25519(shared_mem, sync_mem, private_mem);
                 }
             }
         }
 
         () = #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);
+        (u64)[sync_mem + 16] = return_value;
+
+        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)
     }
 }