Merge pull request #68 from wolfSSL/devin/support_mldsa

Support ML-DSA

Author: dgarske

docs/asymmetric.rst

@@ -112,4 +112,54 @@ ML-KEM
 >>> 
 >>> ss_recv = mlkem_priv.decapsulate(ct)
 >>> ss_send == ss_recv
+True
+
+ML-DSA
+------
+
+.. autoclass:: MlDsaType
+    :show-inheritance:
+
+.. autoclass:: MlDsaPublic
+    :private-members:
+    :members:
+    :inherited-members:
+    
+.. autoclass:: MlDsaPrivate
+    :members:
+    :inherited-members:
+
+**Example:**
+
+>>> ######## Simple Usage
+>>> from wolfcrypt.ciphers import MlDsaType, MlDsaPrivate, MlDsaPublic
+>>> 
+>>> mldsa_type = MlDsaType.ML_DSA_44
+>>> 
+>>> mldsa_priv = MlDsaPrivate.make_key(mldsa_type)
+>>> pub_key = mldsa_priv.encode_pub_key()
+>>> 
+>>> mldsa_pub = MlDsaPublic(mldsa_type)
+>>> mldsa_pub.decode_key(pub_key)
+>>> 
+>>> msg = b"This is an example message"
+>>> 
+>>> sig = mldsa_priv.sign(msg)
+>>> mldsa_pub.verify(sig, msg)
+True
+>>> 
+>>> ######## Export and Import Keys
+>>> exported_key_pair = mldsa_priv.encode_priv_key(), mldsa_priv.encode_pub_key()
+>>> exported_pub_key = mldsa_pub.encode_key()
+>>> exported_key_pair[1] == exported_pub_key
+True
+>>> 
+>>> mldsa_priv2 = MlDsaPrivate(mldsa_type)
+>>> mldsa_priv2.decode_key(exported_key_pair[0], exported_key_pair[1])
+>>> 
+>>> mldsa_pub2 = MlDsaPublic(mldsa_type)
+>>> mldsa_pub2.decode_key(exported_pub_key)
+>>> 
+>>> sig2 = mldsa_priv2.sign(msg)
+>>> mldsa_pub2.verify(sig2, msg)
 True

scripts/build_ffi.py

@@ -235,6 +235,9 @@ def make_flags(prefix, fips):
         # ML-KEM
         flags.append("--enable-kyber")
 
+        # ML-DSA
+        flags.append("--enable-dilithium")
+
         # disabling other configs enabled by default
         flags.append("--disable-oldtls")
         flags.append("--disable-oldnames")
@@ -371,6 +374,7 @@ def get_features(local_wolfssl, features):
     features["AESGCM_STREAM"] = 1 if '#define WOLFSSL_AESGCM_STREAM' in defines else 0
     features["RSA_PSS"] = 1 if '#define WC_RSA_PSS' in defines else 0
     features["CHACHA20_POLY1305"] = 1 if '#define HAVE_CHACHA' and '#define HAVE_POLY1305' in defines else 0
+    features["ML_DSA"] = 1 if '#define HAVE_DILITHIUM'  in defines else 0
 
     if '#define HAVE_FIPS' in defines:
         if not fips:
@@ -447,6 +451,7 @@ def build_ffi(local_wolfssl, features):
         #include <wolfssl/wolfcrypt/chacha20_poly1305.h>
         #include <wolfssl/wolfcrypt/kyber.h>
         #include <wolfssl/wolfcrypt/wc_kyber.h>
+        #include <wolfssl/wolfcrypt/dilithium.h>
     """
 
     init_source_string = """
@@ -484,6 +489,7 @@ def build_ffi(local_wolfssl, features):
         int RSA_PSS_ENABLED = """ + str(features["RSA_PSS"]) + """;
         int CHACHA20_POLY1305_ENABLED = """ + str(features["CHACHA20_POLY1305"]) + """;
         int ML_KEM_ENABLED = """ + str(features["ML_KEM"]) + """;
+        int ML_DSA_ENABLED = """ + str(features["ML_DSA"]) + """;
     """
 
     ffibuilder.set_source( "wolfcrypt._ffi", init_source_string,
@@ -520,6 +526,7 @@ def build_ffi(local_wolfssl, features):
         extern int RSA_PSS_ENABLED;
         extern int CHACHA20_POLY1305_ENABLED;
         extern int ML_KEM_ENABLED;
+        extern int ML_DSA_ENABLED;
 
         typedef unsigned char byte;
         typedef unsigned int word32;
@@ -929,12 +936,16 @@ def build_ffi(local_wolfssl, features):
         int wolfCrypt_GetPrivateKeyReadEnable_fips(enum wc_KeyType);
         """
 
