Complete Implementation of ML-DSA (FIPS 204)

Every code listing from this chapter of Applied Post-Quantum Cryptography — 2 in total, 1 runnable here. Edit any cell and press Run.

The book's snippets build on each other down the chapter, but a Sage Cell kernel runs one cell and keeps no state afterwards, so each cell replays the earlier listings with apqc_book. That call is the only thing added to the book's own code.

← the playground

Listing 1 — SageMath experiment

The core cancellation is easy to verify in miniature. Using small parameters, sample a key, form a masked response, and confirm the verification identity holds before any signature is even released.

Listing 2 — A complete SageMath implementation

The cancellation check above is a fragment. The companion file sage/fips204_mldsa.sage is the whole standard: all forty-nine numbered algorithms of FIPS 204, each as its own function annotated with its algorithm number, for ML-DSA-44, 65, and 87. That includes the parts this chapter has treated as black boxes – ExpandA, ExpandS, ExpandMask, SampleInBall, Power2Round, Decompose, MakeHint, UseHint – as well as the bit packing, the hint encoding, the pre-hash variants HashML-DSA, and Montgomery reduction.

As in FIPS 203 the algebra is carried by Sage. Here, though, the standard is careful to distinguish the ring R_q from the ring R of integer polynomials with a restricted coefficient range, and so is the code: quantities reduced modulo q (the secrets s₁,s₂, the mask y, the response z) live in the Sage quotient ring R_q = 𝔽₈₃₈₀₄₁₇[X]/(X²⁵⁶+1), while t₁, w₁, r₀, and the hint h stay integer lists. The NTT domain T_q is the direct product (ℤ_q)²⁵⁶, so MultiplyNTT (Algorithm 45) is literally a coordinatewise product.

The rejection loop of the algorithm in the book then reads almost exactly as the standard writes it:

Not runnable on its own: the body of Sign_internal, quoted from the companion file -- it refers to self and returns, so it only runs as a method.

kappa = 0
while True:
    y  = self.ExpandMask(rhopp, kappa)
    yh = [NTT(Rq(p)) for p in y]
    w  = [NTTInverse(a) for a in MatrixVectorNTT(Ah, yh)]
    w1 = HighBitsVec(w, self.gamma2)

    ctilde = H(mu + self.w1Encode(w1), self.lam // 4)
    c      = self.SampleInBall(ctilde)
    ch     = NTT(Rq(c))

    cs1  = [NTTInverse(MultiplyNTT(ch, a)) for a in s1h]
    cs2  = [NTTInverse(MultiplyNTT(ch, a)) for a in s2h]
    z    = [Rq(y[i]) + cs1[i] for i in range(self.ell)]
    wcs2 = [w[i] - cs2[i] for i in range(self.k)]
    r0   = LowBitsVec(wcs2, self.gamma2)

    kappa += self.ell
    if inf_norm_vec(z) >= self.gamma1 - self.beta:      # unsafe: retry
        continue
    if inf_norm_ints(r0) >= self.gamma2 - self.beta:
        continue

    ct0 = [NTTInverse(MultiplyNTT(ch, a)) for a in t0h]
    h   = MakeHintVec([-p for p in ct0],
                      [wcs2[i] + ct0[i] for i in range(self.k)],
                      self.gamma2)
    if inf_norm_vec(ct0) >= self.gamma2:
        continue
    if sum(sum(p) for p in h) > self.omega:
        continue

    return self.sigEncode(ctilde, [centered(zi) for zi in z], h)