Comment trouver un bug dans une implémentation Rust de ML-DSA (f.k.a. Dilithium) ?

Contexte

Durant l’été 2024, j’étais en stage chez Fortanix, au Pays-Bas. Le sujet du stage était d’implémenter ML-DSA (fka. CRYSTALS-Dilithium, standardisé dans FIPS-204) en Rust pour ensuite l’intégrer au produit Fortanix DSM. Comme tout développeur qui se respecte, une fois l’implémentation terminée, j’ai écrits différents tests : roundtrip, confirmité, résilience (edge cases) ainsi que des tests d’intéropérabilité. C’est ce dernier qui m’a permis de trouver un bug dans une implémentation Rust de ML-DSA ; la suivante : integritychain/fips204 (le premier résultat google pour “rust fips204” !).

Un test d’intéropérabilité, c’est quoi ?

Un test d’intéropérabilité fonctionne selon le principe suivant : l’algorithme principal prends $n$ librairies à tester en entrée puis les utilisent en concurrence sur des inputs communs générés à la volée ; l’algorithme principal vérifie ensuite que ces librairies ce sont comportées de la même manière (que ce soit en gestion d’erreur ou valeur métier retournée). Voici une image illustrant ce fonctionnement pour $n = 2$ librairies :

Configuration du test

Une fois cet algorithme de test principal implémenté, je l’ai laissé s’exécuter sur deux librairies : celle que j’ai écrite durant mon stage et celle de la crate rust “fips204” de “integritychain” : integritychain/fips204.

Les harnais suivants était implémentés :

• ($\mathbf{sk}$, $\mathbf{vk}$) = $\textsf{KeyGen}(\xi)$ : l’algorithme principale fait varier $\xi$ et vérifie que la sérialization des deux clés sont égales ;
• $\sigma$ = $\textsf{Sign}(\mu, \mathbf{sk})$ : l’algorithme principale fait varier $\mu$ et $\mathbf{sk}$ et vérifie que les deux signatures $\sigma$ sont égales, cela nécessite donc d’utiliser la version déterministe ou bien de hooker le PRNG ;
• $b$ = $\textsf{Verify}(\sigma, \mu, \mathbf{vk})$ : l’algorithme principale fait varier $\sigma$, $\mu$ et $\mathbf{vk}$ et vérifie que le bit $b$ est identique.

Chaque harnais vérifie également que les fonctions $(\textsf{KeyGen}, \textsf{Sign}, \textsf{Verify})$ ne retournent pas d’erreur ou alors retournent la même.

Après avoir lancé l’algorithme de test avec les harnais décrits plus tôt, c’est au bout de 3/4 minutes de test que les harnais de $\textsf{KeyGen}$ a relevé une différence. En effet, les clés sérializés n’étaient pas les mêmes pour une seed $\xi$ donnée.

Analyse de la différence

Pour rappel, voici l’algorithme de génération de clé tel que décrit dans FIPS 204 :

Le problème avec le test d’intéropérabilité est qu’une fois la différence trouvée, il faut maitenant déterminer de quelle librairie vient le bug… Après de multiples println! de valeurs intermédiaires j’ai trouvé que la différence était présente dès l’entrée de la fonction Power2Round. En effet, le code de la crate d’integritychain effectue les opérations suivantes :

// 5: t ← NTT−1(cap_a_hat ◦ NTT(s_1)) + s_2    ▷ Compute t = As1 + s2
let t: [R; K] = {
    let s_1_hat: [T; L] = ntt(&s_1);
    let as1_hat: [T; K] = mat_vec_mul(&cap_a_hat, &s_1_hat);
    vec_add(&inv_ntt(&as1_hat), &s_2)
};

La sortie de la fonction inv_ntt est réduite modulo $q$, i.e. contient des coefficients dans $\lbrace 0, 1, …, q-1 \rbrace$, et $\mathbf{s_2}$ contient des coefficients dans $[-\eta,\eta]$. Le résultat de vec_add(&inv_ntt(&as1_hat), &s_2) contient donc des coefficients dans $ \lbrace -\eta, -\eta +1, …, q-1 + \eta \rbrace $. Problème : la fonction Power2Round attends en entrée un vecteur avec des coefficients modulo $q$.

Pour rappel, la fonction Power2Round décompose un entier $r \in \mathbb{Z}_q$ en deux entiers $(r_1, r_2)$ tel que $r \equiv r_1 2^d + r_0~\mathrm{mod}$ et $r_0~\mathrm{mod}^{\pm}$. Par soucis d’optimisation la fonction fait l’hypothèse que l’entrée est réduite $\mathrm{mod}~q$ et opère directement sur les bits.

Voici l’implémentation de la fonction Power2Round :

pub(crate) fn power2round<const K: usize>(r: &[R; K]) -> ([R; K], [R; K]) {
    // 1: r+ ← r mod q
    // 2: r0 ← r+ mod±2^d
    // 3: return ((r+ − r0)/2^d, r0)

    let r_1: [R; K] = core::array::from_fn(|k| {
        R(core::array::from_fn(|n| (r[k].0[n] + (1 << (D - 1)) - 1) >> D))
    });
    let r_0: [R; K] =
        core::array::from_fn(|k| R(core::array::from_fn(|n| r[k].0[n] - (r_1[k].0[n] << D))));

    debug_assert!(
        {
            let mut result = true;
            for k in 0..K {
                for n in 0..256 {
                    result &= r[k].0[n] == ((r_1[k].0[n] << D) + r_0[k].0[n]);
                }
            }
            result
        },
        "Alg 29: fail"
    );

    (r_1, r_0)
}

Cette implémentation traite un nombre comme un int32_t et fonctionnera différement pour deux entrées différentes mais équivalentes $\mod{q}$. Elle renvoit donc des valeurs $(r_0, r_1)$ incorrectes. TODO : figure conclusion

Impacts du bug

TODO

Probabilité d’apparition du bug

TODO