ML_DSA.hs

-- |
-- Module      : Crypto.PubKey.ML_DSA
-- License     : BSD-3-Clause
-- Maintainer  : Olivier Chéron <olivier.cheron@gmail.com>
-- Stability   : provisional
-- Portability : unknown
--
-- Module-Lattice-based Digital Signature Algorithm (ML-DSA), defined
-- in <https://csrc.nist.gov/pubs/fips/204/final FIPS 204>.
{-# LANGUAGE CPP #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE TypeFamilies #-}
module Crypto.PubKey.ML_DSA
    ( PublicKey, PrivateKey, Signature

    -- * Operations
    , generate, generateWith
    , Context, context, defaultContext, Randomness, randomness, deterministic

    -- ** Pure version
    --
    -- | This is the preferred and most common version of ML-DSA.
    , sign, signWith, verify

    -- ** Pre-hash version
    --
    -- | This version of ML-DSA can be used to sign or verify a message that is
    -- hashed externally.  Typical use case is when the message is too large to
    -- fit into memory, as pre-hashing can be performed through an incremental
    -- API.
    --
    -- Note that signatures produced by HashML-DSA are incompatible with the
    -- "pure" version above.  So both signing and verification must use the same
    -- version of ML-DSA.  Keys can use either version indistinctly but it is
    -- advised to restrict them to one version only to avoid ambiguity.
    , PreHashAlgorithm, signDigest, signDigestWith, verifyDigest

    -- ** External µ version
    --
    -- | An internal signing and verfication API that operates on a message
    -- representative µ.  For testing purpose only.
    , Mu, externalMu, signExternalMu, signExternalMuWith, verifyExternalMu

    -- ** Internal version
    --
    -- | An internal API that exposes algorithms @ML-DSA.Sign_internal@ and
    -- @ML-DSA.Verify_internal@.  For testing purpose only.
    , signInternal, signInternalWith, verifyInternal

    -- * Parameter sets
    , ParamSet, ML_DSA_44, ML_DSA_65, ML_DSA_87

    -- * Conversions and checks
    , Decode(..), Encode(..), toPublic, checkKeyPair

    ) where

import Crypto.Hash (Digest)
import Crypto.Hash.Algorithms
import Crypto.Random

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

import Data.Word

import Builder
import Internal

-- | ML-DSA-44 (security category 2)
data ML_DSA_44 = ML_DSA_44 deriving Show
-- | ML-DSA-65 (security category 3)
data ML_DSA_65 = ML_DSA_65 deriving Show
-- | ML-DSA-87 (security category 5)
data ML_DSA_87 = ML_DSA_87 deriving Show

instance ParamSet ML_DSA_44 where
    type K ML_DSA_44 = 4
    type L ML_DSA_44 = 4
    getParams _ = Params 39 32 17  95232 6 2 3  78 80
instance ParamSet ML_DSA_65 where
    type K ML_DSA_65 = 6
    type L ML_DSA_65 = 5
    getParams _ = Params 49 48 19 261888 4 4 4 196 55
instance ParamSet ML_DSA_87 where
    type K ML_DSA_87 = 8
    type L ML_DSA_87 = 7
    getParams _ = Params 60 64 19 261888 4 2 3 120 75

-- | A byte string used during signature generation and verification for domain
-- separation.  Always 255 bytes or less.  Use 'defaultContext' unless specified
-- otherwise.
data Context = forall ba. ByteArrayAccess ba => Context ba

instance Eq Context where
    Context a == Context b = B.eq a b

instance Show Context where
    showsPrec d (Context b) = showParen (d > 10) $
        showString "Context " . showsPrec 11 (B.convert b :: Bytes)

instance ByteArrayAccess Context where
    length (Context ctx) = B.length ctx
    withByteArray (Context ctx) = B.withByteArray ctx
    copyByteArrayToPtr (Context ctx) = B.copyByteArrayToPtr ctx

-- | Smart constructor for a context value.  Length of input must be 255 bytes
-- maximum.
context :: ByteArrayAccess ba => ba -> Maybe Context
context ctx
    | B.length ctx < 256 = Just $ Context ctx
    | otherwise = Nothing

-- | The default (empty) context.
defaultContext :: Context
defaultContext = Context (B.empty :: Bytes)

-- | A source of randomness to be used during signature generation.
-- Always 32 bytes.
--
-- The use of fresh randomness during signing helps mitigate side-channel
-- attacks.
data Randomness = forall ba. ByteArrayAccess ba => Randomness ba

instance Eq Randomness where
    Randomness a == Randomness b = B.constEq a b

instance Show Randomness where
#ifdef ML_DSA_TESTING
    showsPrec d (Randomness b) = showParen (d > 10) $
        showString "Randomness " . showsPrec 11 (B.convert b :: Bytes)
#else
    showsPrec _ _ = showString "Randomness"
#endif

instance ByteArrayAccess Randomness where
    length (Randomness _) = 32
    withByteArray (Randomness rnd) = B.withByteArray rnd
    copyByteArrayToPtr (Randomness rnd) = B.copyByteArrayToPtr rnd

-- | Smart constructor for randomness.  Length of input must be 32 bytes.
randomness :: ByteArrayAccess ba => ba -> Maybe Randomness
randomness rnd
    | B.length rnd == 32 = Just $ Randomness rnd
    | otherwise = Nothing

-- | Enables a fully deterministic variant of the signing procedure.  Can be
-- used if the signer has no access to a fresh source of randomness at signing
-- time.  However this increases risks side-channel attacks, particularly fault
-- attacks.
deterministic :: Randomness
deterministic = Randomness (B.replicate 32 0 :: Bytes)

-- | Generate an ML-DSA key pair from a random seed.
generate :: (ParamSet a, MonadRandom m)
         => proxy a -> m (PublicKey a, PrivateKey a)
generate p = do
    xi <- getRandomBytes 32
    return $ Internal.keyGen p (xi :: ScrubbedBytes)

-- | Generate an ML-DSA key pair from the specified seed.  Length of input
-- must be 32 bytes.
generateWith :: (ParamSet a, ByteArrayAccess seed)
             => proxy a -> seed -> Maybe (PublicKey a, PrivateKey a)
generateWith p xi
    | B.length xi /= 32 = Nothing
    | otherwise = Just $ Internal.keyGen p xi

-- | Generates an ML-DSA signature.  Fresh randomness is acquired from a
-- 'MonadRandom' instance.
sign :: (ParamSet a, ByteArrayAccess message, MonadRandom m)
     => PrivateKey a -> message -> Context -> m (Signature a)
sign sk m ctx = do
    rnd <- getRandomBytes 32
    return $ Internal.sigGen sk (getPureM' m ctx) (rnd :: ScrubbedBytes)

-- | Generates an ML-DSA signature using explicit randomness.
signWith :: (ParamSet a, ByteArrayAccess message)
         => Randomness -> PrivateKey a -> message -> Context -> Signature a
signWith (Randomness rnd) sk m ctx = Internal.sigGen sk (getPureM' m ctx) rnd

-- | Verifies an ML-DSA signature.  Returns @True@ when the signature is valid
-- for the message.
verify :: (ParamSet a, ByteArrayAccess message)
       => PublicKey a -> message -> Signature a -> Context -> Bool
verify pk m sig ctx = Internal.sigVer pk (getPureM' m ctx) sig

getPureM' :: (ByteArrayAccess message, ByteArrayAccess context)
          => message -> context -> ScrubbedBytes
getPureM' m ctx =
    let ctxLen = fromIntegral (B.length ctx) :: Word16
     in Builder.run $ Builder.storable (B.toBE ctxLen) <>
            Builder.public ctx <> Builder.public m

-- | Generates an HashML-DSA signature.  Fresh randomness is acquired from a
-- 'MonadRandom' instance.
signDigest :: (ParamSet a, PreHashAlgorithm alg, MonadRandom m)
           => PrivateKey a -> Digest alg -> Context -> m (Signature a)
signDigest sk phm ctx = do
    rnd <- getRandomBytes 32
    return $ Internal.sigGen sk (getPreHashM' phm ctx) (rnd :: ScrubbedBytes)

-- | Generates a HashML-DSA signature using explicit randomness.
signDigestWith :: (ParamSet a, PreHashAlgorithm alg)
               => Randomness -> PrivateKey a -> Digest alg -> Context -> Signature a
signDigestWith (Randomness rnd) sk phm ctx =
    Internal.sigGen sk (getPreHashM' phm ctx) rnd

-- | Verifies a HashML-DSA signature.  Returns @True@ when the signature is
-- valid for the message digest.
verifyDigest :: (ParamSet a, PreHashAlgorithm alg)
             => PublicKey a -> Digest alg -> Signature a -> Context -> Bool
verifyDigest pk phm sig ctx = Internal.sigVer pk (getPreHashM' phm ctx) sig

getPreHashM' :: (PreHashAlgorithm alg, ByteArrayAccess context)
             => Digest alg -> context -> ScrubbedBytes
getPreHashM' phm ctx =
    let ctxLen = 256 + fromIntegral (B.length ctx) :: Word16
     in Builder.run $ Builder.storable (B.toBE ctxLen) <> Builder.public ctx <>
            Builder.public (oid phm) <> Builder.public phm

-- | A message representative for ML-DSA.  Always 64 bytes.
data Mu = forall ba. ByteArrayAccess ba => Mu ba

instance Eq Mu where
    Mu a == Mu b = B.eq a b

instance Show Mu where
    showsPrec d (Mu b) = showParen (d > 10) $
        showString "Mu " . showsPrec 11 (B.convert b :: Bytes)

instance ByteArrayAccess Mu where
    length (Mu _) = 64
    withByteArray (Mu mu) = B.withByteArray mu
    copyByteArrayToPtr (Mu mu) = B.copyByteArrayToPtr mu

-- | Smart constructor for a µ value.  Length of input must be 64 bytes.
externalMu :: ByteArrayAccess ba => ba -> Maybe Mu
externalMu mu
    | B.length mu == 64 = Just $ Mu mu
    | otherwise = Nothing

-- | Generates a signature for a message representative µ value.  Fresh
-- randomness is acquired from a 'MonadRandom' instance.
signExternalMu :: (ParamSet a, MonadRandom m)
               => PrivateKey a -> Mu -> m (Signature a)
signExternalMu sk (Mu mu) = do
    rnd <- getRandomBytes 32
    return $ Internal.sigGenMu sk (Builder.public mu) (rnd :: ScrubbedBytes)

-- | Generates a signature for a message representative µ value using explicit
-- randomness.
signExternalMuWith :: (ParamSet a)
                   => Randomness -> PrivateKey a -> Mu -> Signature a
signExternalMuWith (Randomness rnd) sk (Mu mu) =
    Internal.sigGenMu sk (Builder.public mu) rnd

-- | Verifies a signature.  Returns @True@ when the signature is valid for the
-- message representative µ.
verifyExternalMu :: ParamSet a
                 => PublicKey a -> Mu -> Signature a -> Bool
verifyExternalMu pk (Mu mu) = Internal.sigVerMu pk (Builder.public mu)

-- | Internal procedure @ML-DSA.Sign_internal@.  Generates a signature for a
-- formatted message.  Fresh randomness is acquired from a 'MonadRandom'
-- instance.
signInternal :: (ParamSet a, ByteArrayAccess message, MonadRandom m)
             => PrivateKey a -> message -> m (Signature a)
signInternal sk m' = do
    rnd <- getRandomBytes 32
    return $ Internal.sigGen sk m' (rnd :: ScrubbedBytes)

-- | Internal procedure @ML-DSA.Sign_internal@.  Generates a signature for a
-- formatted message using explicit randomness.
signInternalWith :: (ParamSet a, ByteArrayAccess message)
                 => Randomness -> PrivateKey a -> message -> Signature a
signInternalWith (Randomness rnd) sk m' = Internal.sigGen sk m' rnd

-- | Internal procedure @ML-DSA.Verify_internal@.  Returns @True@ when the
-- signature is valid for the formatted message.
verifyInternal :: (ParamSet a, ByteArrayAccess message)
               => PublicKey a -> message -> Signature a -> Bool
verifyInternal = Internal.sigVer

-- | Class of hash algorithms that can be used with HashML-DSA.
--
-- The algorithm should be chosen to provide enough collision resistance for the
-- target security level, for example SHA-384 or stronger with ML-DSA-65.  This
-- is not enforced by the API.
class HashAlgorithm alg => PreHashAlgorithm alg where
    oid :: proxy alg -> Bytes

instance PreHashAlgorithm SHA256 where
    oid _ = nistHashOID 0x01

instance PreHashAlgorithm SHA384 where
    oid _ = nistHashOID 0x02

instance PreHashAlgorithm SHA512 where
    oid _ = nistHashOID 0x03

instance PreHashAlgorithm SHA224 where
    oid _ = nistHashOID 0x04

instance PreHashAlgorithm SHA512t_224 where
    oid _ = nistHashOID 0x05

instance PreHashAlgorithm SHA512t_256 where
    oid _ = nistHashOID 0x06

instance PreHashAlgorithm SHA3_224 where
    oid _ = nistHashOID 0x07

instance PreHashAlgorithm SHA3_256 where
    oid _ = nistHashOID 0x08

instance PreHashAlgorithm SHA3_384 where
    oid _ = nistHashOID 0x09

instance PreHashAlgorithm SHA3_512 where
    oid _ = nistHashOID 0x0A

instance PreHashAlgorithm (SHAKE128 256) where
    oid _ = nistHashOID 0x0B

instance PreHashAlgorithm (SHAKE256 512) where
    oid _ = nistHashOID 0x0C

nistHashOID :: Word8 -> Bytes
nistHashOID w8 =
    B.pack [ 0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, w8 ]