feat: add mlkem keygen

src/lib.rs

@@ -1,100 +1,5 @@
 #![feature(lazy_cell)]
 
 mod agent;
-
-pub use shared_memory_heap::sharedptr::SharedPtr;
-
-use agent::Agent;
-use std::{
-    ops::Deref, path::Path, sync::{LazyLock, Mutex}
-};
-
-static AGENT: LazyLock<Mutex<Agent>> = LazyLock::new(|| {
-    let agent_path = std::env::var("AGENT_PATH").expect("AGENT_PATH environment variable missing");
-    let agent = unsafe { Agent::new(Path::new(&agent_path)).expect("Agent failed to start") };
-    Mutex::new(agent)
-});
-
-#[derive(Debug)]
-pub struct X25519PrivKey(SharedPtr<8>);
-#[derive(Debug)]
-pub struct X25519PubKey(SharedPtr<32>);
-#[derive(Debug)]
-pub struct X25519SharedKey(SharedPtr<32>);
-
-impl From<&[u8; 32]> for X25519PubKey {
-    fn from(value: &[u8; 32]) -> Self {
-        let mut pk = SharedPtr::<32>::new().unwrap();
-        pk.copy_from_slice(value);
-        X25519PubKey(pk)
-    }
-}
-
-impl From<&[u8; 8]> for X25519PrivKey {
-    fn from(value: &[u8; 8]) -> Self {
-        let mut sk = SharedPtr::<8>::new().unwrap();
-        sk.copy_from_slice(value);
-        X25519PrivKey(sk)
-    }
-}
-
-impl Deref for X25519PrivKey {
-    type Target = [u8; 8];
-
-    fn deref(&self) -> &Self::Target {
-        self.0.deref()
-    }
-}
-
-impl Deref for X25519PubKey {
-    type Target = [u8; 32];
-
-    fn deref(&self) -> &Self::Target {
-        self.0.deref()
-    }
-}
-
-impl Deref for X25519SharedKey {
-    type Target = [u8; 32];
-
-    fn deref(&self) -> &Self::Target {
-        self.0.deref()
-    }
-}
-
-pub fn x25519_keygen() -> (X25519PrivKey, X25519PubKey) {
-    let sk = X25519PrivKey(SharedPtr::<8>::new().unwrap());
-    let pk = X25519PubKey(SharedPtr::<32>::new().unwrap()); 
-
-    let mut agent = AGENT.lock().unwrap();
-
-    unsafe {
-        agent.perform_ipc_call(0, &[sk.0.get_offset(), pk.0.get_offset()]);
-    }
-
-    (sk, pk)
-}
-
-pub fn x22519_pubkey(sk: &X25519PrivKey) -> X25519PubKey {
-    let pk = X25519PubKey(SharedPtr::<32>::new().unwrap());
-
-    let mut agent = AGENT.lock().unwrap();
-
-    unsafe {
-        agent.perform_ipc_call(1, &[sk.0.get_offset(), pk.0.get_offset()]);
-    }
-
-    pk
-}
-
-pub fn x25519(sk: &X25519PrivKey, pk: &X25519PubKey) -> X25519SharedKey {
-    let out = X25519SharedKey(SharedPtr::<32>::new().unwrap());
-
-    let mut agent = AGENT.lock().unwrap();
-
-    unsafe {
-        agent.perform_ipc_call(2, &[out.0.get_offset(), sk.0.get_offset(), pk.0.get_offset()]);
-    }
-
-    out
-}
+pub mod x25519;
+pub mod mlkem;

src/mlkem.rs

