From 0966fd5fc7acbf8cfb0124330bc935ff48d7e738 Mon Sep 17 00:00:00 2001
From: Vishal Mhatre <vishal.mhatre@gmail.com>
Date: Mon, 14 Oct 2024 14:11:06 +0530
Subject: [PATCH] [feat] MLDSA-87 driver interface.

This change contains the interface for the MLDSA-87 driver.
---
 drivers/src/array.rs   |   2 +
 drivers/src/lib.rs     |   1 +
 drivers/src/mldsa87.rs | 128 +++++++++++++++++++++++++++++++++++++++++
 3 files changed, 131 insertions(+)
 create mode 100644 drivers/src/mldsa87.rs

diff --git a/drivers/src/array.rs b/drivers/src/array.rs
index 4ad03cb5b5..bbaab1f09c 100644
--- a/drivers/src/array.rs
+++ b/drivers/src/array.rs
@@ -160,6 +160,8 @@ pub type Array4x8 = Array4xN<8, 32>;
 pub type Array4x12 = Array4xN<12, 48>;
 pub type Array4x16 = Array4xN<16, 64>;
 pub type Array4x32 = Array4xN<32, 128>;
+pub type Array4x648 = Array4xN<648, 2592>;
+pub type Array4x1157 = Array4xN<1157, 4628>;
 
 #[cfg(test)]
 mod tests {
diff --git a/drivers/src/lib.rs b/drivers/src/lib.rs
index 23ccd648cd..f3ed62283b 100644
--- a/drivers/src/lib.rs
+++ b/drivers/src/lib.rs
@@ -37,6 +37,7 @@ mod kv_access;
 mod lms;
 mod mailbox;
 pub mod memory_layout;
+mod mldsa87;
 mod okref;
 mod pcr_bank;
 pub mod pcr_log;
diff --git a/drivers/src/mldsa87.rs b/drivers/src/mldsa87.rs
new file mode 100644
index 0000000000..837a5e11e4
--- /dev/null
+++ b/drivers/src/mldsa87.rs
@@ -0,0 +1,128 @@
+/*++
+
+Licensed under the Apache-2.0 license.
+
+File Name:
+
+    mldsa87.rs
+
+Abstract:
+
+    File contains API for MLDSA-87 Cryptography operations
+
+--*/
+#![allow(dead_code)]
+
+use crate::{
+    array::{Array4x1157, Array4x648},
+    Array4x16,
+};
+use crate::{CaliptraResult, KeyReadArgs, Trng};
+
+#[must_use]
+#[repr(u32)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum MlDsa87Result {
+    Success = 0xAAAAAAAA,
+    SigVerifyFailed = 0x55555555,
+}
+
+/// MLDSA-87 Public Key
+pub type MlDsa87PubKey = Array4x648;
+
+/// MLDSA-87 Signature
+pub type MlDsa87Signature = Array4x1157;
+
+/// MLDSA-87 Message (64 Bytes)
+pub type MlDsa87MsgScalar = Array4x16;
+
+/// TEMP: Placeholder for MlDsa87Reg
+pub struct MlDsa87Reg {
+    _priv: (),
+}
+impl MlDsa87Reg {
+    /// # Safety
+    ///
+    /// Caller must ensure that all concurrent use of this
+    /// peripheral in the firmware is done so in a compatible
+    /// way. The simplest way to enforce this is to only call
+    /// this function once.
+    #[inline(always)]
+    pub unsafe fn new() -> Self {
+        Self { _priv: () }
+    }
+}
+/// END - TEMP: Placeholder for MlDsa87Reg
+
+/// MLDSA-87  API
+pub struct MlDsa87 {
+    mldsa87: MlDsa87Reg,
+}
+
+impl MlDsa87 {
+    pub fn new(mldsa87: MlDsa87Reg) -> Self {
+        Self { mldsa87 }
+    }
+
+    /// Generate MLDSA-87 Key Pair
+    ///
+    /// # Arguments
+    ///
+    /// * `seed` - Key Vault slot containing the seed for deterministic MLDSA Key Pair generation.
+    /// * `trng` - TRNG driver instance.
+    ///
+    /// # Returns
+    ///
+    /// * `MlDsa87PubKey` - Generated MLDSA-87 Public Key
+    pub fn key_pair(
+        &mut self,
+        _seed: &KeyReadArgs,
+        _trng: &mut Trng,
+    ) -> CaliptraResult<MlDsa87PubKey> {
+        Ok(MlDsa87PubKey::default())
+    }
+
+    /// Sign the digest with specified private key. To defend against glitching
+    /// attacks that could expose the private key, this function also verifies
+    /// the generated signature.
+    ///
+    /// # Arguments
+    ///
+    /// * `priv_key_in` - Key Vault slot containing the seed for the private key generation.
+    /// * `pub_key` - Public key to verify the signature with.
+    /// * `msg` - Message to sign.
+    /// * `trng` - TRNG driver instance.
+    ///
+    /// # Returns
+    ///
+    /// * `MlDsa87Signature` - Generated signature
+    pub fn sign(
+        &mut self,
+        _priv_key_in: &KeyReadArgs,
+        _pub_key: &MlDsa87PubKey,
+        _msg: &MlDsa87MsgScalar,
+        _trng: &mut Trng,
+    ) -> CaliptraResult<MlDsa87Signature> {
+        Ok(MlDsa87Signature::default())
+    }
+
+    /// Verify the signature with specified public key and message.
+    ///
+    /// # Arguments
+    ///
+    /// * `pub_key` - Public key.
+    /// * `msg` - Message to verify.
+    /// * `signature` - Signature to verify.
+    ///
+    /// # Result
+    ///
+    /// *  `MlDsa87Result` - MlDsa87Result::Success if the signature verification passed else an error code.
+    pub fn verify(
+        &mut self,
+        _pub_key: &MlDsa87PubKey,
+        _msg: &MlDsa87MsgScalar,
+        _signature: &MlDsa87Signature,
+    ) -> CaliptraResult<MlDsa87Result> {
+        Ok(MlDsa87Result::Success)
+    }
+}