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

Author: quexten

ssh-key/src/public.rs

@@ -9,6 +9,8 @@ mod ecdsa;
 mod ed25519;
 mod key_data;
 #[cfg(feature = "alloc")]
+mod mldsa;
+#[cfg(feature = "alloc")]
 mod opaque;
 #[cfg(feature = "alloc")]
 mod rsa;
@@ -20,6 +22,7 @@ pub use self::{ed25519::Ed25519PublicKey, key_data::KeyData, sk::SkEd25519};
 #[cfg(feature = "alloc")]
 pub use self::{
     dsa::DsaPublicKey,
+    mldsa::MlDsaPublicKey,
     opaque::{OpaquePublicKey, OpaquePublicKeyBytes},
     rsa::RsaPublicKey,
 };
@@ -448,6 +451,13 @@ impl From<Ed25519PublicKey> for PublicKey {
     }
 }
 
+#[cfg(feature = "alloc")]
+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

@@ -6,7 +6,7 @@ use encoding::{CheckedSum, Decode, Encode, Reader, Writer};
 
 #[cfg(feature = "alloc")]
 use {
-    super::{DsaPublicKey, OpaquePublicKey, RsaPublicKey},
+    super::{DsaPublicKey, MlDsaPublicKey, OpaquePublicKey, RsaPublicKey},
     crate::Certificate,
     alloc::boxed::Box,
 };
@@ -33,6 +33,10 @@ pub enum KeyData {
     #[cfg(feature = "alloc")]
     Rsa(RsaPublicKey),
 
+    /// ML-DSA public key data.
+    #[cfg(feature = "alloc")]
+    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 +73,8 @@ impl KeyData {
             Self::Ed25519(_) => Algorithm::Ed25519,
             #[cfg(feature = "alloc")]
             Self::Rsa(_) => Algorithm::Rsa { hash: None },
+            #[cfg(feature = "alloc")]
+            Self::MlDsa(key) => key.algorithm(),
             #[cfg(feature = "ecdsa")]
             Self::SkEcdsaSha2NistP256(_) => Algorithm::SkEcdsaSha2NistP256,
             Self::SkEd25519(_) => Algorithm::SkEd25519,
@@ -127,6 +133,16 @@ impl KeyData {
         }
     }
 
+    /// Get ML-DSA public key if this key is the correct type.
+    #[cfg(feature = "alloc")]
+    #[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 +219,13 @@ impl KeyData {
         matches!(self, Self::Rsa(_))
     }
 
+    /// Is this key an ML-DSA key?
+    #[cfg(feature = "alloc")]
+    #[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 +271,10 @@ impl KeyData {
             Algorithm::Ed25519 => Ed25519PublicKey::decode(reader).map(Self::Ed25519),
             #[cfg(feature = "alloc")]
             Algorithm::Rsa { .. } => RsaPublicKey::decode(reader).map(Self::Rsa),
+            #[cfg(feature = "alloc")]
+            Algorithm::MlDsa { params } => {
+                MlDsaPublicKey::decode_as(reader, params).map(Self::MlDsa)
+            }
             #[cfg(feature = "ecdsa")]
             Algorithm::SkEcdsaSha2NistP256 => {
                 SkEcdsaSha2NistP256::decode(reader).map(Self::SkEcdsaSha2NistP256)
@@ -280,6 +307,8 @@ impl KeyData {
             Self::Ed25519(key) => key.encoded_len(),
             #[cfg(feature = "alloc")]
             Self::Rsa(key) => key.encoded_len(),
+            #[cfg(feature = "alloc")]
+            Self::MlDsa(key) => key.encoded_len(),
             #[cfg(feature = "ecdsa")]
             Self::SkEcdsaSha2NistP256(sk) => sk.encoded_len(),
             Self::SkEd25519(sk) => sk.encoded_len(),
@@ -300,6 +329,8 @@ impl KeyData {
             Self::Ed25519(key) => key.encode(writer),
             #[cfg(feature = "alloc")]
             Self::Rsa(key) => key.encode(writer),
+            #[cfg(feature = "alloc")]
+            Self::MlDsa(key) => key.encode(writer),
             #[cfg(feature = "ecdsa")]
             Self::SkEcdsaSha2NistP256(sk) => sk.encode(writer),
             Self::SkEd25519(sk) => sk.encode(writer),
@@ -380,6 +411,13 @@ impl From<RsaPublicKey> for KeyData {
     }
 }
 
+#[cfg(feature = "alloc")]
+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,93 @@
+//! 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::vec::Vec;
+use encoding::{Decode, Encode, Reader, Writer};
+
+/// ML-DSA public key.
+///
+/// SSH encodings for ML-DSA public keys are described in
+/// [draft-sfluhrer-ssh-mldsa]:
+///
+/// Here, 'key' is the public key as described in [FIPS204].
+///
+/// This type represents the `key` portion of the encoding together with the
+/// [`MlDsaParams`] parameter set which is derived from the algorithm name.
+///
+/// [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 struct MlDsaPublicKey {
+    /// ML-DSA parameter set.
+    params: MlDsaParams,
+
+    /// Raw FIPS 204 public key bytes.
+    /// https://sfluhrer.github.io/ssh-mldsa/draft-sfluhrer-ssh-mldsa.html#name-public-key-format
+    key: Vec<u8>,
+}
+
+impl MlDsaPublicKey {
+    /// Create a new ML-DSA public key from raw public key bytes.
+    ///
+    /// # Errors
+    /// Returns an error if the key length does not match the parameter set's
+    /// public key size.
+    pub fn new(params: MlDsaParams, key: impl Into<Vec<u8>>) -> Result<Self> {
+        let key = key.into();
+
+        if key.len() != params.public_key_size() {
+            return Err(encoding::Error::Length.into());
+        }
+
+        Ok(Self { params, key })
+    }
+
+    /// Get the [`MlDsaParams`] parameter set for this public key.
+    #[must_use]
+    pub fn params(&self) -> MlDsaParams {
+        self.params
+    }
+
+    /// 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] {
+        &self.key
+    }
+
+    /// 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 key = Vec::decode(reader)?;
+        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.key.encoded_len()
+    }
+
+    fn encode(&self, writer: &mut impl Writer) -> encoding::Result<()> {
+        self.key.encode(writer)
+    }
+}