From 6bd0a265f8312d6d1447ea17372218b3fb2985c2 Mon Sep 17 00:00:00 2001
From: Aaron Kaiser <a_kaiser@posteo.de>
Date: Tue, 27 Aug 2024 17:33:51 +0200
Subject: [PATCH] feat: add decapsulate

---
 src/agent.jazz | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/src/agent.jazz b/src/agent.jazz
index 000925b..96f57d4 100644
--- a/src/agent.jazz
+++ b/src/agent.jazz
@@ -69,6 +69,20 @@ inline fn encapsulate(reg u64 shared_mem sync_mem) {
     jade_kem_kyber_kyber768_amd64_avx2_enc_derand(ctptr, ssptr, pkptr, ssptr);
 }
 
+inline fn decapsulate(reg u64 shared_mem sync_mem private_mem) {
+   reg u64 ssptr ctptr idptr keyid keyptr;
+
+   ssptr = extract_nth_ptr(sync_mem, shared_mem, 0);
+   ctptr = extract_nth_ptr(sync_mem, shared_mem, 1);
+   idptr = extract_nth_ptr(sync_mem, shared_mem, 2);
+
+   keyid = (u64)[idptr];
+   keyid *= MLKEM_SK_SIZE;
+   keyptr = private_mem + keyid;
+
+   jade_kem_kyber_kyber768_amd64_avx2_dec(ssptr, ctptr, keyptr);
+}
+
 export fn agent_start(reg u64 shared_mem sync_mem private_mem_fd) {
     stack u64 key_id;
     reg u64 private_mem;
@@ -116,6 +130,7 @@ export fn agent_start(reg u64 shared_mem sync_mem private_mem_fd) {
                 encapsulate(shared_mem, sync_mem);
             } else {
                 if (id == 2) {
+                    decapsulate(shared_mem, sync_mem, private_mem);
                 }
             }
         }