feat(ssh-key): add ML-DSA public key type

Author: quexten

ssh-key/src/public.rs

@@ -8,6 +8,8 @@ mod dsa;
 mod ecdsa;
 mod ed25519;
 mod key_data;
+#[cfg(feature = "mldsa")]
+mod mldsa;
 #[cfg(feature = "alloc")]
 mod opaque;
 #[cfg(feature = "alloc")]
@@ -27,6 +29,9 @@ pub use self::{
 #[cfg(feature = "ecdsa")]
 pub use self::{ecdsa::EcdsaPublicKey, sk::SkEcdsaSha2NistP256};
 
+#[cfg(feature = "mldsa")]
+pub use self::mldsa::MlDsaPublicKey;
+
 pub(crate) use self::ssh_format::SshFormat;
 
 use crate::{Algorithm, Error, Fingerprint, HashAlg, Result};
@@ -448,6 +453,13 @@ impl From<Ed25519PublicKey> for PublicKey {
     }
 }
 
+#[cfg(feature = "mldsa")]
+impl From<MlDsaPublicKey> for PublicKey {
+    fn from(public_key: MlDsaPublicKey) -> PublicKey {
+        KeyData::from(public_key).into()
+    }
+}
+
 #[cfg(feature = "alloc")]
 impl From<RsaPublicKey> for PublicKey {
     fn from(public_key: RsaPublicKey) -> PublicKey {

ssh-key/src/public/key_data.rs

@@ -14,6 +14,9 @@ use {
 #[cfg(feature = "ecdsa")]
 use super::{EcdsaPublicKey, SkEcdsaSha2NistP256};
 
+#[cfg(feature = "mldsa")]
+use super::MlDsaPublicKey;
+
 /// Public key data.
 #[derive(Clone, Debug, Eq, Hash, PartialEq, PartialOrd, Ord)]
 #[non_exhaustive]
@@ -33,6 +36,10 @@ pub enum KeyData {
     #[cfg(feature = "alloc")]
     Rsa(RsaPublicKey),
 
+    /// ML-DSA public key data.
+    #[cfg(feature = "mldsa")]
+    MlDsa(MlDsaPublicKey),
+
     /// Security Key (FIDO/U2F) using ECDSA/NIST P-256 as specified in [PROTOCOL.u2f].
     ///
     /// [PROTOCOL.u2f]: https://cvsweb.openbsd.org/src/usr.bin/ssh/PROTOCOL.u2f?annotate=HEAD
@@ -69,6 +76,8 @@ impl KeyData {
             Self::Ed25519(_) => Algorithm::Ed25519,
             #[cfg(feature = "alloc")]
             Self::Rsa(_) => Algorithm::Rsa { hash: None },
+            #[cfg(feature = "mldsa")]
+            Self::MlDsa(key) => key.algorithm(),
             #[cfg(feature = "ecdsa")]
             Self::SkEcdsaSha2NistP256(_) => Algorithm::SkEcdsaSha2NistP256,
             Self::SkEd25519(_) => Algorithm::SkEd25519,
@@ -127,6 +136,16 @@ impl KeyData {
         }
     }
 
+    /// Get ML-DSA public key if this key is the correct type.
+    #[cfg(feature = "mldsa")]
+    #[must_use]
+    pub fn mldsa(&self) -> Option<&MlDsaPublicKey> {
+        match self {
+            Self::MlDsa(key) => Some(key),
+            _ => None,
+        }
+    }
+
     /// Get FIDO/U2F ECDSA/NIST P-256 public key if this key is the correct type.
     #[cfg(feature = "ecdsa")]
     #[must_use]
@@ -203,6 +222,13 @@ impl KeyData {
         matches!(self, Self::Rsa(_))
     }
 
+    /// Is this key an ML-DSA key?
+    #[cfg(feature = "mldsa")]
+    #[must_use]
+    pub fn is_mldsa(&self) -> bool {
+        matches!(self, Self::MlDsa(_))
+    }
+
     /// Is this key a FIDO/U2F ECDSA/NIST P-256 key?
     #[cfg(feature = "ecdsa")]
     #[must_use]
@@ -248,6 +274,10 @@ impl KeyData {
             Algorithm::Ed25519 => Ed25519PublicKey::decode(reader).map(Self::Ed25519),
             #[cfg(feature = "alloc")]
             Algorithm::Rsa { .. } => RsaPublicKey::decode(reader).map(Self::Rsa),
+            #[cfg(feature = "mldsa")]
+            Algorithm::MlDsa { params } => {
+                MlDsaPublicKey::decode_as(reader, params).map(Self::MlDsa)
+            }
             #[cfg(feature = "ecdsa")]
             Algorithm::SkEcdsaSha2NistP256 => {
                 SkEcdsaSha2NistP256::decode(reader).map(Self::SkEcdsaSha2NistP256)
@@ -280,6 +310,8 @@ impl KeyData {
             Self::Ed25519(key) => key.encoded_len(),
             #[cfg(feature = "alloc")]
             Self::Rsa(key) => key.encoded_len(),
+            #[cfg(feature = "mldsa")]
+            Self::MlDsa(key) => key.encoded_len(),
             #[cfg(feature = "ecdsa")]
             Self::SkEcdsaSha2NistP256(sk) => sk.encoded_len(),
             Self::SkEd25519(sk) => sk.encoded_len(),
@@ -300,6 +332,8 @@ impl KeyData {
             Self::Ed25519(key) => key.encode(writer),
             #[cfg(feature = "alloc")]
             Self::Rsa(key) => key.encode(writer),
+            #[cfg(feature = "mldsa")]
+            Self::MlDsa(key) => key.encode(writer),
             #[cfg(feature = "ecdsa")]
             Self::SkEcdsaSha2NistP256(sk) => sk.encode(writer),
             Self::SkEd25519(sk) => sk.encode(writer),
@@ -380,6 +414,13 @@ impl From<RsaPublicKey> for KeyData {
     }
 }
 
+#[cfg(feature = "mldsa")]
+impl From<MlDsaPublicKey> for KeyData {
+    fn from(public_key: MlDsaPublicKey) -> KeyData {
+        Self::MlDsa(public_key)
+    }
+}
+
 #[cfg(feature = "ecdsa")]
 impl From<SkEcdsaSha2NistP256> for KeyData {
     fn from(public_key: SkEcdsaSha2NistP256) -> KeyData {

ssh-key/src/public/mldsa.rs

@@ -0,0 +1,109 @@
+//! ML-DSA public keys.
+//!
+//! Module-Lattice-Based Digital Signature Algorithm (ML-DSA) as specified in [FIPS204].
+//!
+//! [FIPS204]: https://csrc.nist.gov/pubs/fips/204/final
+
+use crate::{Algorithm, MlDsaParams, Result};
+use alloc::boxed::Box;
+use encoding::{Encode, Reader, Writer};
+use ml_dsa::{EncodedVerifyingKey, MlDsa44, MlDsa65, MlDsa87};
+
+/// ML-DSA public key.
+///
+///
+/// [draft-sfluhrer-ssh-mldsa]: https://datatracker.ietf.org/doc/draft-sfluhrer-ssh-mldsa/
+/// [FIPS204]: https://csrc.nist.gov/pubs/fips/204/final
+#[derive(Clone, Debug, Eq, Hash, PartialEq, PartialOrd, Ord)]
+pub enum MlDsaPublicKey {
+    /// ML-DSA-44 public key.
+    MlDsa44(Box<EncodedVerifyingKey<MlDsa44>>),
+
+    /// ML-DSA-65 public key.
+    MlDsa65(Box<EncodedVerifyingKey<MlDsa65>>),
+
+    /// ML-DSA-87 public key.
+    MlDsa87(Box<EncodedVerifyingKey<MlDsa87>>),
+}
+
+impl MlDsaPublicKey {
+    /// Maximum size in bytes of a FIPS 204 public key across all parameter sets
+    /// (i.e. the ML-DSA-87 public key size).
+    const MAX_SIZE: usize = 2592;
+
+    /// Create a new ML-DSA public key from raw FIPS 204 public key bytes for the
+    /// given parameter set.
+    ///
+    /// # Errors
+    /// Returns [`Error::Encoding`] with [`encoding::Error::Length`] if the key
+    /// length does not match the parameter set's public key size.
+    pub fn new(params: MlDsaParams, key: impl AsRef<[u8]>) -> Result<Self> {
+        let key = key.as_ref();
+
+        Ok(match params {
+            MlDsaParams::MlDsa44 => Self::MlDsa44(Box::new(
+                EncodedVerifyingKey::<MlDsa44>::try_from(key).map_err(|_| encoding::Error::Length)?,
+            )),
+            MlDsaParams::MlDsa65 => Self::MlDsa65(Box::new(
+                EncodedVerifyingKey::<MlDsa65>::try_from(key).map_err(|_| encoding::Error::Length)?,
+            )),
+            MlDsaParams::MlDsa87 => Self::MlDsa87(Box::new(
+                EncodedVerifyingKey::<MlDsa87>::try_from(key).map_err(|_| encoding::Error::Length)?,
+            )),
+        })
+    }
+
+    /// Get the [`MlDsaParams`] parameter set for this public key.
+    #[must_use]
+    pub fn params(&self) -> MlDsaParams {
+        match self {
+            Self::MlDsa44(_) => MlDsaParams::MlDsa44,
+            Self::MlDsa65(_) => MlDsaParams::MlDsa65,
+            Self::MlDsa87(_) => MlDsaParams::MlDsa87,
+        }
+    }
+
+    /// Get the [`Algorithm`] for this public key.
+    #[must_use]
+    pub fn algorithm(&self) -> Algorithm {
+        Algorithm::MlDsa {
+            params: self.params(),
+        }
+    }
+
+    /// Get the raw FIPS 204 public key bytes.
+    #[must_use]
+    pub fn as_bytes(&self) -> &[u8] {
+        match self {
+            Self::MlDsa44(key) => key.as_slice(),
+            Self::MlDsa65(key) => key.as_slice(),
+            Self::MlDsa87(key) => key.as_slice(),
+        }
+    }
+
+    /// Decode an ML-DSA public key for the given parameter set.
+    ///
+    /// The parameter set is not encoded in the key body; it is taken from the
+    /// SSH algorithm identifier (see [`MlDsaPublicKey`]).
+    pub(crate) fn decode_as(reader: &mut impl Reader, params: MlDsaParams) -> Result<Self> {
+        let mut buf = [0u8; Self::MAX_SIZE];
+        let key = reader.read_byten(&mut buf)?;
+        Self::new(params, key)
+    }
+}
+
+impl AsRef<[u8]> for MlDsaPublicKey {
+    fn as_ref(&self) -> &[u8] {
+        self.as_bytes()
+    }
+}
+
+impl Encode for MlDsaPublicKey {
+    fn encoded_len(&self) -> encoding::Result<usize> {
+        self.as_bytes().encoded_len()
+    }
+
+    fn encode(&self, writer: &mut impl Writer) -> encoding::Result<()> {
+        self.as_bytes().encode(writer)
+    }
+}