mldsa-0.1.0.0: src/Internal.hs

-- |
-- Module      : Internal
-- License     : BSD-3-Clause
-- Copyright   : (c) 2026 Olivier Chéron
--
-- ML-DSA main internal algorithms
--
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TypeFamilies #-}
module Internal
    ( ParamSet(..), Params(..), Encode(..), Decode(..)
    , PrivateKey, PublicKey, Signature
    , keyGen, sigGen, sigGenMu, sigVer, sigVerMu, toPublic, checkKeyPair
    ) where

import Control.DeepSeq (NFData(..))
import Control.Monad

import Data.ByteArray (ByteArray, ByteArrayAccess, Bytes, ScrubbedBytes)
import qualified Data.ByteArray as B
import qualified Data.Memory.Endian as B

import Data.Bits
import Data.Proxy
import Data.Word

import Auxiliary (Hints, Rq, Tq)
import Base
import BlockN (BlockN)
import Builder (Builder)
import qualified Builder
import qualified Crypto
import Marking (SecurityMarking(..), Leak(..))
import Math
import Vector (Vector)
import qualified Auxiliary as Aux
import qualified Matrix
import qualified Vector

data Params (k :: Nat) (l :: Nat) = Params
    { tau :: {-# UNPACK #-} !Int
    , lambdaDiv4 :: {-# UNPACK #-} !Int
    , gamma1Bits :: {-# UNPACK #-} !Int
    , gamma2 :: {-# UNPACK #-} !Word32
    , wBits :: {-# UNPACK #-} !Int        -- bitlen ((𝑞-1)/(2𝛾2)-1)
    , eta :: {-# UNPACK #-} !Int
    , twoEtaBits :: {-# UNPACK #-} !Int
    , beta :: {-# UNPACK #-} !Word
    , omega :: {-# UNPACK #-} !Int
    }

kdim :: KnownNat k => Params k l -> Int
kdim = fromIntegral . natVal . up
  where
    up :: a b (c :: k) -> Proxy b
    up _ = Proxy

ldim :: KnownNat l => Params k l -> Int
ldim = fromIntegral . natVal

-- | The class of ML-DSA parameter sets.
class (KnownNat (K a), KnownNat (L a)) => ParamSet a where
    type K a :: Nat
    type L a :: Nat
    getParams :: proxy a -> Params (K a) (L a)

-- | Utility class to serialize ML-DSA objects to byte arrays.
class Encode obj where
    -- | Serializes an object to a sequence of bytes.
    encode :: (ParamSet a, ByteArray ba) => obj a -> ba

-- | Utility class to deserialize ML-DSA objects from byte arrays.
class Decode obj where
    -- | Deserializes an object from a sequence of bytes.
    decode :: (ParamSet a, ByteArrayAccess ba) => proxy a -> ba -> Maybe (obj a)

-- | An ML-DSA private key.
data PrivateKey a = PrivateKey
    { skPub    :: {-# UNPACK #-} !(PublicKey a)
    , skK      :: {-# UNPACK #-} !ScrubbedBytes
    , skTr     :: {-# UNPACK #-} !Bytes
    , skE1     :: {-# UNPACK #-} !ScrubbedBytes     -- serialized s1 & s2
    , skE2     :: {-# UNPACK #-} !Bytes             -- serialized t0
    , skS1     :: Vector (L a) (Tq Sec)
    , skS2     :: Vector (K a) (Tq Sec)
    , skT0     :: Vector (K a) (Tq Pub)
    }

skRho :: PrivateKey a -> Bytes
skRho = pkRho . skPub

skT1 :: PrivateKey a -> Vector (K a) (Tq Pub)
skT1 = pkT1 . skPub

skA :: PrivateKey a -> Vector (K a) (Vector (L a) (Tq Pub))
skA = pkA . skPub

-- | An ML-DSA public key.
data PublicKey a = PublicKey
    { pkRho :: {-# UNPACK #-} !Bytes
    , pkE   :: {-# UNPACK #-} !Bytes        -- serialized t1
    , pkT1  :: Vector (K a) (Tq Pub)
    , pkA   :: Vector (K a) (Vector (L a) (Tq Pub))
    }

-- | An ML-DSA signature.
data Signature a = Signature
    { sigCt     :: {-# UNPACK #-} !ScrubbedBytes
    , sigZ      :: {-# UNPACK #-} !(Vector (L a) (Rq Sec))
    , sigH      :: {-# UNPACK #-} !(Vector (K a) Hints)
    }

instance Eq (PrivateKey a) where
    a == b = Crypto.toBool $
        Crypto.constEqW (skRho a) (skRho b) `Crypto.andW`
        Crypto.constEqW (skK a) (skK b) `Crypto.andW`
        Crypto.constEqW (skTr a) (skTr b) `Crypto.andW`
        Crypto.constEqW (skE1 a) (skE1 b) `Crypto.andW`
        Crypto.constEqW (skE2 a) (skE2 b)

instance Eq (PublicKey a) where
    a == b = Crypto.toBool $
        Crypto.constEqW (pkRho a) (pkRho b) `Crypto.andW`
        Crypto.constEqW (pkE a) (pkE b)

instance Eq (Signature a) where
    a == b = Crypto.toBool $
        Crypto.constEqW (sigCt a) (sigCt b) `Crypto.andW`
        Crypto.constEqW (sigZ a) (sigZ b) `Crypto.andW`
        Crypto.constEqW (sigH a) (sigH b)

instance Show (PrivateKey a) where
#ifdef ML_DSA_TESTING
    showsPrec d sk = showParen (d > 10) $
        showString "PrivateKey " . showsPrec 11 (skEncode sk :: Bytes)
#else
    showsPrec _ _ = showString "PrivateKey"
#endif

instance Show (PublicKey a) where
    showsPrec d pk = showParen (d > 10) $
        showString "PublicKey " . showsPrec 11 (pkEncode pk :: Bytes)

instance ParamSet a => Show (Signature a) where
    showsPrec d sig = showParen (d > 10) $
        showString "Signature " . showsPrec 11 (sigEncode sig :: Bytes)

instance NFData (PrivateKey a) where
    rnf sk = rnf (skRho sk) `seq`
             rnf (skK sk) `seq`
             rnf (skTr sk) `seq`
             rnf (skE1 sk) `seq`
             rnf (skE2 sk)
    -- skS1, skS2, skT0, skT1, skA omitted because just for caching

instance NFData (PublicKey a) where
    rnf pk = rnf (pkRho pk) `seq` rnf (pkE pk)
    -- pkT1, pkA omitted because just for caching

instance NFData (Signature a) where
    rnf sig = rnf (sigCt sig) `seq`
              Vector.toNormalForm (sigZ sig) `seq`
              Vector.toNormalForm (sigH sig)

instance Encode PublicKey where
    encode = pkEncode

-- Encodes a public key for ML-DSA into a byte string
pkEncode :: ByteArray ba => PublicKey a -> ba
pkEncode pk = Builder.runRelaxed $
    Builder.bytes (pkRho pk) <> Builder.bytes (pkE pk)

instance Decode PublicKey where
    -- Reverses the procedure pkEncode
    decode p input = do
        guard (B.length input == 32 + 320 * k)
        let rho = B.convert $ B.takeView input 32
            pe = B.convert $ B.dropView input 32
            t1 = Vector.create $ \i -> Aux.simpleBitUnpack10 (view320 i)
            aa = expandA rho
        Just PublicKey { pkRho = rho, pkE = pe, pkT1 = Aux.ntt <$> t1, pkA = aa }
      where
        params = getParams p
        k = kdim params
        view320 (Offset i) = B.view input (32 + 320 * i) 320

instance Encode PrivateKey where
    encode = skEncode

-- Encodes a secret key for ML-DSA into a byte string
skEncode :: ByteArray ba => PrivateKey a -> ba
skEncode sk = Builder.runRelaxed $
    Builder.bytes (skRho sk) <>
    leak (Builder.bytes (skK sk)) <>
    Builder.bytes (skTr sk) <>
    leak (Builder.bytes (skE1 sk)) <>
    Builder.bytes (skE2 sk)

instance Decode PrivateKey where
    -- Reverses the procedure skEncode
    decode p input = do
        guard (B.length input == 128 + teb32 * (l + k) + 416 * k)
        let rho = B.convert $ B.view input  0 32
            kk  = B.convert $ B.view input 32 32
            tr'  = B.convert $ B.view input 64 64
            e1   = B.convert $ B.view input 128 (teb32 * (l + k))
            e2   = B.convert $ B.dropView input (128 + teb32 * (l + k))
            t0'  = Vector.create $ \i -> Aux.bitUnpackSafe 13 (t0Elem i)
            aa   = expandA rho
        s1 <- Vector.createMaybe $ \i -> Aux.bitUnpack (fromIntegral eta) twoEtaBits (s1Elem i)
        s2 <- Vector.createMaybe $ \i -> Aux.bitUnpack (fromIntegral eta) twoEtaBits (s2Elem i)
        let ss1 = Aux.ntt <$> s1
            ss2 = Aux.ntt <$> s2
            t = Aux.nttInv <$> Matrix.mmulAdd aa ss1 ss2
            (t1, t0) = Vector.unzipWith Aux.powerTwoRound t
            pk = PublicKey { pkRho = rho, pkE = pe, pkT1 = Aux.ntt <$> t1, pkA = aa }
            tr = Builder.run $ Crypto.h64 (encode pk :: Bytes)
            sk = PrivateKey { skPub = pk, skK = kk, skTr = tr, skE1 = e1, skE2 = e2, skS1 = ss1, skS2 = ss2, skT0 = Aux.ntt <$> t0 }
            pe = Builder.run $ Vector.concatMap Aux.simpleBitPack10 t1
        guard $ Crypto.toBool $
            Crypto.constEqW t0 t0' `Crypto.andW` Crypto.constEqW tr tr'
        return sk
      where
        s1Elem (Offset i) = B.view input (128 + teb32 * i) teb32
        s2Elem (Offset i) = B.view input (128 + teb32 * (l + i)) teb32
        t0Elem (Offset i) = B.view input (128 + teb32 * (l + k) + 416 * i) 416

        teb32 = 32 * twoEtaBits

        params@Params{..} = getParams p
        k = kdim params
        l = ldim params

instance Encode Signature where
    encode = sigEncode

-- Encodes a signature into a byte string
sigEncode :: (ParamSet a, ByteArray ba) => Signature a -> ba
sigEncode sig = Builder.runRelaxed $ leak $
    Builder.bytes (sigCt sig) <>
    Vector.concatMap (Aux.bitPackSafe (1 + gamma1Bits)) (sigZ sig) <>
    Aux.hintBitPack omega (sigH sig)
  where Params{..} = getParams sig

instance Decode Signature where
    -- Reverses the procedure sigEncode
    decode p input = do
        guard (B.length input == lambdaDiv4 + zElemLen * l + omega + k)
        let ct  = B.convert $ B.view input 0 lambdaDiv4
            z   = Vector.create $ \i -> Aux.bitUnpackSafe (1 + gamma1Bits) (zElem i)
            y   = B.view input (lambdaDiv4 + zElemLen * l) (omega + k)
        h <- Aux.hintBitUnpack omega y
        return Signature { sigCt = ct, sigZ = z, sigH = h }
      where
        zElem (Offset i) = B.view input (lambdaDiv4 + zElemLen * i) zElemLen
        zElemLen = 32 * (1 + gamma1Bits)

        params@Params{..} = getParams p
        k = kdim params
        l = ldim params

-- Generates a public-private key pair from a seed
keyGen :: (ParamSet a, ByteArrayAccess seed) => proxy a -> seed -> (PublicKey a, PrivateKey a)
keyGen p xi = (pk, sk)
  where
    (rho, rho', kk) = Crypto.h128 xi (fromIntegral k) (fromIntegral l)
    aa = expandA rho
    (s1, s2) = expandS l eta rho'
    ss1 = Aux.ntt <$> s1
    ss2 = Aux.ntt <$> s2
    t = Aux.nttInv <$> Matrix.mmulAdd aa ss1 ss2
    (t1, t0) = Vector.unzipWith Aux.powerTwoRound t
    pk = PublicKey { pkRho = rho, pkE = pe, pkT1 = Aux.ntt <$> t1, pkA = aa }
    tr = Builder.run $ Crypto.h64 (encode pk :: Bytes)
    sk = PrivateKey { skPub = pk, skK = kk, skTr = tr, skE1 = e1, skE2 = e2, skS1 = ss1, skS2 = ss2, skT0 = Aux.ntt <$> t0 }
    pe = Builder.run $ Vector.concatMap Aux.simpleBitPack10 t1
    e1 = Builder.run $
            Vector.concatMap (Aux.bitPack (fromIntegral eta) twoEtaBits) s1 <>
            Vector.concatMap (Aux.bitPack (fromIntegral eta) twoEtaBits) s2
    e2 = Builder.run $ Vector.concatMap (Aux.bitPackSafe 13) t0

    params@Params{..} = getParams p
    k = kdim params
    l = ldim params

-- Deterministic algorithm to generate a signature for a formatted message 𝑀′
sigGen :: (ParamSet a, ByteArrayAccess m, ByteArrayAccess rnd) => PrivateKey a -> m -> rnd -> Signature a
sigGen sk m' = sigGenMu sk mu
  where
    tr = Builder.bytes (skTr sk)
    mu = Crypto.h64 (Builder.run $ Builder.promote tr <> Builder.secret m')

sigGenMu :: (ParamSet a, ByteArrayAccess rnd) => PrivateKey a -> Builder Sec -> rnd -> Signature a
sigGenMu sk mu rnd = loop 0
  where
    rhos = Crypto.h64 (Builder.run $ Builder.bytes (skK sk) <> Builder.secret rnd <> mu)

    loop kappa
        | Crypto.toBool cn1 && Crypto.toBool cn2 =
            Signature { sigCt = ct, sigZ = z, sigH = h }
        | otherwise = loop (kappa + l)
      where
        y   = expandMask rhos kappa
        w   = Aux.nttInv <$> Matrix.mmul (skA sk) (Aux.ntt <$> y)
        w1  = Aux.highBits gamma2 <$> w
        ct  = Crypto.h lambdaDiv4 (Builder.run $ mu <> w1Encode wBits w1)
        cc  = Aux.ntt $ Aux.sampleInBall tau ct
        cs1 = Aux.nttInv . (cc ..*) <$> Vector.seq cc (skS1 sk)
        cs2 = Aux.nttInv . (cc ..*) <$> Vector.seq cc (skS2 sk)
        z   = y .+ cs1
        r0  = Aux.lowBits gamma2 <$> (w .- cs2)
        cn1 = Crypto.ltW (normVector z) ((1 `unsafeShiftL` gamma1Bits) - beta) `Crypto.andW`
              Crypto.ltW (normVector r0) (fromIntegral gamma2 - beta)
        ct0 = Aux.nttInv . (..* cc) <$> Vector.seq cc (skT0 sk)
        h   = Vector.zipWith (Aux.makeHint gamma2) (neg ct0) (w .- cs2 .+ ct0)
        cn2 = Crypto.ltW (normVector ct0) (fromIntegral gamma2) `Crypto.andW`
              Crypto.lteW (countHints h) (fromIntegral omega)

    params@Params{..} = getParams sk
    l = ldim params

    -- Samples a vector 𝐲 ∈ 𝑅ℓ such that each polynomial 𝐲[𝑟] has
    -- coefficients between -𝛾1 + 1 and 𝛾1
    expandMask :: KnownNat l => Builder Sec -> Int -> Vector l (Rq Sec)
    expandMask rho mu' = Vector.create $ \(Offset r) ->
        let i = fromIntegral (mu' + r) :: Word16
            rho' = Builder.run $ rho <> Builder.storable (B.toLE i)
            v = Crypto.h (32 * c) rho'
         in Aux.bitUnpackSafe c v
      where c = 1 + gamma1Bits

-- Internal function to verify a signature 𝜎 for a formatted message 𝑀′
sigVer :: (ParamSet a, ByteArrayAccess m) => PublicKey a -> m -> Signature a -> Bool
sigVer pk m' = sigVerMu pk mu
  where
    tr = Crypto.h64 (encode pk :: Bytes)
    mu = Crypto.h64 (Builder.run $ Builder.promote tr <> Builder.secret m')

sigVerMu :: ParamSet a => PublicKey a -> Builder Sec -> Signature a -> Bool
sigVerMu pk mu sig =
    normVector (sigZ sig) < (1 `unsafeShiftL` gamma1Bits) - beta &&
    Crypto.toBool (Crypto.constEqW (sigCt sig) ct')
  where
    cc = Aux.ntt $ neg $ Aux.sampleInBall tau (sigCt sig)
    wApprox = Aux.nttInv <$> Matrix.mmulAdd (pkA pk) (Aux.ntt <$> sigZ sig) ((..* cc) <$> Vector.seq cc (pkT1 pk))
    w1  = Vector.zipWith (Aux.useHint gamma2) (sigH sig) wApprox
    ct' = Crypto.h lambdaDiv4 (Builder.run $ mu <> w1Encode wBits w1)

    Params{..} = getParams pk

-- Encodes a polynomial vector 𝐰1 into a byte string
w1Encode :: Int -> Vector n (BlockN Sec 256 Word32) -> Builder Sec
w1Encode wBits = Vector.concatMap (Aux.simpleBitPack wBits)

-- Samples a 𝑘 × ℓ matrix 𝐀 of elements of 𝑇𝑞
expandA :: (KnownNat k, KnownNat l) => Bytes -> Vector k (Vector l (Tq Pub))
expandA !rho = Matrix.create $ \(Offset s) (Offset r) ->
    Aux.rejNttPoly rho (fromIntegral s) (fromIntegral r)

-- Samples vectors 𝐬1 ∈ 𝑅ℓ and 𝐬2 ∈ 𝑅𝑘, each with polynomial coordinates
-- whose coefficients are in the interval [-𝜂, 𝜂]
expandS :: (KnownNat k, KnownNat l) => Int -> Int -> ScrubbedBytes -> (Vector l (Rq Sec), Vector k (Rq Sec))
expandS l eta rho = (s1, s2)
  where
    s1 = Vector.create $ \(Offset r) -> Aux.rejBoundedPoly rho eta (fromIntegral r)
    s2 = Vector.create $ \(Offset r) -> Aux.rejBoundedPoly rho eta (fromIntegral $ r + l)

countHints :: Vector n Hints -> Word
countHints = Vector.foldl' Aux.countFrom 0

normVector :: Vector n (Rq Sec) -> Word
normVector = fromIntegral . Aux.getNorm . Vector.foldl' Aux.normFrom mempty

-- | Returns the public key associated to the given private key.
toPublic :: PrivateKey a -> PublicKey a
toPublic = skPub

-- | Returns @True@ when the public key and private key both match.  Note that
-- this does not fully guarantee that the key pair was properly generated.
checkKeyPair :: (PublicKey a, PrivateKey a) -> Bool
checkKeyPair (pk, sk) = Crypto.toBool $
    Crypto.constEqW (pkRho pk) (skRho sk) `Crypto.andW` Crypto.constEqW (pkT1 pk) (skT1 sk)
    -- consistency of skT0 with respect to skRho is verified when decoding
    -- the private key