+    if features["ML_KEM"] or features["ML_DSA"]:
+        cdef += """
+        static const int INVALID_DEVID;
+        """
+
     if features["ML_KEM"]:
         cdef += """
         static const int WC_ML_KEM_512;
         static const int WC_ML_KEM_768;
         static const int WC_ML_KEM_1024;
-        static const int INVALID_DEVID;
         typedef struct {...; } KyberKey;
         int wc_KyberKey_CipherTextSize(KyberKey* key, word32* len);
         int wc_KyberKey_SharedSecretSize(KyberKey* key, word32* len);
@@ -950,7 +961,29 @@ def build_ffi(local_wolfssl, features):
         int wc_KyberKey_EncapsulateWithRandom(KyberKey* key, unsigned char* ct, unsigned char* ss, const unsigned char* rand, int len);
         int wc_KyberKey_Decapsulate(KyberKey* key, unsigned char* ss, const unsigned char* ct, word32 len);
         int wc_KyberKey_EncodePrivateKey(KyberKey* key, unsigned char* out, word32 len);
-        int wc_KyberKey_DecodePrivateKey(KyberKey* key, const unsigned char* in, word32 len); 
+        int wc_KyberKey_DecodePrivateKey(KyberKey* key, const unsigned char* in, word32 len);
+        """
+
+    if features["ML_DSA"]:
+        cdef += """
+        static const int WC_ML_DSA_44;
+        static const int WC_ML_DSA_65;
+        static const int WC_ML_DSA_87;
+        typedef struct {...; } dilithium_key;
+        int wc_dilithium_init_ex(dilithium_key* key, void* heap, int devId);
+        int wc_dilithium_set_level(dilithium_key* key, byte level);
+        void wc_dilithium_free(dilithium_key* key);
+        int wc_dilithium_make_key(dilithium_key* key, WC_RNG* rng);
+        int wc_dilithium_export_private(dilithium_key* key, byte* out, word32* outLen);
+        int wc_dilithium_import_private(const byte* priv, word32 privSz, dilithium_key* key);
+        int wc_dilithium_export_public(dilithium_key* key, byte* out, word32* outLen);
+        int wc_dilithium_import_public(const byte* in, word32 inLen, dilithium_key* key);
+        int wc_dilithium_sign_msg(const byte* msg, word32 msgLen, byte* sig, word32* sigLen, dilithium_key* key, WC_RNG* rng);
+        int wc_dilithium_verify_msg(const byte* sig, word32 sigLen, const byte* msg, word32 msgLen, int* res, dilithium_key* key);
+        typedef dilithium_key MlDsaKey;
+        int wc_MlDsaKey_GetPrivLen(MlDsaKey* key, int* len);
+        int wc_MlDsaKey_GetPubLen(MlDsaKey* key, int* len);
+        int wc_MlDsaKey_GetSigLen(MlDsaKey* key, int* len);
         """
 
     ffibuilder.cdef(cdef)
@@ -983,7 +1016,8 @@ def main(ffibuilder):
         "AESGCM_STREAM": 1,
         "RSA_PSS": 1,
         "CHACHA20_POLY1305": 1,
-        "ML_KEM": 1
+        "ML_KEM": 1,
+        "ML_DSA": 1
     }
 
     # Ed448 requires SHAKE256, which isn't part of the Windows build, yet.

tests/test_mldsa.py