@@ -0,0 +1,82 @@
+pub use shared_memory_heap::sharedptr::SharedPtr;
+
+use crate::agent::Agent;
+use std::{
+    ops::Deref, path::Path, sync::{LazyLock, Mutex}
+};
+
+static MLKEMAGENT: LazyLock<Mutex<Agent>> = LazyLock::new(|| {
+    let agent_path = std::env::var("MLKEM_AGENT_PATH").expect("MLKEM_AGENT_PATH environment variable missing");
+    let agent = unsafe { Agent::new(Path::new(&agent_path)).expect("Agent failed to start") };
+    Mutex::new(agent)
+});
+
+#[derive(Debug)]
+pub struct MLKEMPrivKey(SharedPtr<8>);
+#[derive(Debug)]
+pub struct MLKEMPubKey(SharedPtr<1184>);
+#[derive(Debug)]
+pub struct MLKEMCiphertext(SharedPtr<1088>);
+#[derive(Debug)]
+pub struct MLKEMSharedKey(SharedPtr<32>);
+
+impl From<&[u8; 1184]> for MLKEMPubKey {
+    fn from(value: &[u8; 1184]) -> Self {
+        let mut pk = SharedPtr::<1184>::new().unwrap();
+        pk.copy_from_slice(value);
+        MLKEMPubKey(pk)
+    }
+}
+
+impl From<&[u8; 8]> for MLKEMPrivKey {
+    fn from(value: &[u8; 8]) -> Self {
+        let mut sk = SharedPtr::<8>::new().unwrap();
+        sk.copy_from_slice(value);
+        MLKEMPrivKey(sk)
+    }
+}
+
+impl Deref for MLKEMPrivKey {
+    type Target = [u8; 8];
+
+    fn deref(&self) -> &Self::Target {
+        self.0.deref()
+    }
+}
+
+impl Deref for MLKEMPubKey {
+    type Target = [u8; 1184];
+
+    fn deref(&self) -> &Self::Target {
+        self.0.deref()
+    }
+}
+
+impl Deref for MLKEMCiphertext {
+    type Target = [u8; 1088];
+
+    fn deref(&self) -> &Self::Target {
+        self.0.deref()
+    }
+}
+
+impl Deref for MLKEMSharedKey {
+    type Target = [u8; 32];
+
+    fn deref(&self) -> &Self::Target {
+        self.0.deref()
+    }
+}
+
+pub fn mlkem_keygen() -> (MLKEMPrivKey, MLKEMPubKey) {
+    let sk = MLKEMPrivKey(SharedPtr::<8>::new().unwrap());
+    let pk = MLKEMPubKey(SharedPtr::<1184>::new().unwrap());
+
+    let mut agent = MLKEMAGENT.lock().unwrap();
+
+    unsafe {
+        agent.perform_ipc_call(0, &[sk.0.get_offset(), pk.0.get_offset()]);
+    }
+
+    (sk, pk)
+}

src/x25519.rs

@@ -0,0 +1,96 @@
+pub use shared_memory_heap::sharedptr::SharedPtr;
+
+use crate::agent::Agent;
+use std::{
+    ops::Deref, path::Path, sync::{LazyLock, Mutex}
+};
+
+static X25519AGENT: LazyLock<Mutex<Agent>> = LazyLock::new(|| {
+    let agent_path = std::env::var("X25519_AGENT_PATH").expect("X25519_AGENT_PATH environment variable missing");
+    let agent = unsafe { Agent::new(Path::new(&agent_path)).expect("Agent failed to start") };
+    Mutex::new(agent)
+});
+
+#[derive(Debug)]
+pub struct X25519PrivKey(SharedPtr<8>);
+#[derive(Debug)]
+pub struct X25519PubKey(SharedPtr<32>);
+#[derive(Debug)]
+pub struct X25519SharedKey(SharedPtr<32>);
+
+impl From<&[u8; 32]> for X25519PubKey {
+    fn from(value: &[u8; 32]) -> Self {
+        let mut pk = SharedPtr::<32>::new().unwrap();
+        pk.copy_from_slice(value);
+        X25519PubKey(pk)
+    }
+}
+
+impl From<&[u8; 8]> for X25519PrivKey {
+    fn from(value: &[u8; 8]) -> Self {
+        let mut sk = SharedPtr::<8>::new().unwrap();
+        sk.copy_from_slice(value);
+        X25519PrivKey(sk)
+    }
+}
+
+impl Deref for X25519PrivKey {
+    type Target = [u8; 8];
+
+    fn deref(&self) -> &Self::Target {
+        self.0.deref()
+    }
+}
+
+impl Deref for X25519PubKey {
+    type Target = [u8; 32];
+
+    fn deref(&self) -> &Self::Target {
+        self.0.deref()
+    }
+}
+
+impl Deref for X25519SharedKey {
+    type Target = [u8; 32];
+
+    fn deref(&self) -> &Self::Target {
+        self.0.deref()
+    }
+}
+
+pub fn x25519_keygen() -> (X25519PrivKey, X25519PubKey) {
+    let sk = X25519PrivKey(SharedPtr::<8>::new().unwrap());
+    let pk = X25519PubKey(SharedPtr::<32>::new().unwrap()); 
+
+    let mut agent = X25519AGENT.lock().unwrap();
+
+    unsafe {
+        agent.perform_ipc_call(0, &[sk.0.get_offset(), pk.0.get_offset()]);
+    }
+
+    (sk, pk)
+}
+
+pub fn x22519_pubkey(sk: &X25519PrivKey) -> X25519PubKey {
+    let pk = X25519PubKey(SharedPtr::<32>::new().unwrap());
+
+    let mut agent = X25519AGENT.lock().unwrap();
+
+    unsafe {
+        agent.perform_ipc_call(1, &[sk.0.get_offset(), pk.0.get_offset()]);
+    }
+
+    pk
+}
+
+pub fn x25519(sk: &X25519PrivKey, pk: &X25519PubKey) -> X25519SharedKey {
+    let out = X25519SharedKey(SharedPtr::<32>::new().unwrap());
+
+    let mut agent = X25519AGENT.lock().unwrap();
+
+    unsafe {
+        agent.perform_ipc_call(2, &[out.0.get_offset(), sk.0.get_offset(), pk.0.get_offset()]);
+    }
+
+    out
+}