Files
agent_lib/src/mlkem.rs
2024-08-27 17:36:53 +02:00

109 lines
2.7 KiB
Rust

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 keyfile_path = std::env::var("MLKEM_KEYFILE").expect("MLKEM_KEYFILE environment variable missing");
let agent = unsafe { Agent::new(Path::new(&agent_path), Path::new(&keyfile_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)
}
pub fn mlkem_encap(pk: &MLKEMPubKey) -> (MLKEMCiphertext, MLKEMSharedKey) {
let ct = MLKEMCiphertext(SharedPtr::<1088>::new().unwrap());
let ss = MLKEMSharedKey(SharedPtr::<32>::new().unwrap());
let mut agent = MLKEMAGENT.lock().unwrap();
unsafe {
agent.perform_ipc_call(1, &[ct.0.get_offset(), ss.0.get_offset(), pk.0.get_offset()]);
}
(ct, ss)
}
pub fn mlkem_decap(ct: &MLKEMCiphertext, sk: &MLKEMPrivKey) -> MLKEMSharedKey {
let ss = MLKEMSharedKey(SharedPtr::<32>::new().unwrap());
let mut agent = MLKEMAGENT.lock().unwrap();
unsafe {
agent.perform_ipc_call(1, &[ss.0.get_offset(), ct.0.get_offset(), sk.0.get_offset()]);
}
ss
}