@@ -0,0 +1,136 @@
+# test_mldsa.py
+#
+# Copyright (C) 2025 wolfSSL Inc.
+#
+# This file is part of wolfSSL. (formerly known as CyaSSL)
+#
+# wolfSSL is free software; you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation; either version 2 of the License, or
+# (at your option) any later version.
+#
+# wolfSSL is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program; if not, write to the Free Software
+# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA
+
+# pylint: disable=redefined-outer-name
+
+from wolfcrypt._ffi import lib as _lib
+
+if _lib.ML_DSA_ENABLED:
+    import pytest
+
+    from wolfcrypt.ciphers import MlDsaPrivate, MlDsaPublic, MlDsaType
+    from wolfcrypt.random import Random
+
+    @pytest.fixture
+    def rng():
+        return Random()
+
+    @pytest.fixture(
+        params=[MlDsaType.ML_DSA_44, MlDsaType.ML_DSA_65, MlDsaType.ML_DSA_87]
+    )
+    def mldsa_type(request):
+        return request.param
+
+    def test_init_base(mldsa_type):
+        mldsa_priv = MlDsaPrivate(mldsa_type)
+        assert isinstance(mldsa_priv, MlDsaPrivate)
+
+        mldsa_pub = MlDsaPublic(mldsa_type)
+        assert isinstance(mldsa_pub, MlDsaPublic)
+
+    def test_size_properties(mldsa_type):
+        refvals = {
+            MlDsaType.ML_DSA_44: {
+                "sig_size": 2420,
+                "pub_key_size": 1312,
+                "priv_key_size": 2560,
+            },
+            MlDsaType.ML_DSA_65: {
+                "sig_size": 3309,
+                "pub_key_size": 1952,
+                "priv_key_size": 4032,
+            },
+            MlDsaType.ML_DSA_87: {
+                "sig_size": 4627,
+                "pub_key_size": 2592,
+                "priv_key_size": 4896,
+            },
+        }
+
+        mldsa_pub = MlDsaPublic(mldsa_type)
+        assert mldsa_pub.sig_size == refvals[mldsa_type]["sig_size"]
+        assert mldsa_pub.key_size == refvals[mldsa_type]["pub_key_size"]
+
+        mldsa_priv = MlDsaPrivate(mldsa_type)
+        assert mldsa_priv.sig_size == refvals[mldsa_type]["sig_size"]
+        assert mldsa_priv.pub_key_size == refvals[mldsa_type]["pub_key_size"]
+        assert mldsa_priv.priv_key_size == refvals[mldsa_type]["priv_key_size"]
+
+    def test_initializations(mldsa_type, rng):
+        mldsa_priv = MlDsaPrivate.make_key(mldsa_type, rng)
+        assert type(mldsa_priv) is MlDsaPrivate
+
+        mldsa_priv2 = MlDsaPrivate(mldsa_type)
+        assert type(mldsa_priv2) is MlDsaPrivate
+
+        mldsa_pub = MlDsaPublic(mldsa_type)
+        assert type(mldsa_pub) is MlDsaPublic
+
+    def test_key_import_export(mldsa_type, rng):
+        # Generate key pair and export keys
+        mldsa_priv = MlDsaPrivate.make_key(mldsa_type, rng)
+        priv_key = mldsa_priv.encode_priv_key()
+        pub_key = mldsa_priv.encode_pub_key()
+        assert len(priv_key) == mldsa_priv.priv_key_size
+        assert len(pub_key) == mldsa_priv.pub_key_size
+
+        # Export key pair from imported one
+        mldsa_priv2 = MlDsaPrivate(mldsa_type)
+        mldsa_priv2.decode_key(priv_key, pub_key)
+        priv_key2 = mldsa_priv2.encode_priv_key()
+        pub_key2 = mldsa_priv2.encode_pub_key()
+        assert priv_key == priv_key2
+        assert pub_key == pub_key2
+
+        # Export private key from imported one
+        mldsa_priv3 = MlDsaPrivate(mldsa_type)
+        mldsa_priv3.decode_key(priv_key)
+        priv_key3 = mldsa_priv3.encode_priv_key()
+        assert priv_key == priv_key3
+
+        # Export public key from imported one
+        mldsa_pub = MlDsaPublic(mldsa_type)
+        mldsa_pub.decode_key(pub_key)
+        pub_key3 = mldsa_pub.encode_key()
+        assert pub_key == pub_key3
+
+    def test_sign_verify(mldsa_type, rng):
+        # Generate a key pair and export public key
+        mldsa_priv = MlDsaPrivate.make_key(mldsa_type, rng)
+        pub_key = mldsa_priv.encode_pub_key()
+
+        # Import public key
+        mldsa_pub = MlDsaPublic(mldsa_type)
+        mldsa_pub.decode_key(pub_key)
+
+        # Sign a message
+        message = b"This is a test message for ML-DSA signature"
+        signature = mldsa_priv.sign(message, rng)
+        assert len(signature) == mldsa_priv.sig_size
+
+        # Verify the signature by MlDsaPrivate
+        assert mldsa_priv.verify(signature, message)
+
+        # Verify the signature by MlDsaPublic
+        assert mldsa_pub.verify(signature, message)
+
+        # Verify with wrong message
+        wrong_message = b"This is a wrong message for ML-DSA signature"
+        assert not mldsa_pub.verify(signature, wrong_message)