diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,5 @@ +# Changelog for `mldsa` + +## 0.1.0.0 - 2026-07-05 + +* First version. Released on an unsuspecting world. diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,26 @@ +Copyright 2026 Olivier Chรฉron + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +3. Neither the name of the copyright holder nor the names of its contributors + may be used to endorse or promote products derived from this software + without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND +ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR +ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES +(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; +LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON +ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,55 @@ +# ML-DSA + +[![BSD](https://img.shields.io/badge/License-BSD-blue)](https://en.wikipedia.org/wiki/BSD_licenses) +[![Haskell](https://img.shields.io/badge/Language-Haskell-lightgrey)](https://haskell.org/) +[![Hackage](https://img.shields.io/hackage/v/mldsa?label=Hackage&color=green)](https://hackage.haskell.org/package/mldsa) + +_Module-Lattice-based Digital Signature Algorithm_ implemented in Haskell. + +See [FIPS 204](https://csrc.nist.gov/pubs/fips/204/final). + +Example session: + +```haskell +> :set -XOverloadedStrings +> import Crypto.PubKey.ML_DSA +> import Data.ByteString +> import Data.Proxy +> let params = Proxy :: Proxy ML_DSA_65 +> (pubKey, privKey) <- generate params +> sigma <- sign privKey ("original message" :: ByteString) defaultContext +> verify pubKey ("original message" :: ByteString) sigma defaultContext +True +> verify pubKey ("altered message" :: ByteString) sigma defaultContext +False +``` + +## Notes + +The library does its best to destroy secrets and intermediate buffers from +memory after use, despite the implementation in functional style. This relies +on finalization by the garbage collector and is not guaranteed to run before +the program exits. Also, depending on optimizations applied, lambdas may +capture variables and move them to the heap. This could theoretically include +machine words containing secret information that would not then be destroyed. +Cautious users can run the benchmarks with info-table profiling and verify that +closures containing non pointers capture only non-secret variables like loop +indices or algorithm parameters. + +Best performance is obtained with the LLVM code generator. On ARM, define +macro `__ARM_FEATURE_UNALIGNED` if unaligned access is supported by the target. + +Randomness is provided either from explicit inputs or through a user-selected +instance of the `MonadRandom` type class from `crypton`. A good implementation +would combine multiple sources of entropy, reseed periodically, and protect its +internal state in memory. Deterministic signing with no source of randomness is +possible but is not advised. + +## Testing + +The test suite executes all NIST test vectors but necessary files are not +included in the package to limit its size. Instead, three files are downloaded +from the project repository during execution, and this relies on commands `sh` +and `curl` to run the script `tests/get-vectors.sh`. If not applicable to your +environment, please execute the same steps manually. It will be needed only +the first time. diff --git a/benchs/Bench.hs b/benchs/Bench.hs new file mode 100644 --- /dev/null +++ b/benchs/Bench.hs @@ -0,0 +1,67 @@ +{-# LANGUAGE ExistentialQuantification #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE OverloadedStrings #-} +module Main (main) where + +import Criterion.Main + +import Crypto.Random +import Crypto.PubKey.ML_DSA + +import Data.ByteArray (Bytes) +import Data.Maybe (fromJust) +import Data.Proxy + +data PS = forall a . ParamSet a => PS (Proxy a) + +message :: Bytes +message = mempty + +psList :: [(String, PS)] +psList = + [ ("ML-DSA-44", PS (Proxy :: Proxy ML_DSA_44)) + , ("ML-DSA-65", PS (Proxy :: Proxy ML_DSA_65)) + , ("ML-DSA-87", PS (Proxy :: Proxy ML_DSA_87)) + ] + +doBench :: (String, PS) -> Benchmark +doBench (name, PS p) = bgroup name + [ bench "generate" $ perRunEnv setupGenerate (return . runGenerate) + , bench "sign" $ perRunEnv setupSign runSign + , bench "sign (batch)" $ perBatchEnv (const setupSign) runSign + , bench "verify" $ perRunEnv setupVerify (return . runVerify) + , bench "verify (batch)" $ perBatchEnv (const setupVerify) (return . runVerify) + , env (generate p) $ bench "encoding (public)" . nf (encodeDecode p) . fst + , env (generate p) $ bench "encoding (private)" . nf (encodeDecode p) . snd + , env setupVerify $ bench "encoding (signature)" . nf (encodeDecode p) . snd + ] + where + gen32 = getRandomBytes 32 :: IO Bytes + + setupGenerate = gen32 + runGenerate = fromJust . generateWith p + + setupSign = snd <$> generate p + runSign sk = do + -- signature time is highly variable due to the rejection loop, so we + -- want diversity in batch runs and use fresh randomness + rnd <- fromJust . randomness <$> gen32 + return $ signWith rnd sk message defaultContext + + runVerify (pk, sig) = verify pk message sig defaultContext + setupVerify = do + (pk, sk) <- generate p + sig <- sign sk message defaultContext + return (pk, sig) + + encodeDecode :: (Encode obj, Decode obj, Eq (obj a), ParamSet a) + => proxy a -> obj a -> Bool + encodeDecode prx k = decode prx (aligned $ encode k) == Just k + + aligned :: Bytes -> Bytes + aligned = id + +main :: IO () +main = defaultMain + [ bgroup "mldsa" $ map doBench psList + ] diff --git a/mldsa.cabal b/mldsa.cabal new file mode 100644 --- /dev/null +++ b/mldsa.cabal @@ -0,0 +1,206 @@ +cabal-version: 2.2 + +-- This file has been generated from package.yaml by hpack version 0.39.1. +-- +-- see: https://github.com/sol/hpack + +name: mldsa +version: 0.1.0.0 +synopsis: Module-Lattice-based Digital Signature Algorithm +description: Module-Lattice-based Digital Signature Algorithm (ML-DSA) implemented in + Haskell. +category: Cryptography +homepage: https://codeberg.org/ocheron/hs-mldsa#readme +bug-reports: https://codeberg.org/ocheron/hs-mldsa/issues +author: Olivier Chรฉron +maintainer: olivier.cheron@gmail.com +copyright: 2026 Olivier Chรฉron +license: BSD-3-Clause +license-file: LICENSE +build-type: Simple +extra-source-files: + README.md + tests/get-vectors.sh +extra-doc-files: + CHANGELOG.md + +source-repository head + type: git + location: https://codeberg.org/ocheron/hs-mldsa + +flag use_crypton + description: Use crypton instead of cryptonite + manual: True + default: True + +library + exposed-modules: + Crypto.PubKey.ML_DSA + other-modules: + Auxiliary + Base + Block + BlockN + Builder + ByteArrayST + Crypto + Equality + Fusion + Internal + Iterate + Machine + Marking + Math + Matrix + ScrubbedBlock + SecureBlock + SecureBytes + Vector + Paths_mldsa + autogen-modules: + Paths_mldsa + hs-source-dirs: + src + ghc-options: -Wall -Wcompat -Widentities -Wincomplete-record-updates -Wincomplete-uni-patterns -Wmissing-export-lists -Wmissing-home-modules -Wpartial-fields -Wredundant-constraints -Wno-unticked-promoted-constructors -O2 + build-depends: + base >=4.7 && <5 + , deepseq + , primitive >=0.7.2 + default-language: Haskell2010 + if flag(use_crypton) + build-depends: + crypton >=1.1.1 + , ram + else + build-depends: + cryptonite >=0.26 + , memory + +test-suite mldsa-test + type: exitcode-stdio-1.0 + main-is: Tests.hs + other-modules: + KeyGen + SigExt + SigGen + SigVer + Util + Vectors + Paths_mldsa + autogen-modules: + Paths_mldsa + hs-source-dirs: + tests + ghc-options: -Wall -Wcompat -Widentities -Wincomplete-record-updates -Wincomplete-uni-patterns -Wmissing-export-lists -Wmissing-home-modules -Wpartial-fields -Wredundant-constraints -Wno-unticked-promoted-constructors -threaded -rtsopts -with-rtsopts=-N + build-depends: + aeson + , base >=4.7 && <5 + , bytestring + , deepseq + , directory + , filelock + , mldsa + , primitive >=0.7.2 + , process + , tasty + , tasty-hunit + , tasty-quickcheck + , text + , zlib + default-language: Haskell2010 + if flag(use_crypton) + build-depends: + crypton >=1.1.1 + , ram + else + build-depends: + cryptonite >=0.26 + , memory + +test-suite mldsa-test-full + type: exitcode-stdio-1.0 + main-is: Tests.hs + other-modules: + Auxiliary + Base + Block + BlockN + Builder + ByteArrayST + Crypto + Crypto.PubKey.ML_DSA + Equality + Fusion + Internal + Iterate + Machine + Marking + Math + Matrix + ScrubbedBlock + SecureBlock + SecureBytes + Vector + KeyGen + SigExt + SigGen + SigVer + Util + Vectors + Paths_mldsa + autogen-modules: + Paths_mldsa + hs-source-dirs: + src + tests + ghc-options: -Wall -Wcompat -Widentities -Wincomplete-record-updates -Wincomplete-uni-patterns -Wmissing-export-lists -Wmissing-home-modules -Wpartial-fields -Wredundant-constraints -Wno-unticked-promoted-constructors -fno-ignore-asserts -threaded -rtsopts -with-rtsopts=-N + cpp-options: -DML_DSA_TESTING + build-depends: + aeson + , base >=4.7 && <5 + , bytestring + , deepseq + , directory + , filelock + , primitive >=0.7.2 + , process + , tasty + , tasty-hunit + , tasty-quickcheck + , text + , zlib + default-language: Haskell2010 + if flag(use_crypton) + build-depends: + crypton >=1.1.1 + , ram + else + build-depends: + cryptonite >=0.26 + , memory + +benchmark mldsa-bench + type: exitcode-stdio-1.0 + main-is: Bench.hs + other-modules: + Paths_mldsa + autogen-modules: + Paths_mldsa + hs-source-dirs: + benchs + ghc-options: -Wall -Wcompat -Widentities -Wincomplete-record-updates -Wincomplete-uni-patterns -Wmissing-export-lists -Wmissing-home-modules -Wpartial-fields -Wredundant-constraints -Wno-unticked-promoted-constructors -threaded -rtsopts -with-rtsopts=-N -with-rtsopts=-A48m + build-depends: + base >=4.7 && <5 + , criterion + , deepseq + , mldsa + , primitive >=0.7.2 + default-language: Haskell2010 + if flag(use_crypton) + build-depends: + crypton >=1.1.1 + , ram + else + build-depends: + cryptonite >=0.26 + , memory diff --git a/src/Auxiliary.hs b/src/Auxiliary.hs new file mode 100644 --- /dev/null +++ b/src/Auxiliary.hs @@ -0,0 +1,959 @@ +-- | +-- Module : Auxiliary +-- License : BSD-3-Clause +-- Copyright : (c) 2026 Olivier Chรฉron +-- +-- ML-DSA auxiliary functions +-- +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UnboxedTuples #-} +module Auxiliary + ( Zq, Rq, Tq + , Norm(..), normFrom, countFrom, ntt, nttInv + , simpleBitPack, simpleBitPack10, simpleBitUnpack10 + , bitPack, bitUnpack, bitPackSafe, bitUnpackSafe + , hintBitPack, hintBitUnpack + , sampleInBall, rejNttPoly, rejBoundedPoly + , powerTwoRound + , highBits, lowBits + , Hints, makeHint, useHint +#ifdef ML_DSA_TESTING + , Mq, toMontgomery, fromMontgomery, norm + , bitRev8, fromZq, toZq, fromCoeffs, toCoeffs, fromBools, toBools + , powerTwoRoundZq, decomposeZq + , simpleBitUnpack +#endif + ) where + +import Crypto.Hash.Algorithms + +import Data.ByteArray (ByteArrayAccess) +import qualified Data.ByteArray as B +import qualified Data.Memory.Endian as B + +import Data.Primitive.Types (Prim(..)) + +import Control.DeepSeq (NFData(..)) +import Control.Monad +import Control.Monad.ST + +import Data.Bits +import Data.Int +import Data.Proxy +import Data.Word + +#if MIN_VERSION_base(4,16,0) +import GHC.Exts as Exts (eqWord32#, int2Word#, wordToWord8#) +#else +import GHC.Exts as Exts (eqWord#, int2Word#) +#endif +import GHC.TypeNats +import GHC.Word (Word8(W8#), Word32(W32#)) + +import Foreign.Ptr (Ptr, plusPtr) +import Foreign.Storable (pokeByteOff) + +import Base +import Block (blockIndex) +import BlockN (BlockN, MutableBlockN) +import Builder (Builder) +import Crypto (BlockDigest, Digest) +import Fusion +import Machine +import Marking (Classified, SecurityMarking(..)) +import SecureBlock (SecureBlock) +import SecureBytes (SecureBytes) +import Vector (Vector) +import qualified BlockN +import qualified Builder +import qualified ByteArrayST as ST +import qualified Crypto +import qualified Vector +import Math + +type N = 256 + +n :: Int +n = 256 + +q :: Integer +q = 8380417 + +q32 :: Word32 +q32 = fromInteger q + +q64 :: Word64 +q64 = fromInteger q + +bitRev8 :: Word8 -> Word8 +bitRev8 = fromIntegral . rev8 . fromIntegral + where + rev1, rev2, rev4, rev8 :: Word -> Word + rev1 b = b .&. 1 + rev2 b = rev1 (b `unsafeShiftR` 1) .|. rev1 b `unsafeShiftL` 1 + rev4 b = rev2 (b `unsafeShiftR` 2) .|. rev2 b `unsafeShiftL` 2 + rev8 b = rev4 (b `unsafeShiftR` 4) .|. rev4 b `unsafeShiftL` 4 + +unsafeShiftIR :: Word32 -> Int -> Word32 +unsafeShiftIR x s = fromIntegral ((fromIntegral x :: Int32) `unsafeShiftR` s) +{-# INLINE unsafeShiftIR #-} + +select32 :: Word32 -> Word32 -> Word32 -> Word32 +select32 mask yes no = (mask .&. yes) .|. (complement mask .&. no) +{-# INLINE select32 #-} + +max32 :: Word32 -> Word32 -> Word32 +max32 a b = let cond = (b - a) `unsafeShiftIR` 31 in select32 cond a b +{-# INLINE max32 #-} + +min32 :: Word32 -> Word32 -> Word32 +min32 a b = let cond = (a - b) `unsafeShiftIR` 31 in select32 cond a b +{-# INLINE min32 #-} + +-- Reduction ๐‘ฅ mod 5 for 0 โ‰ค ๐‘ฅ < 15 +mod5 :: Word32 -> Word32 +mod5 x = let quotient = (x * 205) `unsafeShiftR` 10 in x - 5 * quotient + +-- Reduction ๐‘ฅ mod ๐‘ž for 0 โ‰ค ๐‘ฅ < 2๐‘ž +reduceSimple :: Word32 -> Word32 +reduceSimple x = select32 mask x subtracted + where + subtracted = x - q32 + mask = subtracted `unsafeShiftIR` 31 +{-# INLINE reduceSimple #-} + +#ifdef ML_DSA_TESTING +-- Reduction ๐‘ฅ mod ๐‘ž for 0 โ‰ค ๐‘ฅ < 1025๐‘ž +reduce :: Word64 -> Word32 +reduce x = reduceSimple $ fromIntegral (x - t * q64) + where t = x `unsafeShiftR` 23 +{-# INLINE reduce #-} +#endif + +-- Computes ๐‘Ž . 2^-32 mod ๐‘ž +reduceMontgomery :: Word64 -> Word32 +reduceMontgomery a = reduceSimple (fromIntegral r) + where + t = (fromIntegral a :: Word32) * 4236238847 -- inverse of -๐‘ž modulo 2^32 + r = (a + fromIntegral t * q64) `unsafeShiftR` 32 +{-# INLINE reduceMontgomery #-} + +newtype Zq = Zq Word32 +#ifdef ML_DSA_TESTING + deriving (Eq, Show) +#else + deriving Eq +#endif + +instance Prim Zq where + sizeOf# (Zq a) = sizeOf# a + {-# INLINE sizeOf# #-} + alignment# (Zq a) = alignment# a + {-# INLINE alignment# #-} +#if MIN_VERSION_primitive(0,9,0) + sizeOfType# _ = sizeOfType# (Proxy :: Proxy Word32) + {-# INLINE sizeOfType# #-} + alignmentOfType# _ = alignmentOfType# (Proxy :: Proxy Word32) + {-# INLINE alignmentOfType# #-} +#endif + indexByteArray# ba i = Zq (indexByteArray# ba i) + {-# INLINE indexByteArray# #-} + readByteArray# mba i s = + case readByteArray# mba i s of + (# s', a #) -> (# s', Zq a #) + {-# INLINE readByteArray# #-} + writeByteArray# mba i (Zq a) = writeByteArray# mba i a + {-# INLINE writeByteArray# #-} + setByteArray# mba i len (Zq a) = setByteArray# mba i len a + {-# INLINE setByteArray# #-} + indexOffAddr# addr i = Zq (indexOffAddr# addr i) + {-# INLINE indexOffAddr# #-} + readOffAddr# addr i s = + case readOffAddr# addr i s of + (# s', a #) -> (# s', Zq a #) + {-# INLINE readOffAddr# #-} + writeOffAddr# addr i (Zq a) = writeOffAddr# addr i a + {-# INLINE writeOffAddr# #-} + setOffAddr# addr i len (Zq a) = setOffAddr# addr i len a + {-# INLINE setOffAddr# #-} + +instance PrimSized Zq where + type PrimSize Zq = 4 + +instance Add Zq where + zero = Zq 0 + Zq a .+ Zq b = Zq $ reduceSimple (a + b) + Zq a .- Zq b = Zq $ reduceSimple (a + q32 - b) + neg (Zq a) = Zq $ reduceSimple (q32 - a) + +#ifdef ML_DSA_TESTING +instance Mul Zq where + one = Zq 1 + a .* b = fromMontgomery (toMontgomery a .* toMontgomery b) + +toZq :: Word32 -> Zq +toZq = Zq . reduce . fromIntegral +#endif + +fromZq :: Zq -> Word32 +fromZq (Zq a) = a + +newtype Mq = Mq Word32 +#ifdef ML_DSA_TESTING + deriving (Eq, Show) +#else + deriving Eq +#endif + +instance Prim Mq where + sizeOf# (Mq a) = sizeOf# a + {-# INLINE sizeOf# #-} + alignment# (Mq a) = alignment# a + {-# INLINE alignment# #-} +#if MIN_VERSION_primitive(0,9,0) + sizeOfType# _ = sizeOfType# (Proxy :: Proxy Word32) + {-# INLINE sizeOfType# #-} + alignmentOfType# _ = alignmentOfType# (Proxy :: Proxy Word32) + {-# INLINE alignmentOfType# #-} +#endif + indexByteArray# ba i = Mq (indexByteArray# ba i) + {-# INLINE indexByteArray# #-} + readByteArray# mba i s = + case readByteArray# mba i s of + (# s', a #) -> (# s', Mq a #) + {-# INLINE readByteArray# #-} + writeByteArray# mba i (Mq a) = writeByteArray# mba i a + {-# INLINE writeByteArray# #-} + setByteArray# mba i len (Mq a) = setByteArray# mba i len a + {-# INLINE setByteArray# #-} + indexOffAddr# addr i = Mq (indexOffAddr# addr i) + {-# INLINE indexOffAddr# #-} + readOffAddr# addr i s = + case readOffAddr# addr i s of + (# s', a #) -> (# s', Mq a #) + {-# INLINE readOffAddr# #-} + writeOffAddr# addr i (Mq a) = writeOffAddr# addr i a + {-# INLINE writeOffAddr# #-} + setOffAddr# addr i len (Mq a) = setOffAddr# addr i len a + {-# INLINE setOffAddr# #-} + +instance PrimSized Mq where + type PrimSize Mq = 4 + +instance Add Mq where + zero = Mq 0 + Mq a .+ Mq b = Mq $ reduceSimple (a + b) + Mq a .- Mq b = Mq $ reduceSimple (a + q32 - b) + neg (Mq a) = Mq $ reduceSimple (q32 - a) + +instance Mul Mq where + one = Mq 4193792 -- toMongomery (Zq 1) + Mq a .* Mq b = Mq $ reduceMontgomery (fromIntegral a * fromIntegral b) + +instance MulAdd Mq where + mulAdd a b c = a .* b .+ c + +#ifdef ML_DSA_TESTING +instance BiMul Mq Mq where + (..*) = (.*) + +instance BiMulAdd Mq Mq where + biMulAdd = mulAdd +#endif + +toMontgomery :: Zq -> Mq +toMontgomery (Zq a) = Mq $ reduceMontgomery $ fromIntegral a * 2365951 -- R^2 mod ๐‘ž + +fromMontgomery :: Mq -> Zq +fromMontgomery (Mq m) = Zq $ reduceMontgomery $ fromIntegral m + +newtype Rq marking = Rq (BlockN marking N Zq) +#ifdef ML_DSA_TESTING + deriving (Eq, Show, NFData) +#else + deriving NFData +#endif + +instance Classified marking => Add (Rq marking) where + zero = Rq zero + Rq a .+ Rq b = Rq (a .+ b) + {-# INLINE (.+) #-} + Rq a .- Rq b = Rq (a .- b) + {-# INLINE (.-) #-} + neg (Rq a) = Rq (neg a) + {-# INLINE neg #-} + +instance Crypto.ConstEqW (Rq Sec) where + constEqW (Rq a) (Rq b) = Crypto.constEqW + (BlockN.unsafeCast a :: SecureBlock Sec Word) + (BlockN.unsafeCast b :: SecureBlock Sec Word) + +instance Crypto.ConstEqW (Rq Pub) where + constEqW (Rq a) (Rq b) = Crypto.constEqW + (BlockN.unsafeCast a :: SecureBlock Pub Word) + (BlockN.unsafeCast b :: SecureBlock Pub Word) + +#ifdef ML_DSA_TESTING +fromCoeffs :: [Zq] -> Maybe (Rq Sec) +fromCoeffs = fmap Rq . BlockN.fromList + +toCoeffs :: Rq Sec -> [Zq] +toCoeffs (Rq a) = BlockN.toList a +#endif + +newtype Tq marking = Tq (BlockN marking N Mq) +#ifdef ML_DSA_TESTING + deriving (Eq, Show, NFData) +#else + deriving NFData +#endif + +instance Classified marking => Add (Tq marking) where + zero = Tq zero + Tq a .+ Tq b = Tq (a .+ b) + {-# INLINE (.+) #-} + Tq a .- Tq b = Tq (a .- b) + {-# INLINE (.-) #-} + neg (Tq a) = Tq (neg a) + {-# INLINE neg #-} + +instance Crypto.ConstEqW (Tq Pub) where + constEqW (Tq a) (Tq b) = Crypto.constEqW + (BlockN.unsafeCast a :: SecureBlock Pub Word) + (BlockN.unsafeCast b :: SecureBlock Pub Word) + +instance BiMul (Tq marking) (Tq Sec) where + Tq a ..* Tq b = Tq $ BlockN.zipWithEqPrimSizeR (.*) a b + {-# INLINE (..*) #-} + +instance BiMulAdd (Tq marking) (Tq Sec) where + biMulFold = multiplyNTTsFold + {-# INLINE biMulFold #-} + +#ifdef ML_DSA_TESTING +instance Mul (Tq Sec) where + one = Tq $ BlockN.replicate one + (.*) = (..*) + +instance MulAdd (Tq Sec) where + mulAdd = biMulAdd +#endif + +newtype Norm = Norm { getNorm :: Word32 } +#ifdef ML_DSA_TESTING + deriving (Eq, Ord, Show, Num) +#endif + +instance Semigroup Norm where + Norm a <> Norm b = Norm (max32 a b) + +instance Monoid Norm where + mempty = Norm 0 + +#ifdef ML_DSA_TESTING +norm :: Rq marking -> Norm +norm = normFrom mempty +#endif + +normFrom :: Norm -> Rq marking -> Norm +normFrom (Norm m0) (Rq w) = Norm $ BlockN.foldl' f m0 w + where + absZq (Zq z) = min32 z (q32 - z) + f m = max32 m . absZq +{-# INLINE normFrom #-} + +countFrom :: Word -> Hints -> Word +countFrom b (Hints w) = let f acc (H h) = acc + fromIntegral h in BlockN.foldl' f b w +{-# INLINE countFrom #-} + +-- Computes the NTT representation of the given polynomial +ntt :: Classified marking => Rq marking -> Tq marking +ntt (Rq a) = Tq $ BlockN.runThaw b mutNtt + where b = BlockN.mapEqPrimSize toMontgomery a +{-# INLINE ntt #-} + +mutNtt :: MutableBlockN marking N Mq s -> ST s () +mutNtt !b = outer 1 128 + where + outer !m len = when (len >= 1) $ inner m len 0 + + inner !m !len start + | start < 256 = do + let z = BlockN.index zetaPowBitRev m -- 1753 ^ bitRev8 m + loop z (start + len) len start + inner (m + 1) len (start + offsetShiftL 1 len) + | otherwise = outer m (offsetShiftR 1 len) + + loop !z end len j = + when (j < end) $ do + t <- (z .*) <$> BlockN.read b (j + len) + x <- BlockN.read b j + BlockN.write b (j + len) (x .- t) + BlockN.write b j (x .+ t) + loop z end len (j + 1) +{-# NOINLINE mutNtt #-} + +-- Computes the polynomial that corresponds to the given NTT representation +nttInv :: Classified marking => Tq marking -> Rq marking +nttInv (Tq a) = Rq $ BlockN.mapEqPrimSize fromMontgomery $ BlockN.runThaw a mutNttInv +{-# INLINE nttInv #-} + +mutNttInv :: MutableBlockN marking N Mq s -> ST s () +mutNttInv !b = do + outer 255 1 + BlockN.iterModify (\x -> x .* Mq 16382) b -- toMontgomery (Zq 8347681) + where + outer !m len = when (len < 256) $ inner m len 0 + + inner !m !len start + | start < 256 = do + let z = BlockN.index zetaPowBitRev m -- 1753 ^ bitRev8 m + loop z (start + len) len start + inner (m - 1) len (start + offsetShiftL 1 len) + | otherwise = outer m (offsetShiftL 1 len) + + loop !z end len j = + when (j < end) $ do + t <- BlockN.read b j + x <- BlockN.read b (j + len) + BlockN.write b j (t .+ x) + BlockN.write b (j + len) (z .* (x .- t)) + loop z end len (j + 1) +{-# NOINLINE mutNttInv #-} + +multiplyNTTsFold :: Foldable t => Tq Sec -> t (Tq marking, Tq Sec) -> Tq Sec +multiplyNTTsFold (Tq c) = Tq . BlockN.runFold c (uncurry multiplyNTTsAdd) +{-# INLINE multiplyNTTsFold #-} + +-- Multiply then add a third term +multiplyNTTsAdd :: Tq marking -> Tq Sec -> MutableBlockN Sec N Mq s -> ST s () +multiplyNTTsAdd (Tq !f) (Tq !g) bb = loop bb 0 + where + loop :: MutableBlockN Sec N Mq s -> Offset Mq -> ST s () + loop !b i = when (i < 256) $ do + c <- BlockN.read b i + BlockN.write b i $ mulAdd (BlockN.index f i) (BlockN.index g i) c + loop b (i + 1) + +-- Values of 1753 ^ BitRev8(๐‘–) mod ๐‘ž for ๐‘– โˆˆ {0, โ€ฆ , 255} +zetaPowBitRev :: BlockN Pub 256 Mq +zetaPowBitRev = BlockN.runNew (Proxy :: Proxy Pub) $ \out -> + foldM_ (loop out) one offsets + where + offsets = Prelude.map (fromIntegral . bitRev8) [0 .. 255] + loop b acc i = BlockN.write b i acc >> return (Mq 2091667 .* acc) + -- toMontgomery (Zq 1753) + +newtype H = H Word8 +#ifdef ML_DSA_TESTING + deriving (Eq, Show) +#else + deriving Eq +#endif + +instance Prim H where + sizeOf# (H a) = sizeOf# a + {-# INLINE sizeOf# #-} + alignment# (H a) = alignment# a + {-# INLINE alignment# #-} +#if MIN_VERSION_primitive(0,9,0) + sizeOfType# _ = sizeOfType# (Proxy :: Proxy Word8) + {-# INLINE sizeOfType# #-} + alignmentOfType# _ = alignmentOfType# (Proxy :: Proxy Word8) + {-# INLINE alignmentOfType# #-} +#endif + indexByteArray# ba i = H (indexByteArray# ba i) + {-# INLINE indexByteArray# #-} + readByteArray# mba i s = + case readByteArray# mba i s of + (# s', a #) -> (# s', H a #) + {-# INLINE readByteArray# #-} + writeByteArray# mba i (H a) = writeByteArray# mba i a + {-# INLINE writeByteArray# #-} + setByteArray# mba i len (H a) = setByteArray# mba i len a + {-# INLINE setByteArray# #-} + indexOffAddr# addr i = H (indexOffAddr# addr i) + {-# INLINE indexOffAddr# #-} + readOffAddr# addr i s = + case readOffAddr# addr i s of + (# s', a #) -> (# s', H a #) + {-# INLINE readOffAddr# #-} + writeOffAddr# addr i (H a) = writeOffAddr# addr i a + {-# INLINE writeOffAddr# #-} + setOffAddr# addr i len (H a) = setOffAddr# addr i len a + {-# INLINE setOffAddr# #-} + +instance PrimSized H where + type PrimSize H = 1 + +-- boolean-less comparison, to avoid code duplication with pattern matching +notEqualH :: Word32 -> Word32 -> H +notEqualH (W32# x#) (W32# y#) = H $ +#if MIN_VERSION_base(4,16,0) + 1 - W8# (Exts.wordToWord8# (Exts.int2Word# (Exts.eqWord32# x# y#))) +#else + 1 - W8# (Exts.int2Word# (Exts.eqWord# x# y#)) +#endif + +newtype Hints = Hints (BlockN Sec N H) +#ifdef ML_DSA_TESTING + deriving (Eq, Show, NFData) +#else + deriving NFData +#endif + +instance Crypto.ConstEqW Hints where + constEqW (Hints a) (Hints b) = Crypto.constEqW + (BlockN.unsafeCast a :: SecureBlock Sec Word) + (BlockN.unsafeCast b :: SecureBlock Sec Word) + +#ifdef ML_DSA_TESTING +fromBools :: [Bool] -> Maybe Hints +fromBools = fmap Hints . BlockN.fromList . map (H . fromIntegral . fromEnum) + +toBools :: Hints -> [Bool] +toBools (Hints a) = map (\(H h) -> toEnum (fromIntegral h)) (BlockN.toList a) +#endif + +peekWordPos :: Ptr WordLE -> BitPos -> ST s WordM +peekWordPos a bp = fromLE <$> ST.peekElemOff a (wordOff bp) + +pokeWordPos :: Ptr WordLE -> BitPos -> WordM -> ST s () +pokeWordPos a bp = ST.pokeElemOff a (wordOff bp) . toLE + +newtype BitPos = BitPos Int + +zeroPos :: BitPos +zeroPos = BitPos 0 + +wordOff :: BitPos -> Int +wordOff (BitPos p) = div p wordBits + +bitPos :: BitPos -> Int +bitPos (BitPos p) = p .&. (wordBits - 1) + +availPos :: Int -> BitPos -> Int +availPos requested (BitPos p) = min available requested + where available = wordBits - (p .&. (wordBits - 1)) + +nextPos :: Int -> BitPos -> (Int, BitPos) +nextPos requested (BitPos p) = (howMany, BitPos $ p + howMany) + where howMany = availPos requested (BitPos p) + +getMask :: Int -> WordM +getMask howMany + | howMany >= wordBits = maxBound + | otherwise = (1 `unsafeShiftL` howMany) - 1 + -- branch useful only when processing one byte at a time due to + -- architecture not supporting unaligned memory access + +-- Encodes an array of ๐‘‘-bit integers for 1 โ‰ค ๐‘‘ โ‰ค 22 +simpleBitPack :: Int -> BlockN marking N Word32 -> Builder marking +simpleBitPack d w = Builder.create (32 * d) (runSimpleBitPack d w) +{-# INLINE simpleBitPack #-} + +runSimpleBitPack :: Int -> BlockN marking N Word32 -> Ptr WordLE -> ST s () +runSimpleBitPack !d !w dst = loop dst 0 zeroPos 0 (get 0) d + where + get = BlockN.index w . Offset + {-# INLINE get #-} + + loop !b !pos !bp !o !a j + | j == 0, pos' == n = return () + | j == 0 = loop b pos' bp o (get pos') d + | bitPos bp + howMany < wordBits = loop b pos bp' o' a' j' + | otherwise = pokeWordPos b bp o' >> loop b pos bp' 0 a' j' + where + pos' = pos + 1 + (howMany, bp') = nextPos j bp + x = fromIntegral a .&. getMask howMany + o' = o .|. (x `unsafeShiftL` bitPos bp) + a' = a `unsafeShiftR` howMany + j' = j - howMany + +-- simpleBitPack with ๐‘‘=10 after conversion from a polynomial and division by 2^13 +simpleBitPack10 :: Rq Pub -> Builder Pub +simpleBitPack10 (Rq w) = simpleBitPack 10 (BlockN.mapEqPrimSize f w) + where f (Zq z) = z `unsafeShiftR` 13 +{-# INLINE simpleBitPack10 #-} + +-- Decodes an array of ๐‘‘-bit integers for 1 โ‰ค ๐‘‘ โ‰ค 22 +simpleBitUnpack :: (Classified marking, ByteArrayAccess ba) => Int -> ba -> BlockN marking N Word32 +simpleBitUnpack d b = BlockN.runNew Proxy $ mutSimpleBitUnpack d b +{-# INLINE simpleBitUnpack #-} + +mutSimpleBitUnpack :: ByteArrayAccess ba => Int -> ba -> MutableBlockN marking N Word32 s -> ST s () +mutSimpleBitUnpack !d b !w = ST.withByteArray b $ \p -> outer p zeroPos 0 + where + outer !p !bp i = when (i < Offset n) $ inner p i bp 0 0 + + inner !p !i !bp !v j + | j == d = BlockN.write w i v >> outer p bp (i + 1) + | otherwise = do + let (howMany, bp') = nextPos (d - j) bp + y <- get p bp howMany + let v' = v .|. (fromIntegral y `unsafeShiftL` j) + j' = j + howMany + inner p i bp' v' j' + + get :: Ptr WordLE -> BitPos -> Int -> ST s WordM + get p bp howMany = do + x <- (`unsafeShiftR` bitPos bp) <$> peekWordPos p bp + return (x .&. getMask howMany) + +-- simpleBitUnpack with ๐‘‘=10 and conversion to a polynomial with multiplication by 2^13 +simpleBitUnpack10 :: ByteArrayAccess ba => ba -> Rq Pub +simpleBitUnpack10 = Rq . BlockN.mapEqPrimSize f . simpleBitUnpack 10 + where f z = Zq (z `unsafeShiftL` 13) +{-# INLINE simpleBitUnpack10 #-} + +-- bitPack when upper bound is a power of two +bitPackSafe :: Classified marking => Int -> Rq marking -> Builder marking +bitPackSafe d = simpleBitPack d . unmirror (1 `unsafeShiftL` (d - 1)) +{-# INLINE bitPackSafe #-} + +-- bitUnpack when upper bound is a power of two +bitUnpackSafe :: (Classified marking, ByteArrayAccess ba) => Int -> ba -> Rq marking +bitUnpackSafe d b = mirror (1 `unsafeShiftL` (d - 1)) (simpleBitUnpack d b) +{-# INLINE bitUnpackSafe #-} + +-- Encodes a polynomial ๐‘ค into a byte string +bitPack :: Word32 -> Int -> Rq Sec -> Builder Sec +bitPack !m d = simpleBitPack d . unmirror m + +-- Reverses the procedure bitPack +bitUnpack :: ByteArrayAccess ba => Word32 -> Int -> ba -> Maybe (Rq Sec) +bitUnpack !m d b + | valid unpacked = Just $ mirror m unpacked + | otherwise = Nothing + where + unpacked = simpleBitUnpack d b + valid = Crypto.toBool . BlockN.foldl' f Crypto.trueW + f acc z = acc `Crypto.andW` inRange z + inRange z = Crypto.ltW (fromIntegral z) (fromIntegral (2 * m + 1)) + +mirror :: Classified marking => Word32 -> BlockN marking N Word32 -> Rq marking +mirror m w = Rq $ BlockN.mapEqPrimSize (\z -> Zq m .- Zq z) w +{-# INLINE mirror #-} + +unmirror :: Classified marking => Word32 -> Rq marking -> BlockN marking N Word32 +unmirror m (Rq w) = BlockN.mapEqPrimSize (\z -> fromZq (Zq m .- z)) w +{-# INLINE unmirror #-} + +-- Encodes a polynomial vector ๐ก with binary coefficients into a byte string +hintBitPack :: forall k. KnownNat k => Int -> Vector k Hints -> Builder Sec +hintBitPack !omega !h = Builder.create (omega + k) $ outer 0 0 + where + k = fromIntegral $ natVal (Proxy :: Proxy k) + + outer :: Int -> Offset Word8 -> Ptr Word8 -> ST s () + outer i !idx !p + | i == k = end p idx + | otherwise = inner p i (Vector.index h $ Offset i) idx 0 + + inner :: Ptr Word8 -> Int -> Hints -> Offset Word8 -> Int -> ST s () + inner !p !i e@(Hints !w) !idx j + | j == n = mark p i idx >> outer (i + 1) idx p + | BlockN.index w (Offset j) == H 0 = inner p i e idx (j + 1) + | otherwise = set p idx j >> inner p i e (idx + 1) (j + 1) + + set :: Ptr Word8 -> Offset Word8 -> Int -> ST s () + set p (Offset idx) j = ST.pokeElemOff p idx (fromIntegral j) + + mark :: Ptr Word8 -> Int -> Offset Word8 -> ST s () + mark p i (Offset idx) = ST.pokeElemOff p (omega + i) (fromIntegral idx) + + end :: Ptr Word8 -> Offset Word8 -> ST s () + end !p (Offset idx) = + when (idx < omega) $ ST.fillBytes (p `plusPtr` idx) 0 (omega - idx) +{-# INLINE hintBitPack #-} + +-- Reverses the procedure hintBitPack +hintBitUnpack :: (KnownNat k, ByteArrayAccess ba) => Int -> ba -> Maybe (Vector k Hints) +hintBitUnpack !omega !y = Vector.createMaybeAccum outer cont 0 + where + outer :: Int -> Offset Hints -> Maybe (Int, Hints) + outer idx (Offset i) = runST $ ST.withByteArray y $ \p -> do + m <- fromIntegral <$> ST.peekElemOff p (omega + i) + if m < idx || m > omega + then return Nothing + else new >>= inner p m idx + + inner :: Ptr Word8 -> Int -> Int -> MutableBlockN Sec N H s -> ST s (Maybe (Int, Hints)) + inner !p !m !first !mb = go first + where + go idx + | idx >= m = unsafeFreezeF mb >>= \b -> + return $ Just (idx, Hints b) + | idx > first = do + val <- ST.peekElemOff p idx + val' <- ST.peekElemOff p (idx - 1) + if val' >= val + then return Nothing + else set mb val >> go (idx + 1) + | otherwise = ST.peekElemOff p idx >>= \val -> + set mb val >> go (idx + 1) + + new :: ST s (MutableBlockN Sec N H s) + new = newF >>= \mb -> BlockN.erase mb >> return mb + + set :: MutableBlockN Sec N H s -> Word8 -> ST s () + set mb val = BlockN.write mb (fromIntegral val) (H 1) + + cont :: forall r. Int -> r -> Maybe r + cont idx h = runST $ ST.withByteArray y $ \p -> go p idx + where + go :: Ptr Word8 -> Int -> ST s (Maybe r) + go !p i + | i >= omega = return (Just h) + | otherwise = do + val <- ST.peekElemOff p i + case val of + 0 -> go p (i + 1) + _ -> return Nothing +{-# INLINE hintBitUnpack #-} + +-- Samples a polynomial ๐‘ โˆˆ ๐‘… with coefficients from {-1, 0, 1} and +-- Hamming weight ๐œ โ‰ค 64 +sampleInBall :: Int -> SecureBytes Sec -> Rq Sec +sampleInBall tau rho = Rq $ + BlockN.runNew (Proxy :: Proxy Sec) $ \c -> runXof c 136 + where + runXof :: MutableBlockN Sec N Zq s -> Int -> ST s () + runXof !c !xofLen = case someNatVal (fromIntegral (8 * xofLen)) of + SomeNat proxy -> do + let bytes = Crypto.unDigest (doHash proxy) + BlockN.erase c + loop c xofLen bytes 8 (256 - tau) + + loop :: MutableBlockN Sec N Zq s -> Int -> SecureBytes Sec -> Int -> Int -> ST s () + loop !c !xofLen !bytes !pos i + | i == n = return () + | pos >= xofLen = runXof c (xofLen + 136) + | j > Offset i = loop c xofLen bytes (pos + 1) i + | otherwise = do + BlockN.read c j >>= BlockN.write c (Offset i) + let idx = i + tau - 256 + iByte = idx `unsafeShiftR` 3 + iBit = idx .&. 7 + b = B.index bytes iByte `unsafeShiftR` iBit + BlockN.write c j (Zq $ if even b then 1 else q32 - 1) + loop c xofLen bytes (pos + 1) (i + 1) + where + j = fromIntegral $ B.index bytes pos + + doHash :: KnownNat bitlen => proxy bitlen -> Digest (SHAKE256 bitlen) + doHash _ = Crypto.hash rho + +-- Samples a polynomial โˆˆ ๐‘‡๐‘ž +rejNttPoly :: SecureBytes Pub -> Word8 -> Word8 -> Tq Pub +rejNttPoly rho !s !r = Tq $ + BlockN.runNew (Proxy :: Proxy Pub) $ \b -> runXof b (280 * 3) 0 0 + where + runXof !b !xofLen !pos !j = case someNatVal (fromIntegral (8 * xofLen)) of + SomeNat proxy -> do + let bytes = Crypto.unBlockDigest (doHash proxy) + loop b xofLen bytes pos j + + loop !b !xofLen !bytes !pos j + | j == 256 = return () + | pos >= Offset xofLen = runXof b (xofLen + 56 * 3) pos j + | otherwise = do + let s0 = fromIntegral $ blockIndex bytes pos + s1 = fromIntegral $ blockIndex bytes (pos + 1) + s2 = fromIntegral $ blockIndex bytes (pos + 2) + z = s0 + (s1 `unsafeShiftL` 8) + + ((s2 .&. 0x7F) `unsafeShiftL` 16) + poke b j z >>= loop b xofLen bytes (pos + 3) + + poke b j z + | z < q32 = BlockN.write b j (toMontgomery $ Zq z) >> return (j + 1) + | otherwise = return j + + doHash :: KnownNat bitlen => proxy bitlen -> BlockDigest (SHAKE128 bitlen) + doHash _ = Crypto.hashToBlock input + + input :: SecureBytes Pub + !input = B.unsafeCreate (len + 2) $ \d -> do + B.copyByteArrayToPtr rho d + pokeByteOff d len s + pokeByteOff d (len + 1) r + len = B.length rho + +-- Samples an element ๐‘Ž โˆˆ ๐‘… with coefficients in [-๐œ‚, ๐œ‚] computed via rejection +-- sampling from ๐œŒ +rejBoundedPoly :: SecureBytes Sec -> Int -> Word16 -> Rq Sec +rejBoundedPoly rho !eta !r = Rq $ + BlockN.runNew (Proxy :: Proxy Sec) $ \b -> runXof b 408 0 0 + where + runXof :: MutableBlockN Sec N Zq s -> Int -> Int -> Offset Zq -> ST s () + runXof !b !xofLen !pos !j = case someNatVal (fromIntegral (8 * xofLen)) of + SomeNat proxy -> case eta of + 2 -> loop2 b xofLen bytes pos j + 4 -> loop4 b xofLen bytes pos j + _ -> error "rejBoundedPoly: unexpected eta value" + where bytes = Crypto.unDigest (doHash proxy) + + loop2 :: MutableBlockN Sec N Zq s -> Int -> SecureBytes Sec -> Int -> Offset Zq -> ST s () + loop2 !b !xofLen !bytes !pos j + | j == 256 = return () + | pos >= xofLen = runXof b (xofLen + 136) pos j + | otherwise = do + let z = fromIntegral $ B.index bytes pos + z0 = z .&. 0xF + z1 = z `unsafeShiftR` 4 + j2 <- poke2 b j z0 + when (j2 < 256) $ poke2 b j2 z1 >>= loop2 b xofLen bytes (pos + 1) + + loop4 :: MutableBlockN Sec N Zq s -> Int -> SecureBytes Sec -> Int -> Offset Zq -> ST s () + loop4 !b !xofLen !bytes !pos j + | j == 256 = return () + | pos >= xofLen = runXof b (xofLen + 136) pos j + | otherwise = do + let z = fromIntegral $ B.index bytes pos + z0 = z .&. 0xF + z1 = z `unsafeShiftR` 4 + j2 <- poke4 b j z0 + when (j2 < 256) $ poke4 b j2 z1 >>= loop4 b xofLen bytes (pos + 1) + + -- rejection sampling from {-2, โ€ฆ , 2} + poke2 b j z + | z < 15 = BlockN.write b j (Zq 2 .- Zq (mod5 z)) >> return (j + 1) + | otherwise = return j + + -- rejection sampling from {-4, โ€ฆ , 4} + poke4 b j z + | z < 9 = BlockN.write b j (Zq 4 .- Zq z) >> return (j + 1) + | otherwise = return j + + doHash :: KnownNat bitlen => proxy bitlen -> Digest (SHAKE256 bitlen) + doHash _ = Crypto.hash input + + input :: SecureBytes Sec + !input = B.unsafeCreate (len + 2) $ \d -> do + B.copyByteArrayToPtr rho d + pokeByteOff d len (B.toLE r) + len = B.length rho + +powerTwoRound :: Rq Sec -> (Rq Pub, Rq Pub) +powerTwoRound (Rq !a) = runST $ do + mb <- newF + mc <- newF + loop mb mc 0 + b <- unsafeFreezeF mb + c <- unsafeFreezeF mc + return (Rq b, Rq c) + where + loop :: MutableBlockN Pub N Zq s -> MutableBlockN Pub N Zq s -> Int -> ST s () + loop !mb !mc i = when (i < n) $ do + let r = BlockN.index a (Offset i) + (r1, r0) = powerTwoRoundZq r + BlockN.write mb (Offset i) r1 + BlockN.write mc (Offset i) r0 + loop mb mc (i + 1) + +-- Decomposes ๐‘Ÿ into (๐‘Ÿ1, ๐‘Ÿ0) such that ๐‘Ÿ โ‰ก ๐‘Ÿ1.2^๐‘‘ + ๐‘Ÿ0 mod ๐‘ž (for ๐‘‘=13) +powerTwoRoundZq :: Zq -> (Zq, Zq) +powerTwoRoundZq (Zq r) = (Zq r1, Zq r0) + where + r0' = r .&. 0x1fff + r1' = r .&. 0xffffe000 -- keep multiplied by 2^13 + + -- if r0 > 4096, increment r1 and subtract 8192 from r0 + cond = (4096 - r0') `unsafeShiftIR` 31 + r0 = select32 cond (q32 - 0x2000 + r0') r0' + r1 = reduceSimple (r1' + (cond .&. 0x2000)) + +-- Decomposes ๐‘Ÿ into (๐‘Ÿ1, ๐‘Ÿ0) such that ๐‘Ÿ โ‰ก ๐‘Ÿ1.(2๐›พ2) + ๐‘Ÿ0 mod ๐‘ž +decomposeZq :: Word32 -> Zq -> (Word32, Zq) +decomposeZq gamma2 r@(Zq rr) = (r1, Zq r0) + where + r1 = highZq gamma2 r + r0' = rr - r1 * (2 * gamma2) + !r0 = r0' + ((r0' `unsafeShiftIR` 31) .&. q32) + +highZq :: Word32 -> Zq -> Word32 +highZq gamma2 (Zq r) = case gamma2 of + 95232 -> high190464 r + _ -> high523776 r + +high190464 :: Word32 -> Word32 +high190464 r = r1 `xor` (((43 - r1) `unsafeShiftIR` 31) .&. r1) + where + r' = (r + 127) `unsafeShiftR` 7 + r1 = (r' * 11275 + (1 `unsafeShiftL` 23)) `unsafeShiftR` 24; + +high523776 :: Word32 -> Word32 +high523776 r = r1 .&. 15 + where + r' = (r + 127) `unsafeShiftR` 7 + r1 = (r' * 1025 + (1 `unsafeShiftL` 21)) `unsafeShiftR` 22; + +-- Returns ๐‘Ÿ1 from the output of decomposeZq ๐‘Ÿ +highBits :: Word32 -> Rq Sec -> BlockN Sec N Word32 +highBits gamma2 (Rq w) = BlockN.seq gamma2 $ + BlockN.mapEqPrimSize (highZq gamma2) w +{-# INLINE highBits #-} + +-- Returns ๐‘Ÿ0 from the output of decomposeZq ๐‘Ÿ +lowBits :: Word32 -> Rq Sec -> Rq Sec +lowBits gamma2 (Rq w) = Rq $ BlockN.seq gamma2 $ + BlockN.mapEqPrimSize (snd . decomposeZq gamma2) w +{-# INLINE lowBits #-} + +-- Computes hint bit indicating whether adding ๐‘ง to ๐‘Ÿ alters the high bits of ๐‘Ÿ +makeHint :: Word32 -> Rq Sec -> Rq Sec -> Hints +makeHint gamma2 (Rq z) (Rq r) = Hints $ BlockN.seq gamma2 $ + BlockN.zipWith (makeHintH gamma2) z r +{-# INLINE makeHint #-} + +-- Returns the high bits of ๐‘Ÿ adjusted according to hint โ„Ž +useHint :: Word32 -> Hints -> Rq Sec -> BlockN Sec N Word32 +useHint gamma2 (Hints h) (Rq r) = BlockN.seq gamma2 $ + BlockN.zipWithEqPrimSizeR (useHintH gamma2) h r +{-# INLINE useHint #-} + +makeHintH :: Word32 -> Zq -> Zq -> H +makeHintH gamma2 z r = notEqualH r1 v1 + where + r1 = highZq gamma2 r + v1 = highZq gamma2 (r .+ z) + +useHintH :: Word32 -> H -> Zq -> Word32 +useHintH !gamma2 h r + | h == H 0 = r1 + | positive r0 = decr gamma2 r1 + | otherwise = incr gamma2 r1 + where + (r1, Zq r0) = decomposeZq gamma2 r + + -- useHintH already has a 6-way code path with just 3 possibilities for + -- output times the match on gamma2, so for the other tests we do not use + -- boolean expressions, otherwise code size would grow even more + + -- z > 0 and z <= (๐‘ž-1)/2 + positive z = + let x = negate z + m = (q32 - 1) `div` 2 + w = x .&. xor (m + x) 0x80000000 + in w < 0x80000000 + + incr43, decr43 :: Word32 -> Word32 + decr43 x = 43 - incr43 (43 - x) + incr43 x = (x + 1) .&. complement wrapAroundMask + where wrapAroundMask = (42 - x) `unsafeShiftIR` 31 + + incr 95232 z = incr43 z -- if z == 43 then 0 else z + 1 + incr _ z = (z + 1) .&. 15 + + decr 95232 z = decr43 z -- if z == 0 then 43 else z - 1 + decr _ z = (z - 1) .&. 15 +{-# INLINE useHintH #-} diff --git a/src/Base.hs b/src/Base.hs new file mode 100644 --- /dev/null +++ b/src/Base.hs @@ -0,0 +1,62 @@ +-- | +-- Module : Base +-- License : BSD-3-Clause +-- Copyright : (c) 2026 Olivier Chรฉron +-- +-- Utilities similar to ones provided in basement but adapted for primitive +-- +{-# LANGUAGE CPP #-} +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE TypeFamilies #-} +module Base + ( CountOf(..), KnownNat, Nat, Offset(..), PrimMonad, PrimType, PrimSized(..) + , PrimState, natVal, offsetShiftL, offsetShiftR, unsafePrimFromIO, (.==#) +#ifdef ML_DSA_TESTING + , checkBounds +#endif + ) where + +import Control.Monad.Primitive + +import Data.Primitive.Types (Prim) + +import Data.Bits +import Data.Word + +import GHC.TypeNats + +type PrimType = Prim + +class PrimType a => PrimSized a where + type PrimSize a :: Nat + +instance PrimSized Word32 where + type PrimSize Word32 = 4 + +newtype CountOf ty = CountOf Int + deriving (Show, Eq, Ord, Num) + +newtype Offset ty = Offset Int + deriving (Show, Eq, Ord, Num) + +offsetShiftL :: Int -> Offset ty -> Offset ty2 +offsetShiftL n (Offset o) = Offset (o `unsafeShiftL` n) + +offsetShiftR :: Int -> Offset ty -> Offset ty2 +offsetShiftR n (Offset o) = Offset (o `unsafeShiftR` n) + +unsafePrimFromIO :: PrimMonad m => IO a -> m a +unsafePrimFromIO = unsafeIOToPrim + +(.==#) :: Offset ty -> CountOf ty -> Bool +(.==#) (Offset ofs) (CountOf sz) = ofs == sz +{-# INLINE (.==#) #-} + +#ifdef ML_DSA_TESTING +checkBounds :: CountOf ty -> Offset ty -> a -> a +checkBounds (CountOf sz) (Offset ofs) a + | ofs >= 0 && ofs < sz = a + | otherwise = error ("offset not in valid range: " ++ show ofs) +#endif diff --git a/src/Block.hs b/src/Block.hs new file mode 100644 --- /dev/null +++ b/src/Block.hs @@ -0,0 +1,127 @@ +-- | +-- Module : Block +-- License : BSD-3-Clause +-- Copyright : (c) 2025 Olivier Chรฉron +-- +-- An array of primitive (unlifted) elements. This module exposes the +-- t'PrimArray' implementation from primitive but through an API similar to +-- basement @Block@. +-- +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE UnboxedTuples #-} +module Block + ( Block, MutableBlock, blockIndex, blockRead, blockWrite, erase + , foldZipWith, Block.length, mutableContents, getMutableLength + , Block.new, Block.newPinned, Block.thaw, thawPinned + , Block.unsafeCast, unsafeCastMut, Block.unsafeFreeze +#ifdef ML_DSA_TESTING + , Block.toList +#endif + ) where + +import Control.Monad.Primitive + +import Data.Primitive.ByteArray +import Data.Primitive.PrimArray +import Data.Primitive.Types (sizeOf#) + +import Control.Exception (assert) + +import Foreign.Ptr (Ptr) + +import Base + +import GHC.Base (Int(I#), setByteArray#) + +type Block = PrimArray +type MutableBlock ty s = MutablePrimArray s ty + +blockIndex :: PrimType ty => Block ty -> Offset ty -> ty +blockRead :: (PrimMonad prim, PrimType ty) => MutableBlock ty (PrimState prim) -> Offset ty -> prim ty +blockWrite :: (PrimMonad prim, PrimType ty) => MutableBlock ty (PrimState prim) -> Offset ty -> ty -> prim () +#ifdef ML_DSA_TESTING +blockIndex b off@(Offset i) = + checkBounds (Block.length b) off $ indexPrimArray b i +blockRead mb off@(Offset i) = getSizeofMutablePrimArray mb >>= \sz -> + checkBounds (CountOf sz) off $ readPrimArray mb i +blockWrite mb off@(Offset i) a = getSizeofMutablePrimArray mb >>= \sz -> + checkBounds (CountOf sz) off $ writePrimArray mb i a +#else +blockIndex b (Offset i) = indexPrimArray b i +blockRead mb (Offset i) = readPrimArray mb i +blockWrite mb (Offset i) = writePrimArray mb i +#endif + +erase :: (PrimMonad prim, PrimType ty) => CountOf ty -> MutableBlock ty (PrimState prim) -> prim () +erase len@(CountOf n) = eraseBytes (I# (sizeOf# (toType len)) * n) + where + toType :: CountOf ty -> ty + toType = undefined + +eraseBytes :: PrimMonad prim => Int -> MutableBlock ty (PrimState prim) -> prim () +eraseBytes (I# len) (MutablePrimArray mbarr) = primitive $ \s1 -> + case setByteArray# mbarr 0# len 0# s1 of + s2 -> (# s2, () #) + +foldZipWith :: (PrimType a, PrimType b) + => (c -> a -> b -> c) -> c -> Block a -> Block b -> c +foldZipWith f c a b = assert (sa == sb) $ + loop c 0 + where + sa = sizeofPrimArray a + sb = sizeofPrimArray b + + loop !acc i + | i == sa = acc + | otherwise = do + let va = indexPrimArray a i + let vb = indexPrimArray b i + loop (f acc va vb) (i + 1) +{-# INLINE foldZipWith #-} + +length :: PrimType ty => Block ty -> CountOf ty +length = CountOf . sizeofPrimArray + +getMutableLength :: (PrimMonad prim, PrimType ty) => MutableBlock ty (PrimState prim) -> prim (CountOf ty) +getMutableLength = fmap CountOf . getSizeofMutablePrimArray + +mutableContents :: MutableBlock ty s -> Ptr ty +mutableContents = mutablePrimArrayContents -- pinned only + +new :: (PrimMonad prim, PrimType ty) => CountOf ty -> prim (MutableBlock ty (PrimState prim)) +new (CountOf n) = newPrimArray n + +newPinned :: (PrimMonad prim, PrimType ty) => CountOf ty -> prim (MutableBlock ty (PrimState prim)) +newPinned (CountOf n) = newPinnedPrimArray n + +thaw :: PrimMonad prim => Block ty -> prim (MutableBlock ty (PrimState prim)) +thaw (PrimArray !barr) = unsafeSTToPrim $ + -- as optimization, combine both steps in a known monad and avoid + -- round trip between byte length and element count + thawByteArray ba 0 (sizeofByteArray ba) >>= \(MutableByteArray mbarr) -> + return (MutablePrimArray mbarr) + where ba = ByteArray barr + +thawPinned :: PrimMonad prim => Block ty -> prim (MutableBlock ty (PrimState prim)) +thawPinned (PrimArray !barr) = unsafeSTToPrim $ do + let ba = ByteArray barr + n = sizeofByteArray ba + mb@(MutableByteArray mbarr) <- newPinnedByteArray n + copyByteArray mb 0 ba 0 n + return (MutablePrimArray mbarr) + +#ifdef ML_DSA_TESTING +toList :: PrimType ty => Block ty -> [ty] +toList = primArrayToList +#endif + +unsafeCast :: Block a -> Block b +unsafeCast (PrimArray b) = PrimArray b + +unsafeCastMut :: MutableBlock a m -> MutableBlock b m +unsafeCastMut (MutablePrimArray mb) = MutablePrimArray mb + +unsafeFreeze :: PrimMonad prim => MutableBlock ty (PrimState prim) -> prim (Block ty) +unsafeFreeze = unsafeFreezePrimArray diff --git a/src/BlockN.hs b/src/BlockN.hs new file mode 100644 --- /dev/null +++ b/src/BlockN.hs @@ -0,0 +1,262 @@ +-- | +-- Module : BlockN +-- License : BSD-3-Clause +-- Copyright : (c) 2025 Olivier Chรฉron +-- +-- A secure block with length at type level +-- +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE ScopedTypeVariables #-} +module BlockN + ( BlockN, MutableBlockN, BlockN.erase, BlockN.foldl', index, iterModify + , iterSet, mapEqPrimSize, BlockN.read, BlockN.replicate, runNew, runThaw + , runFold, BlockN.seq, BlockN.unsafeCast, BlockN.write, BlockN.zipWith + , zipWithEqPrimSizeR +#ifdef ML_DSA_TESTING + , create, BlockN.fromList, BlockN.toList +#endif + ) where + +import Control.DeepSeq (NFData(..)) +import Control.Monad.ST + +import Data.Proxy + +import Base +import Block (MutableBlock, blockRead, blockWrite, erase, unsafeCastMut) +import Equality +import Fusion +import Marking (Classified, SecurityMarking) +import SecureBlock (SecureBlock) +import qualified SecureBlock +import Math + +newtype BlockN marking (n :: Nat) a = BlockN { unBlockN :: SecureBlock marking a } + +#ifdef ML_DSA_TESTING +instance (Classified marking, Eq a, PrimType a) => Eq (BlockN marking n a) where + BlockN a == BlockN b = SecureBlock.eq a b + +instance (Classified marking, PrimType a, Show a) => Show (BlockN marking n a) where + showsPrec d = SecureBlock.showsPrec d . unBlockN +#endif + +instance NFData (BlockN marking n a) where + rnf = SecureBlock.toNormalForm . unBlockN + +instance (Classified marking, KnownNat n, PrimType a, Add a) => Add (BlockN marking n a) where + zero = BlockN.replicate zero + {-# INLINE zero #-} + (.+) = BlockN.zipWithEndoL (.+) + {-# INLINE (.+) #-} + (.-) = BlockN.zipWithEndoL (.-) + {-# INLINE (.-) #-} + neg = BlockN.mapEndo neg + {-# INLINE neg #-} + +newtype MutableBlockN (marking :: SecurityMarking) (n :: Nat) a m = MutableBlockN { unMutableBlockN :: MutableBlock a m } + +instance (Classified marking, KnownNat n, PrimType a) => Fusion (BlockN marking n a) where + type Mut (BlockN marking n a) s = MutableBlockN marking n a s + newF = new Proxy + thawF = thaw + unsafeFreezeF = unsafeFreeze + +-- Endomorphism specialization: a different implementation is substituted +-- wherever possible with rewrite rules. Identical input and output types give +-- a chance for a transformation to be fused on an existing mutable block. + +{-# RULES +"mapEndo" [~2] forall f. BlockN.map f = mapEndo f +"zipWithEndoL" [~2] forall f. BlockN.zipWith f = zipWithEndoL f + #-} + +foldl' :: forall marking n a b. (KnownNat n, PrimType a) => (b -> a -> b) -> b -> BlockN marking n a -> b +foldl' f b (BlockN !a) = loop b 0 + where + !sz = fromIntegral $ natVal (Proxy :: Proxy n) + loop !acc i + | i .==# sz = acc + | otherwise = loop (acc `f` SecureBlock.index a i) (i + 1) +{-# INLINE foldl' #-} + +index :: PrimType a => BlockN marking n a -> Offset a -> a +index = SecureBlock.index . unBlockN + +replicate :: (Classified marking, KnownNat n, PrimType a) => a -> BlockN marking n a +replicate = create . const + +#ifdef ML_DSA_TESTING +fromList :: forall marking n a. (Classified marking, KnownNat n, PrimType a) => [a] -> Maybe (BlockN marking n a) +fromList elems + | Prelude.length elems /= sz = Nothing + | otherwise = Just $ runNew (Proxy :: Proxy marking) $ \mb -> go mb 0 elems + where + !sz = fromIntegral $ natVal (Proxy :: Proxy n) + + go !mb !i list = case list of + [] -> return () + (x:xs) -> write mb i x >> go mb (i + 1) xs + +toList :: PrimType a => BlockN marking n a -> [a] +toList = SecureBlock.toList . unBlockN +#endif + +create :: (Classified marking, KnownNat n, PrimType ty) + => (Offset ty -> ty) + -> BlockN marking n ty +create initializer = runNew Proxy $ iterSet initializer +{-# INLINE create #-} + +map :: (Classified marking, KnownNat n, PrimType a, PrimType b) + => (a -> b) -> BlockN marking n a -> BlockN marking n b +map f (BlockN a) = BlockN.seq a $ + create $ \(Offset i) -> f (SecureBlock.index a (Offset i)) +{-# INLINE [2] map #-} + +mapEndo :: (Classified marking, KnownNat n, PrimType a) + => (a -> a) -> BlockN marking n a -> BlockN marking n a +mapEndo = mapEqPrimSize +{-# INLINE mapEndo #-} + +mapEqPrimSize :: (Classified marking, KnownNat n, EqPrimSize a b) => (a -> b) -> BlockN marking n a -> BlockN marking n b +mapEqPrimSize f = runContext . iterMapContext f . thawContext +{-# INLINE mapEqPrimSize #-} + +erase :: forall marking n ty prim. (PrimType ty, KnownNat n, PrimMonad prim) + => MutableBlockN marking n ty (PrimState prim) -> prim () +erase (MutableBlockN ma) = Block.erase sz ma + where !sz = fromIntegral $ natVal (Proxy :: Proxy n) + +iterModify :: forall marking n ty prim. (PrimType ty, KnownNat n, PrimMonad prim) + => (ty -> ty) + -> MutableBlockN marking n ty (PrimState prim) + -> prim () +iterModify f = iterModifyIx (\_ x -> f x) +{-# INLINE iterModify #-} + +iterModifyIx :: forall marking n ty prim. (PrimType ty, KnownNat n, PrimMonad prim) + => (Offset ty -> ty -> ty) + -> MutableBlockN marking n ty (PrimState prim) + -> prim () +iterModifyIx f (MutableBlockN !ma) = loop 0 + where + !sz = fromIntegral $ natVal (Proxy :: Proxy n) + + loop i + | i .==# sz = pure () + | otherwise = blockRead ma i >>= \x -> blockWrite ma i (f i x) >> loop (i + 1) +{-# INLINE iterModifyIx #-} + +iterSet :: forall marking n ty prim. (PrimType ty, KnownNat n, PrimMonad prim) + => (Offset ty -> ty) + -> MutableBlockN marking n ty (PrimState prim) + -> prim () +iterSet f (MutableBlockN !ma) = loop 0 + where + !sz = fromIntegral $ natVal (Proxy :: Proxy n) + + loop i + | i .==# sz = pure () + | otherwise = blockWrite ma i (f i) >> loop (i + 1) +{-# INLINE iterSet #-} + +zipWith :: (Classified mc, KnownNat n, PrimType a, PrimType b, PrimType c) + => (a -> b -> c) -> BlockN ma n a -> BlockN mb n b -> BlockN mc n c +zipWith f (BlockN a) (BlockN b) = + BlockN.seq a $ BlockN.seq b $ create $ \(Offset i) -> + f (SecureBlock.index a (Offset i)) (SecureBlock.index b (Offset i)) +{-# INLINE [2] zipWith #-} + +zipWithEndoL :: (Classified ma, KnownNat n, PrimType a, PrimType b) + => (a -> b -> a) -> BlockN ma n a -> BlockN mb n b -> BlockN ma n a +zipWithEndoL = flip . zipWithEqPrimSizeR . flip +{-# INLINE zipWithEndoL #-} + +zipWithEqPrimSizeR :: (Classified marking, KnownNat n, PrimType a, EqPrimSize b c) + => (a -> b -> c) -> BlockN ma n a -> BlockN marking n b -> BlockN marking n c +zipWithEqPrimSizeR f a b = runContext (seqContext a (iterMapIxContext g (thawContext b))) + where g = f . index a . Offset +{-# INLINE zipWithEqPrimSizeR #-} + +unsafeCast :: BlockN marking n a -> SecureBlock marking b +unsafeCast = SecureBlock.unsafeCast . unBlockN + +read :: (PrimMonad prim, PrimType a) => MutableBlockN marking n a (PrimState prim) -> Offset a -> prim a +read = blockRead . unMutableBlockN + +write :: (PrimMonad prim, PrimType a) => MutableBlockN marking n a (PrimState prim) -> Offset a -> a -> prim () +write = blockWrite . unMutableBlockN + +new :: forall proxy marking n a prim. (Classified marking, KnownNat n, PrimMonad prim, PrimType a) => proxy marking -> prim (MutableBlockN marking n a (PrimState prim)) +new prx = MutableBlockN <$> SecureBlock.new prx (CountOf sz) + where !sz = fromIntegral $ natVal (Proxy :: Proxy n) +{-# INLINE new #-} + +runThaw :: (Classified marking, KnownNat n, PrimType a) => BlockN marking n a -> (forall s. MutableBlockN marking n a s -> ST s ()) -> BlockN marking n a +runThaw a f = runContext (modifyContext f (thawContext a)) +{-# INLINE runThaw #-} + +runNew :: (Classified marking, KnownNat n, PrimType a) => proxy marking -> (forall s. MutableBlockN marking n a s -> ST s ()) -> BlockN marking n a +runNew _ f = runContext (modifyContext f newContext) +{-# INLINE runNew #-} + +runFold :: (Classified marking, KnownNat n, PrimType a, Foldable t) => BlockN marking n a -> (forall s. b -> MutableBlockN marking n a s -> ST s ()) -> t b -> BlockN marking n a +runFold a f = runContext . foldContext f (thawContext a) +{-# INLINE runFold #-} + +seq :: (Classified marking, KnownNat n, PrimType a) => b -> BlockN marking n a -> BlockN marking n a +seq b a = runContext (seqContext b (thawContext a)) +{-# INLINE seq #-} + +thaw :: (Classified marking, PrimMonad prim) => BlockN marking n a -> prim (MutableBlockN marking n a (PrimState prim)) +thaw = fmap MutableBlockN . SecureBlock.thaw . unBlockN + +unsafeFreeze :: (Classified marking, PrimMonad prim) => MutableBlockN marking n a (PrimState prim) -> prim (BlockN marking n a) +unsafeFreeze = fmap BlockN . SecureBlock.unsafeFreeze . unMutableBlockN + +unsafeMapIx :: forall marking n a b prim. (KnownNat n, EqPrimSize a b, PrimMonad prim) => (Int -> a -> b) -> MutableBlockN marking n a (PrimState prim) -> prim (MutableBlockN marking n b (PrimState prim)) +unsafeMapIx f (MutableBlockN !ma) = MutableBlockN . ensureEqPrimSize witness <$> loop 0 + where + witness = undefined :: a -> b + !sz = fromIntegral $ natVal (Proxy :: Proxy n) + loop i + | i == sz = return (unsafeCastMut ma) + | otherwise = do + a <- blockRead ma (Offset i) + blockWrite (unsafeCastMut ma) (Offset i) (f i a) + loop (i + 1) +{-# INLINE unsafeMapIx #-} + +-- + +iterMapContext :: (EqPrimSize a b, Classified marking, KnownNat n) => (a -> b) -> Context (BlockN marking n a) -> Context (BlockN marking n b) +iterMapContext f = iterMapIxContext (\_ x -> f x) +{-# INLINE iterMapContext #-} + +iterMapIxContext :: (EqPrimSize a b, Classified marking, KnownNat n) => (Int -> a -> b) -> Context (BlockN marking n a) -> Context (BlockN marking n b) +iterMapIxContext f = mapContext m + where m = MapF { mapUpdate = unsafeMapIx f + , mapInit = \x -> newF >>= \mb -> Prelude.seq x (iterSet (g x) mb) >> return mb + } + g x (Offset i) = f i (index x (Offset i)) +{-# INLINE [1] iterMapIxContext #-} + + +-- Fusion rules +-- +-- "iterMapIxContext/iterMapIxContext" merges element-wise transformations as +-- single operations. For example @a .+ b .+ c@ becomes a single loop that +-- processes all input blocks in parallel and writes to the destination block. +-- +-- "iterMapIxContext/seqContext" moves strictness annotations upstream so that +-- they do not prevent other rules from firing. + +{-# RULES +"iterMapIxContext/seqContext" [~1] forall a f c. iterMapIxContext f (seqContext a c) = seqContext a (iterMapIxContext f c) +"iterMapIxContext/iterMapIxContext" [~1] forall f g c. iterMapIxContext f (iterMapIxContext g c) = iterMapIxContext (\i a -> f i (g i a)) c + #-} diff --git a/src/Builder.hs b/src/Builder.hs new file mode 100644 --- /dev/null +++ b/src/Builder.hs @@ -0,0 +1,99 @@ +-- | +-- Module : Builder +-- License : BSD-3-Clause +-- Copyright : (c) 2025 Olivier Chรฉron +-- +-- Eliminate intermediate allocations when concatenating several buffers +-- +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE KindSignatures #-} +{-# LANGUAGE RankNTypes #-} +module Builder + ( Builder, builderLength, bytes, copyBuilderToPtr, create, promote, public + , run, runRelaxed, runToBlock, secret, storable, unsafeCreate + ) where + +import Data.ByteArray (ByteArray, ByteArrayAccess) +import qualified Data.ByteArray as B + +import Control.Monad.ST +import Control.Monad.ST.Unsafe + +import Data.Semigroup +import Data.Word + +import Foreign.Ptr (Ptr, castPtr, plusPtr) +import Foreign.Storable (Storable(..)) + +import qualified GHC.Exts as Exts + +import Base +import Block (Block) +import Marking (Classified, Leak(..), SecurityMarking(..)) +import SecureBytes (SecureBytes) +import qualified Block +import qualified ByteArrayST as ST +import qualified SecureBytes + +data Builder (marking :: SecurityMarking) = Builder + { builderLength :: {-# UNPACK #-} !Int + , copyBuilderToPtr :: forall s. Ptr Word8 -> ST s () + } + +instance Semigroup (Builder marking) where + b1 <> b2 = create (n1 + n2) $ \p -> + copyBuilderToPtr b1 p >> copyBuilderToPtr b2 (p `plusPtr` n1) + where + n1 = builderLength b1 + n2 = builderLength b2 + +instance Monoid (Builder marking) where + mempty = empty + mconcat builders = create (getSize (Exts.inline builders)) $ + go (Exts.inline builders) + where + getSize = getSum . Prelude.mconcat . map (Sum . builderLength) + go = foldr c (const $ return ()) + c b k = Exts.oneShot $ \p -> copyBuilderToPtr b p >> k (p `plusPtr` builderLength b) + {-# INLINE mconcat #-} + +instance Leak Builder + +bytes :: Classified marking => SecureBytes marking -> Builder marking +bytes b = unsafeCreate (SecureBytes.length b) (SecureBytes.copyByteArrayToPtr b) + +create :: Int -> (forall s. Ptr a -> ST s ()) -> Builder marking +create n f = Builder n (f . castPtr) +{-# INLINE create #-} + +empty :: Builder marking +empty = Builder 0 $ \_ -> return () + +public :: ByteArrayAccess ba => ba -> Builder marking +public b = unsafeCreate (B.length b) (B.copyByteArrayToPtr b) + +promote :: Builder Pub -> Builder Sec +promote (Builder n f) = Builder n f + +secret :: ByteArrayAccess ba => ba -> Builder Sec +secret = public + +run :: Classified marking => Builder marking -> SecureBytes marking +run b = SecureBytes.unsafeCreate (builderLength b) (copyBuilderToPtr b) + +runRelaxed :: ByteArray ba => Builder Pub -> ba +runRelaxed b = ST.unsafeCreate (builderLength b) (copyBuilderToPtr b) + +runToBlock :: Builder Pub -> Block Word8 +runToBlock b = runST $ do + mb <- Block.newPinned (CountOf $ builderLength b) + copyBuilderToPtr b (Block.mutableContents mb) + Block.unsafeFreeze mb + +storable :: Storable a => a -> Builder marking +storable !a = unsafeCreate (sizeOf a) (`poke` a) + +unsafeCreate :: Int -> (Ptr a -> IO ()) -> Builder marking +unsafeCreate n f = create n (unsafeIOToST . f) +{-# INLINE unsafeCreate #-} diff --git a/src/ByteArrayST.hs b/src/ByteArrayST.hs new file mode 100644 --- /dev/null +++ b/src/ByteArrayST.hs @@ -0,0 +1,46 @@ +-- | +-- Module : ByteArrayST +-- License : BSD-3-Clause +-- Copyright : (c) 2025 Olivier Chรฉron +-- +-- Byte array primitives in the @ST@ monad instead of @IO@ +-- +{-# LANGUAGE RankNTypes #-} +module ByteArrayST + ( unsafeCreate, withByteArray + , peek, peekElemOff, pokeElemOff + , fillBytes + ) where + +import Data.ByteArray (ByteArray, ByteArrayAccess) +import qualified Data.ByteArray as B + +import Control.Monad.ST +import Control.Monad.ST.Unsafe + +import Data.Word + +import Foreign.Ptr (Ptr) +import qualified Foreign.Marshal.Utils as S +import Foreign.Storable (Storable) +import qualified Foreign.Storable as S + +unsafeCreate :: ByteArray ba => Int -> (forall s. Ptr p -> ST s ()) -> ba +unsafeCreate sz f = B.unsafeCreate sz (stToIO . f) +{-# INLINE unsafeCreate #-} + +withByteArray :: ByteArrayAccess ba => ba -> (Ptr p -> ST s a) -> ST s a +withByteArray b f = unsafeIOToST $ B.withByteArray b (unsafeSTToIO . f) +{-# INLINE withByteArray #-} + +peek :: Storable a => Ptr a -> ST s a +peek = unsafeIOToST . S.peek + +peekElemOff :: Storable a => Ptr a -> Int -> ST s a +peekElemOff a = unsafeIOToST . S.peekElemOff a + +pokeElemOff :: Storable a => Ptr a -> Int -> a -> ST s () +pokeElemOff a off = unsafeIOToST . S.pokeElemOff a off + +fillBytes :: Ptr a -> Word8 -> Int -> ST s () +fillBytes a v = unsafeIOToST . S.fillBytes a v diff --git a/src/Crypto.hs b/src/Crypto.hs new file mode 100644 --- /dev/null +++ b/src/Crypto.hs @@ -0,0 +1,214 @@ +-- | +-- Module : Crypto +-- License : BSD-3-Clause +-- Copyright : (c) 2026 Olivier Chรฉron +-- +-- Crypto-related utilities like the ML-DSA hash functions, or more general +-- concerns like constant-time equality and selection. +-- +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE ScopedTypeVariables #-} +module Crypto + ( ConstEqW(..), BoolW, trueW, andW, ltW, lteW, toBool + , h, h64, h128 + , BlockDigest, unBlockDigest, hashToBlock + , Digest, unDigest, hash + ) where + +import Crypto.Hash (Context) +import Crypto.Hash.Algorithms +import Crypto.Hash.IO + +import Control.Exception (assert, mask_) +import Control.Monad.ST + +import Data.ByteArray (ByteArrayAccess, Bytes, MemView(..), ScrubbedBytes) +import qualified Data.ByteArray as B + +import Data.Bits +import Data.Word + +import GHC.TypeNats + +import Foreign.Marshal.Utils (fillBytes) +import Foreign.Ptr (Ptr, castPtr, plusPtr) +import Foreign.Storable (pokeByteOff) + +import Block (Block) +import Builder (Builder) +import Marking (Classified) +import ScrubbedBlock (ScrubbedBlock) +import SecureBytes (SecureBytes) +import Vector (Vector) +import qualified Block +import qualified Builder +import qualified ByteArrayST as ST +import qualified ScrubbedBlock +import qualified Vector + +unsafeShiftIR :: Word -> Int -> Word +unsafeShiftIR x s = fromIntegral ((fromIntegral x :: Int) `unsafeShiftR` s) +{-# INLINE unsafeShiftIR #-} + +newtype BoolW = BoolW Word + +#ifdef ML_DSA_TESTING +instance Show BoolW where + showsPrec d = showsPrec d . toBool +#endif + +toBool :: BoolW -> Bool +toBool (BoolW mask) = mask /= 0 + +falseW, trueW :: BoolW +falseW = BoolW 0 +trueW = BoolW maxBound + +notW :: BoolW -> BoolW +notW (BoolW a) = BoolW (complement a) + +andW :: BoolW -> BoolW -> BoolW +andW (BoolW a) (BoolW b) = BoolW (a .&. b) + +bitsW :: Int +bitsW = let BoolW x = falseW in finiteBitSize x + +bytesW :: Int +bytesW = div bitsW 8 + +msbW :: Word -> Word +msbW x = x `unsafeShiftIR` (bitsW - 1) + +eqW :: Word -> Word -> BoolW +eqW a b = isZeroW (a `xor` b) + where isZeroW x = BoolW $ msbW (complement x .&. (x - 1)) + +ltW :: Word -> Word -> BoolW +ltW a b = BoolW $ msbW (a - b) + +lteW :: Word -> Word -> BoolW +lteW a b = notW (ltW b a) + +assertMultW :: Int -> a -> a +assertMultW n = assert (n .&. mask == 0) + where mask = bytesW - 1 + +class ConstEqW a where + constEqW :: a -> a -> BoolW + +instance ConstEqW a => ConstEqW (Vector n a) where + constEqW = + Vector.fold1ZipWith (\mask x y -> mask `andW` constEqW x y) constEqW + +instance ConstEqW (Block Word) where + constEqW a b + | Block.length a /= Block.length b = falseW + | otherwise = Block.foldZipWith (\mask x y -> mask `andW` eqW x y) trueW a b + +instance ConstEqW (ScrubbedBlock Word) where + constEqW a b + | ScrubbedBlock.length a /= ScrubbedBlock.length b = falseW + | otherwise = ScrubbedBlock.foldZipWith (\mask x y -> mask `andW` eqW x y) trueW a b + +instance ConstEqW Bytes where + constEqW = bytesConstEqW + +instance ConstEqW ScrubbedBytes where + constEqW = bytesConstEqW + +bytesConstEqW :: (ByteArrayAccess bs1, ByteArrayAccess bs2) => bs1 -> bs2 -> BoolW +bytesConstEqW a b + | B.length a /= B.length b = falseW + | otherwise = foldZipWith (\mask x y -> mask `andW` eqW x y) trueW a b + +foldZipWith :: (ByteArrayAccess bs1, ByteArrayAccess bs2) + => (c -> Word -> Word -> c) -> c -> bs1 -> bs2 -> c +foldZipWith f c a b = assert (sa == sb) $ assertMultW sa $ assertMultW sb $ + runST $ ST.withByteArray a $ \pa -> ST.withByteArray b $ \pb -> + loop (pa :: Ptr Word) (pb :: Ptr Word) c 0 + where + !sa = B.length a + !sb = B.length b + + loop !pa !pb !acc i + | i == sa = return acc + | otherwise = do + va <- ST.peek pa + vb <- ST.peek pb + loop (pa `plusPtr` bytesW) (pb `plusPtr` bytesW) (f acc va vb) (i + bytesW) +{-# INLINE foldZipWith #-} + +h :: (Classified marking, ByteArrayAccess (SecureBytes marking)) + => Int -> SecureBytes marking -> SecureBytes marking +h !len input = case someNatVal (fromIntegral (8 * len)) of + SomeNat proxy -> Builder.run $ hashWith (alg proxy) input + where + alg :: proxy bitlen -> SHAKE256 bitlen + alg _ = SHAKE256 + +h64 :: ByteArrayAccess (SecureBytes marking) + => SecureBytes marking -> Builder marking +h64 = hashWith (SHAKE256 :: SHAKE256 512) + +h128 :: ByteArrayAccess a + => a -> Word8 -> Word8 -> (Bytes, ScrubbedBytes, ScrubbedBytes) +h128 xi !k !l = (a, b, c) + where + rho = B.takeView bs 32 + rho' = B.view bs 32 64 + kk = B.dropView bs 96 + + !a = B.convert rho + !b = B.convert rho' + !c = B.convert kk + + bs :: ScrubbedBytes + bs = Builder.run $ hashWith (SHAKE256 :: SHAKE256 1024) input + + input :: ScrubbedBytes + input = B.unsafeCreate (len + 2) $ \p -> do + B.copyByteArrayToPtr xi p + pokeByteOff p len k + pokeByteOff p (len + 1) l + len = B.length xi + +-- Override cryptonite/crypton types and hashing functions. +-- +-- Standard type Digest is a newtype over an unpinned Block Word8, which +-- requires a trampoline to implement most Ptr access to the underlying byte +-- array. Instead we re-implement here the Digest type over ScrubbedBytes as +-- well as pinned Block backends, to avoid trampoline costs. Additionnally +-- we use the mutable API to avoid copying the hashing Context in between +-- steps init/update/finalize and then clear the content. + +newtype Digest a = Digest { unDigest :: ScrubbedBytes } +newtype BlockDigest a = BlockDigest { unBlockDigest :: Block Word8 } + +hash :: forall a ba. (HashAlgorithm a, ByteArrayAccess ba) => ba -> Digest a +hash = Digest . Builder.run . hashWith (undefined :: a) + +hashToBlock :: forall a. HashAlgorithm a => Bytes -> BlockDigest a +hashToBlock = BlockDigest . Builder.runToBlock . hashWith (undefined :: a) + +hashWith :: forall marking a ba. (HashAlgorithm a, ByteArrayAccess ba) => a -> ba -> Builder marking +hashWith a ba = Builder.unsafeCreate (hashDigestSize a) $ \dig -> + hashMutableInit >>= \ctx -> mask_ $ do + hashUpdateChunked (ctx :: MutableContext a) ba + B.withByteArray ctx $ \pctx -> do + hashInternalFinalize (castPtr pctx :: Ptr (Context a)) dig + fillBytes pctx 0 (B.length ctx) + +hashUpdateChunked :: (HashAlgorithm a, ByteArrayAccess ba) => MutableContext a -> ba -> IO () +hashUpdateChunked ctx ba = B.withByteArray ba $ goChunked (B.length ba) + where + chunkSize = 1073741824 -- 1 GB + + goChunked remaining p + | remaining > chunkSize = do + hashMutableUpdate ctx (MemView p chunkSize) + goChunked (remaining - chunkSize) (p `plusPtr` chunkSize) + | otherwise = hashMutableUpdate ctx (MemView p remaining) diff --git a/src/Crypto/PubKey/ML_DSA.hs b/src/Crypto/PubKey/ML_DSA.hs new file mode 100644 --- /dev/null +++ b/src/Crypto/PubKey/ML_DSA.hs @@ -0,0 +1,334 @@ +-- | +-- Module : Crypto.PubKey.ML_DSA +-- License : BSD-3-Clause +-- Maintainer : Olivier Chรฉron +-- Stability : provisional +-- Portability : unknown +-- +-- Module-Lattice-based Digital Signature Algorithm (ML-DSA), defined +-- in . +{-# 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 ] diff --git a/src/Equality.hs b/src/Equality.hs new file mode 100644 --- /dev/null +++ b/src/Equality.hs @@ -0,0 +1,25 @@ +-- | +-- Module : Equality +-- License : BSD-3-Clause +-- Copyright : (c) 2026 Olivier Chรฉron +-- +-- Generate a constraint @'PrimSize' a ~ 'PrimSize' b@ +-- +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE TypeOperators #-} +module Equality + ( EqPrimSize, ensureEqPrimSize + ) where + +import Data.Type.Equality + +import Base + +type EqPrimSize a b = (PrimType a, PrimType b, PrimSize a ~ PrimSize b) + +eqPrimSize :: PrimSize a ~ PrimSize b => k a b -> PrimSize a :~: PrimSize b +eqPrimSize _ = Refl + +ensureEqPrimSize :: EqPrimSize a b => k a b -> c -> c +ensureEqPrimSize op = case eqPrimSize op of Refl -> id diff --git a/src/Fusion.hs b/src/Fusion.hs new file mode 100644 --- /dev/null +++ b/src/Fusion.hs @@ -0,0 +1,102 @@ +-- | +-- Module : Fusion +-- License : BSD-3-Clause +-- Copyright : (c) 2026 Olivier Chรฉron +-- +-- Infrastructure to decrease intermediate allocations and prefer in-place +-- mutation when possible +-- +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE TypeFamilyDependencies #-} +module Fusion + ( Fusion(..), MapF(..) + , Context, runContext, newContext, thawContext, mapContext, modifyContext + , foldContext, seqContext + ) where + +import Control.Monad ( forM_, (>=>) ) +import Control.Monad.ST + +-- class of values that can be mutated in the ST monad +class Fusion a where + type Mut a s = mut | mut -> a + newF :: ST s (Mut a s) + thawF :: a -> ST s (Mut a s) + unsafeFreezeF :: Mut a s -> ST s a + +-- a transformation step in the fusion pipeline, with two implementations +-- provided: one that operates on an existing mutation context, and one that +-- initiates a new context from the input +data MapF a b = MapF + { mapUpdate :: forall s. Mut a s -> ST s (Mut b s) + , mapInit :: forall s. a -> ST s (Mut b s) + } + +-- MapF is almost a category except for the 'Fusion' constraint on objects +-- +-- idMapF :: Fusion a => MapF a a +-- idMapF = MapF { mapUpdate = pure, mapInit = thawF } + +composeMapF :: MapF b c -> MapF a b -> MapF a c +composeMapF m2 m1 = MapF + { mapUpdate = mapUpdate m1 >=> mapUpdate m2 + , mapInit = mapInit m1 >=> mapUpdate m2 + } + +-- fusion context +newtype Context a = Context (forall s. ST s (Mut a s)) + +newContext :: Fusion a => Context a +newContext = Context newF + +thawContext :: Fusion a => a -> Context a +thawContext a = Context $ thawF a +{-# INLINE [0] thawContext #-} + +modifyContext :: (forall s. Mut a s -> ST s ()) -> Context a -> Context a +modifyContext f = bindContext $ \ma -> f ma >> return ma + +mapContext :: MapF a b -> Context a -> Context b +mapContext m = bindContext (mapUpdate m) +{-# INLINE [0] mapContext #-} + +initContext :: MapF a b -> a -> Context b +initContext m a = Context $ mapInit m a + +bindContext :: (forall s. Mut a s -> ST s (Mut b s)) -> Context a -> Context b +bindContext f (Context ctx) = Context $ ctx >>= f + +foldContext :: Foldable t => (forall s. b -> Mut a s -> ST s ()) -> Context a -> t b -> Context a +foldContext f c bs = modifyContext (\ma -> forM_ bs $ \b -> f b ma) c + +runContext :: Fusion a => Context a -> a +runContext (Context ctx) = runST (ctx >>= unsafeFreezeF) +{-# INLINE [0] runContext #-} + +seqContext :: a -> Context b -> Context b +seqContext = seq +{-# INLINE [0] seqContext #-} + + +-- Fusion rules +-- +-- "thawContext/runContext" is the canonical optimization that eliminates an +-- allocation + value copy. Instead, it sequences two transformations on the +-- same mutation context. +-- +-- "mapContext/seqContext" moves strictness annotations upstream so that they +-- do not prevent other rules from firing. +-- +-- "mapContext/mapContext" is not strictly needed: the function is ultimately +-- inlined to the same code. But we keep it so that simplifications fire early +-- and do not wait for the final phase. +-- +-- "mapContext/thawContext" is the rule that invokes mapInit instead of copying +-- the input and calling mapUpdate. + +{-# RULES +"thawContext/runContext" [~0] forall c. thawContext (runContext c) = c +"mapContext/seqContext" [~0] forall a m c. mapContext m (seqContext a c) = seqContext a (mapContext m c) +"mapContext/mapContext" [~0] forall m1 m2 c. mapContext m2 (mapContext m1 c) = mapContext (composeMapF m2 m1) c +"mapContext/thawContext" [1] forall m a. mapContext m (thawContext a) = initContext m a + #-} diff --git a/src/Internal.hs b/src/Internal.hs new file mode 100644 --- /dev/null +++ b/src/Internal.hs @@ -0,0 +1,389 @@ +-- | +-- 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 diff --git a/src/Iterate.hs b/src/Iterate.hs new file mode 100644 --- /dev/null +++ b/src/Iterate.hs @@ -0,0 +1,47 @@ +-- | +-- Module : Iterate +-- License : BSD-3-Clause +-- Copyright : (c) 2025 Olivier Chรฉron +-- +-- @'offsetsFrom' x y@ is similar to @[x .. y - 1]@ but does not verify that +-- @x <= y@. This removes an unnecessary branch. +-- +{-# LANGUAGE MagicHash #-} +module Iterate + ( offsets, offsetsFrom + ) +where + +import Control.Exception (assert) + +import GHC.Exts + +{-# INLINE offsets #-} +offsets :: Int -> [Int] +offsets = offsetsFrom 0 + +{-# INLINE offsetsFrom #-} +offsetsFrom :: Int -> Int -> [Int] +offsetsFrom xx@(I# x) yy@(I# y) = assert (xx <= yy) (indices x y) + +-- Implementation is derived from enumFromTo, this acts as good producer for +-- list fusion. + +{-# RULES +"indices" [~1] forall x y. indices x y = build (\c n -> indicesFB c n x y) +"indicesList" [1] indicesFB (:) [] = indices + #-} + +indices :: Int# -> Int# -> [Int] +indices x0 y = go x0 + where + go x = if isTrue# (x <# y) then I# x : go (x +# 1#) else [] +{-# NOINLINE [1] indices #-} + +indicesFB :: (Int -> r -> r) -> r -> Int# -> Int# -> r +indicesFB c n x0 y = go x0 + where + go x = if isTrue# (x <# y) then I# x `c` go (x +# 1#) else n + -- Be very careful not to have more than one "c" so that when indicesFB is + -- inlined we can inline whatever is bound to "c" +{-# INLINE [0] indicesFB #-} diff --git a/src/Machine.hs b/src/Machine.hs new file mode 100644 --- /dev/null +++ b/src/Machine.hs @@ -0,0 +1,92 @@ +-- | +-- Module : Machine +-- License : BSD-3-Clause +-- Copyright : (c) 2025 Olivier Chรฉron +-- +-- Architecture-dependent utilities, to read/write unaligned machine words +-- in little-endian order +-- +{-# LANGUAGE CPP #-} +module Machine + ( WordM, WordLE, assertMultM, fromLE, toLE, wordBits, wordBytes + ) where + +#include "MachDeps.h" + +-- Taken from `bytestring`, a list of architectures known to accept +-- unaligned loads and stores +#if defined(i386_HOST_ARCH) || defined(x86_64_HOST_ARCH) \ + || ((defined(arm_HOST_ARCH) || defined(aarch64_HOST_ARCH)) \ + && defined(__ARM_FEATURE_UNALIGNED)) \ + || defined(powerpc_HOST_ARCH) || defined(powerpc64_HOST_ARCH) \ + || defined(powerpc64le_HOST_ARCH) +#define MLDSA_ALLOW_UNALIGNED_OP 1 + +-- Little-endian conversion in `memory` / `ram` is avoided at compile +-- time only for AMD/Intel, here we will short circuit on ARM too +#if (defined(arm_HOST_ARCH) || defined(aarch64_HOST_ARCH)) \ + && !defined(WORDS_BIGENDIAN) +#define MLDSA_FORCE_LITTLE_ENDIAN_ARCH 1 +#endif + +#endif + +import Control.Exception (assert) + +#ifdef MLDSA_ALLOW_UNALIGNED_OP +import qualified Data.Memory.Endian as B +#endif + +import Data.Bits +import Data.Word + +#ifdef MLDSA_ALLOW_UNALIGNED_OP + +-- our preferred word size +#if WORD_SIZE_IN_BITS == 64 +type WordM = Word64 +#else +type WordM = Word32 +#endif + +type WordLE = B.LE WordM + +fromLE :: WordLE -> WordM +#ifdef MLDSA_FORCE_LITTLE_ENDIAN_ARCH +fromLE = B.unLE -- unwrap constructor with no byte swapping +#else +fromLE = B.fromLE -- byte swap if necessary +#endif + +toLE :: WordM -> WordLE +#ifdef MLDSA_FORCE_LITTLE_ENDIAN_ARCH +toLE = B.LE -- wrap constructor with no byte swapping +#else +toLE = B.toLE -- byte swap if necessary +#endif + +#else + +-- unaligned memory access is not allowed so we fallback to one byte at a time +-- and endianness does not matter + +type WordM = Word8 +type WordLE = WordM + +fromLE :: WordLE -> WordM +fromLE = id + +toLE :: WordM -> WordLE +toLE = id + +#endif + +wordBits :: Int +wordBits = finiteBitSize (0 :: WordM) + +wordBytes :: Int +wordBytes = div wordBits 8 + +assertMultM :: Int -> a -> a +assertMultM n = assert (n .&. mask == 0) + where mask = wordBytes - 1 diff --git a/src/Marking.hs b/src/Marking.hs new file mode 100644 --- /dev/null +++ b/src/Marking.hs @@ -0,0 +1,133 @@ +-- | +-- Module : Marking +-- License : BSD-3-Clause +-- Copyright : (c) 2025 Olivier Chรฉron +-- +-- Infrastructure that associates a security marking at type level to all +-- buffers created by the library. This determines which buffers need the +-- scrubbed (Sec) or regular (Pub) variants. +-- +{-# LANGUAGE CPP #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE PolyKinds #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilyDependencies #-} +module Marking + ( SecurityMarking(..), Classified(..), Leak(..), index + , Marking.toNormalForm, unsafeCast +#ifdef ML_DSA_TESTING + , Marking.toList +#endif + ) where + +import Control.DeepSeq (NFData(..)) +import Control.Monad.ST + +import Data.ByteArray (Bytes, ScrubbedBytes) +import qualified Data.ByteArray as B + +import Data.Kind + +import Foreign.Ptr (Ptr) + +import Unsafe.Coerce + +import Base +import Block (Block, MutableBlock, blockIndex) +import ScrubbedBlock (ScrubbedBlock) +import qualified Block +import qualified ByteArrayST as ST +import qualified ScrubbedBlock + +data SecurityMarking = Sec | Pub -- secret or public information + +-- Transformation called only at expected location in the LWE problem, after +-- adding noise to secret information. +-- +-- Block and ScrubbedBlock have the same representation, we can force coercion +-- from Sec to Pub even though the block will be actually scrubbed. This is +-- simpler than copying to a real non-scrubbed block. +class Leak t where + leak :: t Sec -> t Pub + leak = unsafeCoerce + +class Classified (marking :: SecurityMarking) where + type SecureBlock marking = (block :: Type -> Type) | block -> marking + + new :: (PrimType ty, PrimMonad prim) => proxy marking -> CountOf ty -> prim (MutableBlock ty (PrimState prim)) + thaw :: PrimMonad m => SecureBlock marking ty -> m (MutableBlock ty (PrimState m)) + unsafeFreeze :: PrimMonad prim => MutableBlock ty (PrimState prim) -> prim (SecureBlock marking ty) + +#ifdef ML_DSA_TESTING + eq :: (Eq ty, PrimType ty) => SecureBlock marking ty -> SecureBlock marking ty -> Bool + showsPrec :: (PrimType ty, Show ty) => Int -> SecureBlock marking ty -> ShowS + lengthBlock :: PrimType ty => SecureBlock marking ty -> CountOf ty +#endif + + type SecureBytes marking = bytes | bytes -> marking + unsafeCreate :: Int -> (forall s. Ptr a -> ST s ()) -> SecureBytes marking + lengthBytes :: SecureBytes marking -> Int + copyByteArrayToPtr :: SecureBytes marking -> Ptr a -> IO () + +instance Classified Pub where + type SecureBlock Pub = Block + + new _ = Block.new + thaw = Block.thaw + unsafeFreeze = Block.unsafeFreeze + +#ifdef ML_DSA_TESTING + eq = (==) + showsPrec = Prelude.showsPrec + lengthBlock = Block.length +#endif + + type SecureBytes Pub = Bytes + unsafeCreate = ST.unsafeCreate + {-# INLINE unsafeCreate #-} + lengthBytes = B.length + copyByteArrayToPtr = B.copyByteArrayToPtr + +instance Classified Sec where + type SecureBlock Sec = ScrubbedBlock + + new _ = ScrubbedBlock.new + thaw = ScrubbedBlock.thaw + unsafeFreeze = ScrubbedBlock.unsafeFreeze + +#ifdef ML_DSA_TESTING + eq = (==) + showsPrec = Prelude.showsPrec + lengthBlock = ScrubbedBlock.length +#endif + + type SecureBytes Sec = ScrubbedBytes + unsafeCreate = ST.unsafeCreate + {-# INLINE unsafeCreate #-} + lengthBytes = B.length + copyByteArrayToPtr = B.copyByteArrayToPtr + + +-- for some functions we use the fact that Block and SecureBlock have the same +-- representation and implementation + +unwrap :: SecureBlock marking a -> Block a +unwrap = unsafeCoerce + +wrap :: Block b -> SecureBlock marking b +wrap = unsafeCoerce + +index :: PrimType ty => SecureBlock marking ty -> Offset ty -> ty +index = blockIndex . unwrap + +#ifdef ML_DSA_TESTING +toList :: PrimType ty => SecureBlock marking ty -> [ty] +toList = Block.toList . unwrap +#endif + +toNormalForm :: SecureBlock marking ty -> () +toNormalForm = rnf . unwrap + +unsafeCast :: SecureBlock marking a -> SecureBlock marking b +unsafeCast = wrap . Block.unsafeCast . unwrap diff --git a/src/Math.hs b/src/Math.hs new file mode 100644 --- /dev/null +++ b/src/Math.hs @@ -0,0 +1,55 @@ +-- | +-- Module : Math +-- License : BSD-3-Clause +-- Copyright : (c) 2025 Olivier Chรฉron +-- +-- Type classes that define additive and multiplicative operations, as well as +-- a multiply-then-add operation that will often optimize chaining. +-- +-- The module also defines a non-homogenous multiplication that combines +-- typically a public operand (left) with a secret operand (right), producing a +-- secret output. +-- +{-# LANGUAGE CPP #-} +{-# LANGUAGE MultiParamTypeClasses #-} +module Math + ( Add(..), Mul(..), MulAdd(..), BiMul(..), BiMulAdd(..) + ) where + +#if !(MIN_VERSION_base(4,20,0)) +import Data.List (foldl') +#endif + +infixl 7 .* +infixr 7 ..* +infixl 6 .+, .- + +class Add a where + zero :: a + (.+) :: a -> a -> a + (.-) :: a -> a -> a + neg :: a -> a + +class Add a => Mul a where + one :: a + (.*) :: a -> a -> a + +class Mul a => MulAdd a where + -- invariant: mulAdd a b c == a .* b .+ c + mulAdd :: a -> a -> a -> a + +class Add a => BiMul b a where + (..*) :: b -> a -> a + +class BiMul b a => BiMulAdd b a where + {-# MINIMAL biMulAdd | biMulFold #-} + + -- invariant: biMulAdd a b c == a ..* b .+ c + biMulAdd :: b -> a -> a -> a + biMulAdd b a x = biMulFold x [(b, a)] + {-# INLINE biMulAdd #-} + + -- repeated biMulAdd + biMulFold :: Foldable t => a -> t (b, a) -> a + biMulFold = foldl' $ \c (b, a) -> biMulAdd b a c + {-# INLINE biMulFold #-} diff --git a/src/Matrix.hs b/src/Matrix.hs new file mode 100644 --- /dev/null +++ b/src/Matrix.hs @@ -0,0 +1,35 @@ +-- | +-- Module : Matrix +-- License : BSD-3-Clause +-- Copyright : (c) 2026 Olivier Chรฉron +-- +-- A matrix here is simply a vector of vectors. The module also implements two +-- utility functions 'mmul' and 'mmulAdd' that multiply a matrix and a vector. +-- +module Matrix + ( create, mmul, mmulAdd + ) where + +import Base +import Math +import Vector (Vector) +import qualified Vector + +create :: (KnownNat m, KnownNat n) => (Offset ty -> Offset (Vector n ty) -> ty) -> Vector m (Vector n ty) +create f = Vector.create $ \j -> Vector.create (`f` j) +{-# INLINE create #-} + +index :: Vector m (Vector n ty) -> Offset (Vector n ty) -> Offset ty -> ty +index a i = Vector.index (Vector.index a i) + +mmul :: BiMulAdd b a => Vector m (Vector n b) -> Vector n a -> Vector m a +mmul a u = Vector.seq u (fmap (`Vector.dot` u) a) +{-# INLINE mmul #-} + +mmulAdd :: BiMulAdd b a => Vector m (Vector n b) -> Vector n a -> Vector m a -> Vector m a +mmulAdd a u b = Vector.seq u (Vector.mapIx f b) + where + f (Offset i) bv = + let g (Offset j) vu = (index a (Offset i) (Offset j), vu) + in Vector.biMulFoldIndexWith g bv u +{-# INLINE mmulAdd #-} diff --git a/src/ScrubbedBlock.hs b/src/ScrubbedBlock.hs new file mode 100644 --- /dev/null +++ b/src/ScrubbedBlock.hs @@ -0,0 +1,90 @@ +-- | +-- Module : ScrubbedBlock +-- License : BSD-3-Clause +-- Copyright : (c) 2025 Olivier Chรฉron +-- +-- A block that is always pinned in memory and automatically erased by a +-- finalizer when not referenced anymore. Same pattern as ScrubbedBytes from +-- package memory but for blocks. +-- +-- A complication here is that we distinguish between mutable and immutable +-- values. And for resiliency against asynchronous exceptions, we need to +-- schedule block scrubbing with a finalizer right at the beginning when the +-- block is still in mutable form. Fortunately, for the perspective of the GC, +-- ByteArray# and MutableByteArray# are really the same heap object in disguise +-- and unsafeFreezeByteArray# is a true no-op. So the finalizer set on the +-- initial MutableByteArray# value gets transferred transparently to the final +-- ByteArray# form. +-- +-- See GHC note [primOpEffect of unsafe freezes and thaws] +-- +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE UnboxedTuples #-} +module ScrubbedBlock + ( ScrubbedBlock, foldZipWith, ScrubbedBlock.length + , new, thaw, unsafeFreeze + ) where + +import Data.Primitive.PrimArray as Block + +import Control.Exception (assert) +import Control.Monad.ST + +import Data.Word + +import Unsafe.Coerce + +import Base +import Block (Block, MutableBlock) +import qualified Block + +import GHC.Base (IO(IO)) +import GHC.Exts (mkWeak#) + +newtype ScrubbedBlock ty = ScrubbedBlock (Block ty) + deriving (Eq, Show) + +foldZipWith :: (PrimType a, PrimType b) + => (c -> a -> b -> c) -> c -> ScrubbedBlock a -> ScrubbedBlock b -> c +foldZipWith f c (ScrubbedBlock a) (ScrubbedBlock b) = + Block.foldZipWith f c a b +{-# INLINE foldZipWith #-} + +length :: PrimType ty => ScrubbedBlock ty -> CountOf ty +length (ScrubbedBlock b) = Block.length b + +new :: (PrimType ty, PrimMonad prim) => CountOf ty -> prim (MutableBlock ty (PrimState prim)) +new n = Block.newPinned n >>= scrubbed -- always pinned + +thaw :: PrimMonad m => ScrubbedBlock ty -> m (MutableBlock ty (PrimState m)) +thaw (ScrubbedBlock b) = Block.thawPinned b >>= scrubbed -- always pinned + +unsafeFreeze :: PrimMonad prim => MutableBlock ty (PrimState prim) -> prim (ScrubbedBlock ty) +unsafeFreeze mb = checkPinned <$> Block.unsafeFreeze mb + + +{- internal -} + +assertPinned :: Block ty -> a -> a +assertPinned mb = assert (Block.isPrimArrayPinned mb) + +checkPinned :: Block ty -> ScrubbedBlock ty +checkPinned b = assertPinned b (ScrubbedBlock b) + +scrubbed :: PrimMonad prim => MutableBlock ty (PrimState prim) -> prim (MutableBlock ty (PrimState prim)) +scrubbed b = unsafePrimFromIO (scheduleBlockScrubbing b >> return b) + +wakeUpAfterInception :: MutableBlock ty s -> MutableBlock ty RealWorld +wakeUpAfterInception = unsafeCoerce -- sometimes disappointing + +scheduleBlockScrubbing :: MutableBlock ty s -> IO () +scheduleBlockScrubbing b = addBlockFinalizer b (scrub $ Block.unsafeCastMut b') + where b' = wakeUpAfterInception b +{-# NOINLINE scheduleBlockScrubbing #-} + +scrub :: MutableBlock Word8 RealWorld -> IO () +scrub b = Block.getMutableLength b >>= \len -> Block.erase len b + +addBlockFinalizer :: MutableBlock ty s -> IO () -> IO () +addBlockFinalizer (Block.MutablePrimArray mbarr) (IO finalizer) = IO $ \s -> + case mkWeak# mbarr () finalizer s of { (# s1, _ #) -> (# s1, () #) } diff --git a/src/SecureBlock.hs b/src/SecureBlock.hs new file mode 100644 --- /dev/null +++ b/src/SecureBlock.hs @@ -0,0 +1,18 @@ +-- | +-- Module : SecureBlock +-- License : BSD-3-Clause +-- Copyright : (c) 2025 Olivier Chรฉron +-- +-- Use a type-level annotation to decide between a scrubbed block or a regular +-- block +-- +{-# LANGUAGE CPP #-} +module SecureBlock + ( SecureBlock, index, new, thaw + , unsafeCast, unsafeFreeze, toNormalForm +#ifdef ML_DSA_TESTING + , eq, Marking.showsPrec, toList +#endif + ) where + +import Marking diff --git a/src/SecureBytes.hs b/src/SecureBytes.hs new file mode 100644 --- /dev/null +++ b/src/SecureBytes.hs @@ -0,0 +1,16 @@ +-- | +-- Module : SecureBytes +-- License : BSD-3-Clause +-- Copyright : (c) 2025 Olivier Chรฉron +-- +-- Use a type-level annotation to decide between a scrubbed byte array or +-- a regular byte array +-- +module SecureBytes + ( SecureBytes, unsafeCreate, SecureBytes.length, copyByteArrayToPtr + ) where + +import Marking + +length :: Classified marking => SecureBytes marking -> Int +length = lengthBytes diff --git a/src/Vector.hs b/src/Vector.hs new file mode 100644 --- /dev/null +++ b/src/Vector.hs @@ -0,0 +1,253 @@ +-- | +-- Module : Vector +-- License : BSD-3-Clause +-- Copyright : (c) 2025 Olivier Chรฉron +-- +-- A vector of lifted elements with the vector dimension at type level. +-- Backed by type t'SmallArray' from primitive. +-- +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE TupleSections #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE ScopedTypeVariables #-} +module Vector + ( Vector, Vector.concatMap, Vector.dot, Vector.index, Vector.toNormalForm + , Vector.create, Vector.createMaybe, Vector.createMaybeAccum, Vector.foldl' + , Vector.fold1ZipWith, Vector.biMulFoldIndexWith, mapIx, Vector.seq + , Vector.unzipWith, Vector.zipWith +#ifdef ML_DSA_TESTING + , Vector.replicateM +#endif + ) where + +import Data.Primitive.SmallArray + +import Control.DeepSeq (NFData(..)) +#ifdef ML_DSA_TESTING +import Control.Monad +#endif +import Control.Monad.ST + +#if !(MIN_VERSION_base(4,20,0)) +import Data.List as Prelude (foldl') +#endif +import Data.Proxy + +import Base +import Iterate +import Math + +type Array = SmallArray +type MArray ty s = SmallMutableArray s ty + +newtype Vector (n :: Nat) a = Vector { unVector :: Array a } + deriving (Eq, Show) + +instance Functor (Vector n) where + fmap = mapVector + {-# INLINE fmap #-} + +instance (Add a, KnownNat n) => Add (Vector n a) where + zero = create (const zero) + {-# INLINE zero #-} + (.+) = Vector.zipWith (.+) + {-# INLINE (.+) #-} + (.-) = Vector.zipWith (.-) + {-# INLINE (.-) #-} + neg = mapVector neg + {-# INLINE neg #-} + +arrayCreate :: forall ty. CountOf ty -> (Offset ty -> ty) -> Array ty +arrayCreate n initializer = runST (arrayNew n >>= iter initializer) + where + iter :: PrimMonad prim => (Offset ty -> ty) -> MArray ty (PrimState prim) -> prim (Array ty) + iter f ma = loop 0 + where + loop s@(Offset i) + | s .==# n = unsafeFreezeSmallArray ma + | otherwise = writeSmallArray ma i (f s) >> loop (s + 1) + {-# INLINE loop #-} + {-# INLINE iter #-} + +arrayLength :: Array ty -> CountOf ty +arrayLength = CountOf . sizeofSmallArray + +arrayMapIx :: (Offset a -> a -> b) -> Array a -> Array b +arrayMapIx f a = arrayCreate (CountOf sz) $ \(Offset i) -> + let off = Offset i in f off (arrayIndex a off) + where CountOf sz = arrayLength a + +arrayNew :: PrimMonad prim => CountOf ty -> prim (MArray ty (PrimState prim)) +arrayNew (CountOf c) = newSmallArray c placeholder + where placeholder = error "arrayNew: unexpected evaluation" + +create :: forall n a. KnownNat n => (Offset a -> a) -> Vector n a +create f = Vector $ arrayCreate (CountOf sz) (\(Offset !i) -> f (Offset i)) + where !sz = fromIntegral $ natVal (Proxy :: Proxy n) +{-# INLINE [2] create #-} + +mapIx :: (Offset a -> a -> b) -> Vector n a -> Vector n b +mapIx = mapVectorIx +{-# INLINE [2] mapIx #-} + +mapVector :: (a -> b) -> Vector n a -> Vector n b +mapVector f = mapVectorIx $ \_ x -> f x +{-# INLINE [2] mapVector #-} + +mapVectorIx :: (Offset a -> a -> b) -> Vector n a -> Vector n b +mapVectorIx f = Vector <$> arrayMapIx f . unVector +{-# INLINE [1] mapVectorIx #-} + +arrayIndex :: Array a -> Offset a -> a +#ifdef ML_DSA_TESTING +arrayIndex a off@(Offset i) = + checkBounds (arrayLength a) off $ indexSmallArray a i + +replicateM :: forall n m a. (KnownNat n, Applicative m) => m a -> m (Vector n a) +replicateM f = Vector . smallArrayFromList <$> Control.Monad.replicateM sz f + where !sz = fromIntegral $ natVal (Proxy :: Proxy n) +#else +arrayIndex a (Offset i) = indexSmallArray a i +#endif + +index :: Vector n a -> Offset a -> a +index = arrayIndex . unVector + +concatMap :: Monoid b => (a -> b) -> Vector n a -> b +concatMap f = mconcat . mapToList f +{-# INLINE concatMap #-} + +mapToList :: (a -> b) -> Vector n a -> [b] +mapToList f (Vector a) = Prelude.map (f . arrayIndex a . Offset) (offsets sa) + where CountOf sa = arrayLength a + +foldl' :: (b -> a -> b) -> b -> Vector n a -> b +foldl' f b (Vector a) = Prelude.foldl' g b (offsets sa) + where + g acc i = acc `f` arrayIndex a (Offset i) + CountOf !sa = arrayLength a +{-# INLINE foldl' #-} + +seq :: b -> Vector n a -> Vector n a +seq = Prelude.seq +{-# INLINE [1] seq #-} + +zipWith :: (a -> b -> c) -> Vector n a -> Vector n b -> Vector n c +zipWith f a (Vector !b) = mapVectorIx g a + where g (Offset i) x = f x $ arrayIndex b (Offset i) +{-# INLINE [2] zipWith #-} + +fold1ZipWith :: (c -> a -> b -> c) -> (a -> b -> c) -> Vector n a -> Vector n b -> c +fold1ZipWith f g (Vector a) (Vector !b) = + Prelude.foldl' ff gg (offsetsFrom 1 sa) + where + ff x i = f x (arrayIndex a (Offset i)) (arrayIndex b (Offset i)) + gg = g (arrayIndex a 0) (arrayIndex b 0) + CountOf !sa = arrayLength a +{-# INLINE fold1ZipWith #-} + +biMulFoldIndexWith :: BiMulAdd b a => (Offset ty -> t -> (b, a)) -> a -> Vector n t -> a +biMulFoldIndexWith f c (Vector a) = + biMulFold c (map g $ offsets sa) + where + g i = f (Offset i) (arrayIndex a (Offset i)) + CountOf !sa = arrayLength a +{-# INLINE biMulFoldIndexWith #-} + +dot :: BiMulAdd b a => Vector n b -> Vector n a -> a +dot (Vector b) (Vector a) = + biMulFold (arrayIndex b 0 ..* arrayIndex a 0) (map g $ offsetsFrom 1 sb) + where + g i = (arrayIndex b (Offset i), arrayIndex a (Offset i)) + CountOf !sb = arrayLength b +{-# INLINE dot #-} + +toNormalForm :: NFData a => Vector n a -> () +toNormalForm = Vector.foldl' (\acc x -> acc `Prelude.seq` rnf x) () + +createMaybe :: KnownNat n => (Offset a -> Maybe a) -> Maybe (Vector n a) +createMaybe f = createMaybeAccum (\_ off -> ((), ) <$> f off) (const Just) () +{-# INLINE createMaybe #-} + +createMaybeAccum :: forall n a b r. KnownNat n => (b -> Offset a -> Maybe (b, a)) -> (b -> Vector n a -> Maybe r) -> b -> Maybe r +createMaybeAccum f k b0 = runST $ arrayNew (CountOf n) >>= \ma -> loop ma b0 0 + where + !n = fromIntegral $ natVal (Proxy :: Proxy n) + + loop :: MArray a s -> b -> Offset a -> ST s (Maybe r) + loop !ma b s@(Offset i) + | s .==# CountOf n = + unsafeFreezeSmallArray ma >>= \a -> return (b `k` Vector a) + | otherwise = do + case f b s of + Nothing -> return Nothing + Just (b', a) -> do + writeSmallArray ma i a + loop ma b' (s + 1) +{-# INLINE createMaybeAccum #-} + +unzipWith :: forall n a b c. KnownNat n => (a -> (b, c)) -> Vector n a -> (Vector n b, Vector n c) +unzipWith f (Vector !a) = runST $ do + mb <- arrayNew (CountOf n) + mc <- arrayNew (CountOf n) + loop mb mc 0 + where + !n = fromIntegral $ natVal (Proxy :: Proxy n) + + loop :: MArray b s -> MArray c s -> Offset a -> ST s (Vector n b, Vector n c) + loop !mb !mc s@(Offset i) + | s .==# CountOf n = do + b <- unsafeFreezeSmallArray mb + c <- unsafeFreezeSmallArray mc + return (Vector b, Vector c) + | otherwise = do + let (b, c) = f (arrayIndex a s) + writeSmallArray mb i b + writeSmallArray mc i c + loop mb mc (s + 1) +{-# INLINE unzipWith #-} + + +-- Rewrite rules +-- +-- A first set of rules before Phase 2 performs simplifications between the four +-- main functions: create, mapVector, mapIx, and zipWith. +-- +-- During Phase 2, the functions are then inlined. While function create +-- is replaced with its final form as call to arrayCreate, the three other +-- functions mapVector, mapIx, and zipWith all become calls to mapVectorIx. +-- Then, nested calls to mapVectorIx can further be simplified using rule +-- "mapVectorIx/mapVectorIx". +-- +-- Finally at Phase 1 the function mapVectorIx is replaced with a call to +-- arrayCreate. +-- +-- Both layers of rules have transformations that push calls to Vector.seq +-- outwards. Normal 'Prelude.seq' or the use of bang patterns would prevent +-- rewrite rules from firing. + +{-# RULES +"mapVector/mapVector" [~2] forall f g a. mapVector f (mapVector g a) = mapVector (f . g) a +"mapVector/mapIx" [~2] forall f g a. mapVector f (mapIx g a) = mapIx (\i -> f . g i) a +"mapVector/create" [~2] forall f g. mapVector f (create g) = create (\(Offset i) -> f (g (Offset i))) +"mapVector/zipWith" [~2] forall f g a b. mapVector f (Vector.zipWith g a b) = Vector.zipWith (\x -> f . g x) a b +"mapVector/seq" [~2] forall f a b. mapVector f (Vector.seq b a) = Vector.seq b (mapVector f a) + +"zipWith/mapVector left" [~2] forall f g a. Vector.zipWith f (mapVector g a) = Vector.zipWith (f . g) a +"zipWith/mapVector right" [~2] forall f g a b. Vector.zipWith f a (mapVector g b) = Vector.zipWith (\aa bb -> f aa (g bb)) a b +"zipWith/create left" [~2] forall f g. Vector.zipWith f (create g) = mapIx (\(Offset i) -> f (g (Offset i))) +"zipWith/create right" [~2] forall f g a. Vector.zipWith f a (create g) = mapIx (\(Offset i) x -> f x (g (Offset i))) a +"zipWith/zipWith right" [~2] forall f g a b c. Vector.zipWith f a (Vector.zipWith g b c) = Vector.zipWith (flip f) (Vector.zipWith g b c) a +"zipWith/seq left" [~2] forall f a b c. Vector.zipWith f (Vector.seq c a) b = Vector.seq c (Vector.zipWith f a b) +"zipWith/seq right" [~2] forall f a b c. Vector.zipWith f a (Vector.seq c b) = Vector.seq c (Vector.zipWith f a b) + +"mapIx/mapVector" [~2] forall f g a. mapIx f (mapVector g a) = mapIx (\(Offset i) -> f (Offset i) . g) a +"mapIx/mapIx" [~2] forall f g a. mapIx f (mapIx g a) = mapIx (\(Offset i) -> f (Offset i) . g (Offset i)) a +"mapIx/create" [~2] forall f g. mapIx f (create g) = create (\(Offset i) -> f (Offset i) (g (Offset i))) +"mapIx/seq" [~2] forall f a b. mapIx f (Vector.seq b a) = Vector.seq b (mapIx f a) + +"mapVectorIx/mapVectorIx" [~1] forall f g a. mapVectorIx f (mapVectorIx g a) = mapVectorIx (\(Offset i) -> f (Offset i) . g (Offset i)) a +"mapVectorIx/seq" [~1] forall f a b. mapVectorIx f (Vector.seq b a) = Vector.seq b (mapVectorIx f a) + #-} diff --git a/tests/KeyGen.hs b/tests/KeyGen.hs new file mode 100644 --- /dev/null +++ b/tests/KeyGen.hs @@ -0,0 +1,46 @@ +-- | +-- Module : KeyGen +-- License : BSD-3-Clause +-- Copyright : (c) 2025 Olivier Chรฉron +-- +-- Types that are specific to key-generation test vectors +-- +{-# LANGUAGE OverloadedStrings #-} +module KeyGen + ( TestGroup(..), Test(..) + ) where + +import Data.Aeson +import Data.ByteString (ByteString) + +import Util + +data TestGroup = TestGroup + { tgId :: Int + , testType :: String + , parameterSet :: String + , tests :: [Test] + } deriving Show + +instance FromJSON TestGroup where + parseJSON = withObject "TestGroup" $ \o -> TestGroup + <$> o .: "tgId" + <*> o .: "testType" + <*> o .: "parameterSet" + <*> o .: "tests" + +data Test = Test + { tcId :: Int + , deferred :: Bool + , seed :: ByteString + , pk :: ByteString + , sk :: ByteString + } deriving Show + +instance FromJSON Test where + parseJSON = withObject "Test" $ \o -> Test + <$> o .: "tcId" + <*> o .: "deferred" + <*> o .:: "seed" + <*> o .:: "pk" + <*> o .:: "sk" diff --git a/tests/SigExt.hs b/tests/SigExt.hs new file mode 100644 --- /dev/null +++ b/tests/SigExt.hs @@ -0,0 +1,57 @@ +-- | +-- Module : SigExt +-- License : BSD-3-Clause +-- Copyright : (c) 2026 Olivier Chรฉron +-- +-- Types that are common to signature-generation and verification test vectors +-- +{-# LANGUAGE OverloadedStrings #-} +module SigExt + ( TestExt(..) , PureExt(..), PreHashExt(..), ExternalMuExt(..) + , InternalExt(..) + ) where + +import Data.Aeson +import Data.Aeson.Types +import Data.ByteString (ByteString) + +import Util + +class TestExt ext where + parseExt :: Object -> Parser ext + +data PureExt = PureExt + { msgEncP :: ByteString + , ctxEncP :: ByteString + } deriving Show + +instance TestExt PureExt where + parseExt o = PureExt + <$> o .:: "message" + <*> o .:: "context" + +data PreHashExt = PreHashExt + { msgEncPH :: ByteString + , ctxEncPH :: ByteString + , hashPH :: String + } deriving Show + +instance TestExt PreHashExt where + parseExt o = PreHashExt + <$> o .:: "message" + <*> o .:: "context" + <*> o .: "hashAlg" + +newtype ExternalMuExt = ExternalMuExt + { muEncEM :: ByteString + } deriving Show + +instance TestExt ExternalMuExt where + parseExt o = ExternalMuExt <$> o .:: "mu" + +newtype InternalExt = InternalExt + { msgEncIM :: ByteString + } deriving Show + +instance TestExt InternalExt where + parseExt o = InternalExt <$> o .:: "message" diff --git a/tests/SigGen.hs b/tests/SigGen.hs new file mode 100644 --- /dev/null +++ b/tests/SigGen.hs @@ -0,0 +1,81 @@ +-- | +-- Module : SigVer +-- License : BSD-3-Clause +-- Copyright : (c) 2026 Olivier Chรฉron +-- +-- Types that are specific to signature-generation test vectors +-- +{-# LANGUAGE OverloadedStrings #-} +module SigGen + ( TestGroup(..), Test(..), TestGroupPayload(..) + ) where + +import Data.Aeson +import Data.Aeson.Types +import Data.ByteString (ByteString) + +import SigExt +import Util + +data TestGroup = TestGroup + { tgId :: Int + , testType :: String + , parameterSet :: String + , deterministic :: Bool + , signatureInterface :: String + , preHash :: String + , externalMu :: Bool + , cornerCase :: String + , payload :: TestGroupPayload + } deriving Show + +data TestGroupPayload + = ModePure [Test PureExt] + | ModePreHash [Test PreHashExt] + | ModeExternalMu [Test ExternalMuExt] + | ModeInternal [Test InternalExt] + deriving Show + +parsePayload :: Object -> Parser TestGroupPayload +parsePayload o = do + si <- o .: "signatureInterface" + ph <- o .: "preHash" + em <- o .: "externalMu" + case (si :: String, ph :: String, em) of + ("external", "pure", False) -> ModePure <$> (o .: "tests") + ("external", "preHash", False) -> ModePreHash <$> (o .: "tests") + ("internal", "none", True) -> ModeExternalMu <$> (o .: "tests") + ("internal", "none", False) -> ModeInternal <$> (o .: "tests") + unknown -> fail ("parsePayload: unknown mode " ++ show unknown) + +instance FromJSON TestGroup where + parseJSON = withObject "TestGroup" $ \o -> TestGroup + <$> o .: "tgId" + <*> o .: "testType" + <*> o .: "parameterSet" + <*> o .: "deterministic" + <*> o .: "signatureInterface" + <*> o .: "preHash" + <*> o .: "externalMu" + <*> o .: "cornerCase" + <*> parsePayload o + +data Test ext = Test + { tcId :: Int + , deferred :: Bool + , pkEnc :: ByteString + , skEnc :: ByteString + , sigEnc :: ByteString + , rndEnc :: Maybe ByteString + , tcExt :: ext + } deriving Show + +instance TestExt ext => FromJSON (Test ext) where + parseJSON = withObject "Test" $ \o -> Test + <$> o .: "tcId" + <*> o .: "deferred" + <*> o .:: "pk" + <*> o .:: "sk" + <*> o .:: "signature" + <*> o .::? "rnd" + <*> parseExt o diff --git a/tests/SigVer.hs b/tests/SigVer.hs new file mode 100644 --- /dev/null +++ b/tests/SigVer.hs @@ -0,0 +1,79 @@ +-- | +-- Module : SigVer +-- License : BSD-3-Clause +-- Copyright : (c) 2026 Olivier Chรฉron +-- +-- Types that are specific to signature-verification test vectors +-- +{-# LANGUAGE OverloadedStrings #-} +module SigVer + ( TestGroup(..), Test(..), TestGroupPayload(..) + ) where + +import Data.Aeson +import Data.Aeson.Types +import Data.ByteString (ByteString) + +import SigExt +import Util + +data TestGroup = TestGroup + { tgId :: Int + , testType :: String + , parameterSet :: String + , signatureInterface :: String + , preHash :: String + , externalMu :: Bool + , payload :: TestGroupPayload + } deriving Show + +data TestGroupPayload + = ModePure [Test PureExt] + | ModePreHash [Test PreHashExt] + | ModeExternalMu [Test ExternalMuExt] + | ModeInternal [Test InternalExt] + deriving Show + +parsePayload :: Object -> Parser TestGroupPayload +parsePayload o = do + si <- o .: "signatureInterface" + ph <- o .: "preHash" + em <- o .: "externalMu" + case (si :: String, ph :: String, em) of + ("external", "pure", False) -> ModePure <$> (o .: "tests") + ("external", "preHash", False) -> ModePreHash <$> (o .: "tests") + ("internal", "none", True) -> ModeExternalMu <$> (o .: "tests") + ("internal", "none", False) -> ModeInternal <$> (o .: "tests") + unknown -> fail ("parsePayload: unknown mode " ++ show unknown) + +instance FromJSON TestGroup where + parseJSON = withObject "TestGroup" $ \o -> TestGroup + <$> o .: "tgId" + <*> o .: "testType" + <*> o .: "parameterSet" + <*> o .: "signatureInterface" + <*> o .: "preHash" + <*> o .: "externalMu" + <*> parsePayload o + +data Test ext = Test + { tcId :: Int + , testPassed :: Bool + , deferred :: Bool + , pkEnc :: ByteString + , skEnc :: ByteString + , sigEnc :: ByteString + , reason :: String + , tcExt :: ext + } deriving Show + +instance TestExt ext => FromJSON (Test ext) where + parseJSON = withObject "Test" $ \o -> Test + <$> o .: "tcId" + <*> o .: "testPassed" + <*> o .: "deferred" + <*> o .:: "pk" + <*> o .:: "sk" + <*> o .:: "signature" + <*> o .: "reason" + <*> parseExt o diff --git a/tests/Tests.hs b/tests/Tests.hs new file mode 100644 --- /dev/null +++ b/tests/Tests.hs @@ -0,0 +1,739 @@ +-- | +-- Module : Main +-- License : BSD-3-Clause +-- Copyright : (c) 2026 Olivier Chรฉron +-- +-- The ML-DSA test suite. Can be instanciated twice, with and without the +-- @ML_DSA_TESTING@ macro to run property testing with assertions enabled in +-- the internal modules. +-- +{-# LANGUAGE CPP #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE ExistentialQuantification #-} +{-# LANGUAGE KindSignatures #-} +{-# LANGUAGE RankNTypes #-} +module Main (main) where + +import Data.ByteArray (Bytes) +import qualified Data.ByteArray as B + +import Crypto.Hash (hashWith) +import Crypto.Hash.Algorithms + +import Control.Monad + +import Data.List (isPrefixOf) +import Data.Maybe (fromJust) +import Data.Proxy + +import GHC.IO.Exception (IOErrorType(..)) + +import System.Directory (doesFileExist) +import System.FileLock (FileLock, SharedExclusive(..), unlockFile, withFileLock) +import System.IO.Error (catchIOError, mkIOError) +import System.Process (readProcess) + +#ifdef ML_DSA_TESTING + +import Data.Bits +import Data.Word + +import GHC.TypeNats + +import Foreign.Ptr (plusPtr) + +import Auxiliary +import BlockN (BlockN) +import Builder (Builder) +import Marking (Leak(..), SecurityMarking(..)) +import Math +import Matrix +import Vector +import qualified BlockN +import qualified Builder +#endif + +import Crypto.PubKey.ML_DSA as Lib + +import Test.Tasty +import Test.Tasty.HUnit +import Test.Tasty.QuickCheck + +import qualified KeyGen +import qualified SigExt +import qualified SigGen +import qualified SigVer +import qualified Vectors + +arbitraryBytes :: Int -> Gen Bytes +arbitraryBytes n = B.pack <$> vectorOf n arbitrary + +#ifdef ML_DSA_TESTING + +truncateRq :: Word32 -> Rq Sec -> Rq Sec +truncateRq m = fromJust . fromCoeffs . map f . toCoeffs + where f = toZq . (`mod` m) . fromZq + +truncateRq' :: Word32 -> Rq Sec -> Rq Sec +truncateRq' m = fromJust . fromCoeffs . map f . toCoeffs + where f = toZq . (\x -> x `mod` (2 * m - 1) + 8380418 - m) . fromZq + +newtype FE = FE { unFE :: Zq } deriving Show + +instance Arbitrary FE where +#if (MIN_VERSION_tasty_quickcheck(0,10,2)) + arbitrary = FE . toZq <$> chooseBoundedIntegral (0, 8380416) +#else + arbitrary = FE . toZq <$> choose (0, 8380416) +#endif + +newtype ME = ME { unME :: Mq } deriving Show + +instance Arbitrary ME where + arbitrary = ME . toMontgomery . unFE <$> arbitrary + +newtype Poly = Poly (Rq Sec) deriving Show + +instance Arbitrary Poly where + arbitrary = do + coeffs <- map unFE <$> vectorOf 256 arbitrary + let a = fromJust (fromCoeffs coeffs) + return (Poly a) + +newtype PolyNTT = PolyNTT (Tq Sec) deriving Show + +instance Arbitrary PolyNTT where + arbitrary = (\(Poly f) -> PolyNTT (ntt f)) <$> arbitrary + +arbitraryHints :: Gen Hints +arbitraryHints = fromJust . fromBools <$> vectorOf 256 h + where h = elements (True : replicate 9 False) -- probability 1/10 + +arbitraryHintsVector :: KnownNat k => proxy k -> Gen (Vector k Hints) +arbitraryHintsVector _ = Vector.replicateM arbitraryHints + +countHints :: Vector n Hints -> Word +countHints = Vector.foldl' countFrom 0 + +newtype D = D Int deriving Show + +instance Arbitrary D where + arbitrary = sized $ \n -> D <$> choose (0, min n 22) + +data Dim = forall (n :: Nat). KnownNat n => Dim (Proxy n) + +instance Show Dim where + show (Dim n) = show n + +instance Arbitrary Dim where + arbitrary = sized $ \n -> toDim <$> choose (1, min (1 + n) 9) + +toDim :: Int -> Dim +toDim n = case someNatVal (fromIntegral n) of SomeNat p -> Dim p + +type VElem = Mq -- test with any ring but Tq would also work here + +arbitraryVector :: KnownNat n => proxy n -> Gen (Vector n VElem) +arbitraryVector _ = Vector.replicateM (unME <$> arbitrary) + +arbitraryMatrix :: (KnownNat m, KnownNat n) => proxy n -> proxy m -> Gen (Vector n (Vector m VElem)) +arbitraryMatrix _ m = Vector.replicateM (arbitraryVector m) + +simpleBitPackBytes :: Int -> BlockN Sec 256 Word32 -> Bytes +simpleBitPackBytes d = runBytes . simpleBitPack d + +runPubBytes :: Builder Pub -> Bytes +runPubBytes = Builder.run + +runPubUnaligned :: Builder Pub -> UnalignedBytes +runPubUnaligned = Builder.runRelaxed + +runBytes :: Builder Sec -> Bytes +runBytes = runPubBytes . leak + +runUnaligned :: Builder Sec -> UnalignedBytes +runUnaligned = runPubUnaligned . leak + +newtype UnalignedBytes = UnalignedBytes (B.View Bytes) + deriving (Eq,Ord,Show) + +instance Semigroup UnalignedBytes where + UnalignedBytes a <> UnalignedBytes b = + B.allocAndFreeze (na + nb) $ \p -> do + B.copyByteArrayToPtr a p + B.copyByteArrayToPtr b (p `plusPtr` na) + where + na = B.length a + nb = B.length b + +instance Monoid UnalignedBytes where + mempty = B.convert (B.empty :: Bytes) + +instance B.ByteArrayAccess UnalignedBytes where + length (UnalignedBytes v) = B.length v + withByteArray (UnalignedBytes v) = B.withByteArray v + +instance B.ByteArray UnalignedBytes where + allocRet n f = do + let build ba = UnalignedBytes (B.dropView ba offset) + (a, ba) <- B.allocRet (n + offset) $ \p -> f (p `plusPtr` offset) + return (a, build ba) + where offset = 1 + +#endif + +data P = forall a. (ParamSet a, Show a) => P (Proxy a) + +instance Show P where + show (P p) = show p + +instance Arbitrary P where + arbitrary = elements + [ P (Proxy :: Proxy ML_DSA_44) + , P (Proxy :: Proxy ML_DSA_65) + , P (Proxy :: Proxy ML_DSA_87) + ] + +toP :: String -> P +toP "ML-DSA-44" = P (Proxy :: Proxy ML_DSA_44) +toP "ML-DSA-65" = P (Proxy :: Proxy ML_DSA_65) +toP "ML-DSA-87" = P (Proxy :: Proxy ML_DSA_87) +toP paramSet = error ("unknown parameter set " ++ paramSet) + +data PH = forall alg. (PreHashAlgorithm alg, Show alg) => PH (Proxy alg) + +instance Show PH where + show (PH ph) = show ph + +instance Arbitrary PH where + arbitrary = elements + [ PH (Proxy :: Proxy SHA224) + , PH (Proxy :: Proxy SHA256) + , PH (Proxy :: Proxy SHA384) + , PH (Proxy :: Proxy SHA512) + , PH (Proxy :: Proxy SHA512t_224) + , PH (Proxy :: Proxy SHA512t_256) + , PH (Proxy :: Proxy SHA3_224) + , PH (Proxy :: Proxy SHA3_256) + , PH (Proxy :: Proxy SHA3_384) + , PH (Proxy :: Proxy SHA3_512) + , PH (Proxy :: Proxy (SHAKE128 256)) + , PH (Proxy :: Proxy (SHAKE256 512)) + ] + +toPH :: String -> PH +toPH "SHA2-224" = PH (Proxy :: Proxy SHA224) +toPH "SHA2-256" = PH (Proxy :: Proxy SHA256) +toPH "SHA2-384" = PH (Proxy :: Proxy SHA384) +toPH "SHA2-512" = PH (Proxy :: Proxy SHA512) +toPH "SHA2-512/224" = PH (Proxy :: Proxy SHA512t_224) +toPH "SHA2-512/256" = PH (Proxy :: Proxy SHA512t_256) +toPH "SHA3-224" = PH (Proxy :: Proxy SHA3_224) +toPH "SHA3-256" = PH (Proxy :: Proxy SHA3_256) +toPH "SHA3-384" = PH (Proxy :: Proxy SHA3_384) +toPH "SHA3-512" = PH (Proxy :: Proxy SHA3_512) +toPH "SHAKE-128" = PH (Proxy :: Proxy (SHAKE128 256)) +toPH "SHAKE-256" = PH (Proxy :: Proxy (SHAKE256 512)) +toPH hashAlg = error ("unknown hash algorithm " ++ hashAlg) + +preHash :: Proxy alg -> alg +preHash _ = undefined + +withPreHash :: String -> (forall alg. PreHashAlgorithm alg => alg -> r) -> r +withPreHash s f = case toPH s of + PH ph -> f (preHash ph) + +newtype Ctx = Ctx { unCtx :: Context } deriving Show + +instance Arbitrary Ctx where + arbitrary = sized $ \n -> do + len <- choose (0, min n 255) + Ctx . fromJust . Lib.context <$> arbitraryBytes len + +newtype Msg = Msg { unMsg :: Bytes } deriving Show + +instance Arbitrary Msg where + arbitrary = Msg . B.pack <$> arbitrary + +newtype ExternalMu = ExternalMu { unExternalMu :: Mu } deriving Show + +instance Arbitrary ExternalMu where + arbitrary = ExternalMu . fromJust . Lib.externalMu <$> arbitraryBytes 64 + +withVectors :: (IO () -> TestTree) -> TestTree +withVectors = withResource alloc free + where + scriptPath = "tests/get-vectors.sh" + free _ = return () + whenNeeded action = do + keyGenExists <- doesFileExist "tests/keyGen.json.gz" + sigGenExists <- doesFileExist "tests/sigGen.json.gz" + sigVerExists <- doesFileExist "tests/sigVer.json.gz" + unless (keyGenExists && sigGenExists && sigVerExists) action + alloc = withTestLock Shared $ \lock -> whenNeeded $ do + unlockFile lock -- sanity before lock upgrade + withTestLock Exclusive $ \_ -> whenNeeded $ catchIOError + (void $ readProcess "/bin/sh" [scriptPath] "") + (\e -> + let msg = "Could not download test vectors, you will need to run the script `" ++ + scriptPath ++ "' manually. Script failure was: " ++ show e + in ioError (mkIOError OtherError msg Nothing Nothing) + ) + +withTestLock :: SharedExclusive -> (FileLock -> IO a) -> IO a +withTestLock mode what = do + -- locking a hidden file in the directory prevents a race condition between + -- two instances of the test suite trying both to download the test vectors + -- (otherwise one instance may try to run with files not fully downloaded + -- yet by the other instance) + let path = "tests/.lock" + exists <- doesFileExist path + unless exists $ writeFile path "DO NOT DELETE" + withFileLock path mode what + +keyGenVectors :: (String -> IO ()) -> Assertion +keyGenVectors step = do + step "Reading test vectors ..." + file <- Vectors.readJson "tests/keyGen.json.gz" + forM_ (Vectors.testGroups file) $ \group -> do + let paramSet = KeyGen.parameterSet group + step paramSet + case toP paramSet of + P p -> forM_ (KeyGen.tests group) $ \t -> do + let tcId = KeyGen.tcId t + pks = Lib.encode pk + sks = Lib.encode sk + (pk, sk) = fromJust $ Lib.generateWith p (KeyGen.seed t) + assertEqual ("pk mismatch for tcId=" ++ show tcId) (KeyGen.pk t) pks + assertEqual ("sk mismatch for tcId=" ++ show tcId) (KeyGen.sk t) sks + +sigGenVectors :: (String -> IO ()) -> Assertion +sigGenVectors step = do + step "Reading test vectors ..." + file <- Vectors.readJson "tests/sigGen.json.gz" + forM_ (Vectors.testGroups file) $ \group -> do + let paramSet = SigGen.parameterSet group + step (paramSet ++ " (" ++ mode group ++ ")") + case toP paramSet of + P p -> case SigGen.payload group of + SigGen.ModePure tests -> + forM_ tests (testExt doSignPure p) + SigGen.ModePreHash tests -> + forM_ tests (testExt doSignPreHash p) + SigGen.ModeExternalMu tests -> + forM_ tests (testExt doSignExternalMu p) + SigGen.ModeInternal tests -> + forM_ tests (testExt doSignInternal p) + where + mode group + | SigGen.signatureInterface group == "external" = SigGen.preHash group + | SigGen.externalMu group = "external ยต" + | otherwise = SigGen.signatureInterface group + testExt signExtWith p test = + assertEqual ("sig mismatch for tcId=" ++ show tcId) sig' sig + where + tcId = SigGen.tcId test + ext = SigGen.tcExt test + sk = fromJust $ Lib.decode p (SigGen.skEnc test) + sig' = fromJust $ Lib.decode p (SigGen.sigEnc test) + rnd = maybe Lib.deterministic (fromJust . Lib.randomness) (SigGen.rndEnc test) + sig = signExtWith ext rnd sk + doSignPure ext rnd sk = Lib.signWith rnd sk m ctx + where + m = SigExt.msgEncP ext + ctx = fromJust $ Lib.context (SigExt.ctxEncP ext) + doSignPreHash ext rnd sk = withPreHash (SigExt.hashPH ext) $ \alg -> + let phm = hashWith alg m + in Lib.signDigestWith rnd sk phm ctx + where + m = SigExt.msgEncPH ext + ctx = fromJust $ Lib.context (SigExt.ctxEncPH ext) + doSignExternalMu ext rnd sk = Lib.signExternalMuWith rnd sk mu + where mu = fromJust $ Lib.externalMu (SigExt.muEncEM ext) + doSignInternal ext rnd sk = Lib.signInternalWith rnd sk m + where m = SigExt.msgEncIM ext + +sigVerVectors :: (String -> IO ()) -> Assertion +sigVerVectors step = do + step "Reading test vectors ..." + file <- Vectors.readJson "tests/sigVer.json.gz" + forM_ (Vectors.testGroups file) $ \group -> do + let paramSet = SigVer.parameterSet group + step (paramSet ++ " (" ++ mode group ++ ")") + case toP paramSet of + P p -> case SigVer.payload group of + SigVer.ModePure tests -> + forM_ tests (testExt doVerifyPure p) + SigVer.ModePreHash tests -> + forM_ tests (testExt doVerifyPreHash p) + SigVer.ModeExternalMu tests -> + forM_ tests (testExt doVerifyExternalMu p) + SigVer.ModeInternal tests -> + forM_ tests (testExt doVerifyInternal p) + where + mode group + | SigVer.signatureInterface group == "external" = SigVer.preHash group + | SigVer.externalMu group = "external ยต" + | otherwise = SigVer.signatureInterface group + testExt doVerify p test = + case Lib.decode p (SigVer.sigEnc test) of + Just sig -> do + assertBool ("opposite outcome for tcId=" ++ show tcId) + (SigVer.testPassed test == doVerify ext pk sig) + Nothing -> + assertBool ("could not decode signature for tcId=" ++ show tcId) + ("modified signature - " `isPrefixOf` SigVer.reason test) + where + tcId = SigVer.tcId test + ext = SigVer.tcExt test + pk = fromJust $ Lib.decode p (SigVer.pkEnc test) + doVerifyPure ext pk sig = Lib.verify pk m sig ctx + where + m = SigExt.msgEncP ext + ctx = fromJust $ Lib.context (SigExt.ctxEncP ext) + doVerifyPreHash ext pk sig = withPreHash (SigExt.hashPH ext) $ \alg -> + let phm = hashWith alg m + in Lib.verifyDigest pk phm sig ctx + where + m = SigExt.msgEncPH ext + ctx = fromJust $ Lib.context (SigExt.ctxEncPH ext) + doVerifyExternalMu ext pk = Lib.verifyExternalMu pk mu + where mu = fromJust $ Lib.externalMu (SigExt.muEncEM ext) + doVerifyInternal ext pk = Lib.verifyInternal pk m + where m = SigExt.msgEncIM ext + +main :: IO () +main = defaultMain $ testGroup "mldsa" + [ withVectors $ \_ -> testGroup "vectors" + [ testCaseSteps "keyGen" keyGenVectors + , testCaseSteps "sigGen" sigGenVectors + , testCaseSteps "sigVer" sigVerVectors + ] + , testGroup "properties" + [ testGroup "ML-DSA" + [ testProperty "sign/verify (pure)" $ \(P p) (Msg m) (Ctx ctx) -> ioProperty $ do + (pk, sk) <- Lib.generate p + sig <- Lib.sign sk m ctx + return (Lib.verify pk m sig ctx) + , testProperty "sign/verify (preHash)" $ \(P p) (PH ph) (Msg m) (Ctx ctx) -> ioProperty $ do + (pk, sk) <- Lib.generate p + let phm = hashWith (preHash ph) m + sig <- Lib.signDigest sk phm ctx + return (Lib.verifyDigest pk phm sig ctx) + , testProperty "sign/verify (external ยต)" $ \(P p) (ExternalMu mu) -> ioProperty $ do + (pk, sk) <- Lib.generate p + sig <- Lib.signExternalMu sk mu + return (Lib.verifyExternalMu pk mu sig) + , testProperty "sign/verify (internal)" $ \(P p) (Msg m) -> ioProperty $ do + (pk, sk) <- Lib.generate p + sig <- Lib.signInternal sk m + return (Lib.verifyInternal pk m sig) + , testProperty "encode/decode keys" $ \(P p) -> ioProperty $ do + (pk, sk) <- Lib.generate p + return $ conjoin + [ Just pk === Lib.decode p (Lib.encode pk :: Bytes) + , Just sk === Lib.decode p (Lib.encode sk :: Bytes) + ] + , testProperty "encode/decode signatures" $ \(P p) (Msg m) (Ctx ctx) -> ioProperty $ do + (_, sk) <- Lib.generate p + sig <- Lib.sign sk m ctx + return $ Just sig === Lib.decode p (Lib.encode sig :: Bytes) + , testProperty "toPublic" $ \(P p) -> ioProperty $ do + (ek, sk) <- Lib.generate p + return (ek === toPublic sk) + , testProperty "checkKeyPair" $ \(P p) -> ioProperty $ + checkKeyPair <$> Lib.generate p + ] +#ifdef ML_DSA_TESTING + , testGroup "bitRev8" + [ testCase "powers of two" $ + let powers = [1, 2, 4, 8, 16, 32, 64, 128] + in reverse powers @=? map bitRev8 powers + , testProperty "or" $ \a b -> + bitRev8 (a .|. b) === bitRev8 a .|. bitRev8 b + , testProperty "not" $ \a -> + let comp = xor 255 + in bitRev8 (comp a) === comp (bitRev8 a) + , testProperty "involutive" $ \a -> + a === bitRev8 (bitRev8 a) + , testProperty "preserves bit count" $ \a -> + popCount a === popCount (bitRev8 a) + ] + , testGroup "compression" + [ testProperty "powerTwoRound" $ \(FE r) -> + let (r1, r0) = powerTwoRoundZq r + in fromZq r === (fromZq r1 + fromZq r0) `mod` 8380417 .&&. + (fromZq r1 `mod` 8192 == 0) .&&. + (fromZq r0 <= 4096 .||. fromZq r0 > 8380417 - 4096) + , testProperty "decompose" $ \(FE r) -> do + gamma2 <- elements [ 95232, 261888 ] + let (r1, r0) = decomposeZq gamma2 r + return $ r === toZq (2 * gamma2 * r1) .+ r0 .&&. + (fromZq r0 <= gamma2 .||. fromZq r0 > 8380417 - gamma2) + ] + , testGroup "hints" + [ testProperty "fromBools . toBools == id " $ do + a <- arbitraryHints + return $ Just a === fromBools (toBools a) + , testProperty "toBools . fromBools == id " $ do + a <- vector 256 + return $ Just a === (toBools <$> fromBools a) + , testProperty "useHint . makeHint" $ \(Poly z') (Poly r) -> do + gamma2 <- elements [ 95232, 261888 ] + let z = truncateRq' gamma2 z' -- to ensure |z| < gamma2 + return $ useHint gamma2 (makeHint gamma2 z r) r === highBits gamma2 (r .+ z) + ] + , testGroup "conversions" + [ testProperty "simpleBitPack . simpleBitUnpack == id" $ \(D d) -> do + b <- arbitraryBytes (32 * d) + return (b === runBytes (simpleBitPack d (simpleBitUnpack d b))) + , testProperty "simpleBitPack . simpleBitUnpack == id (unaligned)" $ \(D d) -> do + b <- arbitraryBytes (32 * d) + return (B.convert b === runUnaligned (simpleBitPack d (simpleBitUnpack d b))) + , testProperty "simpleBitPack 8" $ \x -> + B.replicate 256 x === simpleBitPackBytes 8 (BlockN.replicate $ fromIntegral x) + , testCase "simpleBitPack 1 (zeros)" $ + B.replicate 32 0 @=? simpleBitPackBytes 1 (BlockN.replicate 0) + , testCase "simpleBitPack 1 (ones)" $ + B.replicate 32 255 @=? simpleBitPackBytes 1 (BlockN.replicate 1) + , testProperty "simpleBitPack10 . simpleBitUnpack10 == id" $ do + b <- arbitraryBytes 320 + return (b === runPubBytes (simpleBitPack10 (simpleBitUnpack10 b))) + , testProperty "simpleBitPack10 . simpleBitUnpack10 == id (unaligned)" $ do + b <- arbitraryBytes 320 + return (B.convert b === runPubUnaligned (simpleBitPack10 (simpleBitUnpack10 b))) + , testProperty "bitPackSafe . bitUnpackSafe == id" $ \(D d) -> do + b <- arbitraryBytes (32 * d) + return (b === runBytes (bitPackSafe d (bitUnpackSafe d b))) + , testProperty "bitPackSafe . bitUnpackSafe == id (unaligned)" $ \(D d) -> do + b <- arbitraryBytes (32 * d) + return (B.convert b === runUnaligned (bitPackSafe d (bitUnpackSafe d b))) + , testProperty "bitUnpack . bitPack == id" $ \(FE t) (Poly w') -> + let w = truncateRq (fromZq t) w' + m = getNorm (norm w) + d = 33 - countLeadingZeros m + in Just w === bitUnpack m d (runBytes $ bitPack m d w) + , testProperty "bitUnpack . bitPack == id (unaligned)" $ \(FE t) (Poly w') -> + let w = truncateRq (fromZq t) w' + m = getNorm (norm w) + d = 33 - countLeadingZeros m + in Just w === bitUnpack m d (runUnaligned $ bitPack m d w) + , testProperty "hintBitUnpack . hintBitPack == id" $ \(Dim k) extra -> do + h <- arbitraryHintsVector k + let omega = fromIntegral (extra + countHints h) + return $ omega < 256 ==> + Just h === hintBitUnpack omega (runBytes $ hintBitPack omega h) + , testProperty "hintBitUnpack . hintBitPack == id (unaligned)" $ \(Dim k) extra -> do + h <- arbitraryHintsVector k + let omega = fromIntegral (extra + countHints h) + return $ omega < 256 ==> + Just h === hintBitUnpack omega (runUnaligned $ hintBitPack omega h) + ] + , testGroup "Zq" + [ testProperty "toZq . fromZq == id " $ \(FE a) -> + a === toZq (fromZq a) + , testProperty "fromZq . toZq == id " $ \a -> + mod a 8380417 === fromZq (toZq a) + , testCase "field order" $ zero @=? toZq 8380417 + , testProperty "addition with zero" $ \(FE a) -> + conjoin [ a === zero .+ a + , a === a .+ zero + ] + , testProperty "addition associative" $ \(FE a) (FE b) (FE c) -> + a .+ (b .+ c) === (a .+ b) .+ c + , testProperty "addition commutative" $ \(FE a) (FE b) -> + a .+ b === b .+ a + , testProperty "substraction with zero" $ \(FE a) -> + a === a .- zero + , testProperty "substraction non-associative" $ \(FE a) (FE b) (FE c) -> + a .- (b .- c) === (a .- b) .+ c + , testProperty "substraction anti-commutative" $ \(FE a) (FE b) -> + a .- b === neg (b .- a) + , testProperty "negation" $ \(FE a) -> + neg a === zero .- a + , testProperty "double negation" $ \(FE a) -> + a === neg (neg a) + , testProperty "multiplication with zero" $ \(FE a) -> + conjoin [ zero === zero .* a + , zero === a .* zero + ] + , testProperty "multiplication with one" $ \(FE a) -> + conjoin [ a === one .* a + , a === a .* one + ] + , testProperty "multiplication associative" $ \(FE a) (FE b) (FE c) -> + a .* (b .* c) === (a .* b) .* c + , testProperty "multiplication commutative" $ \(FE a) (FE b) -> + a .* b === b .* a + , testProperty "multiplication distributive" $ \(FE a) (FE b) (FE c) -> + conjoin [ (a .* b) .+ (a .* c) === a .* (b .+ c) + , (b .* a) .+ (c .* a) === (b .+ c) .* a + ] + ] + , testGroup "Mq" + [ testProperty "toMontgomery . fromMontgomery == id " $ \(ME a) -> + a === toMontgomery (fromMontgomery a) + , testProperty "fromMontgomery . toMontgomery == id" $ \(FE a) -> + a === fromMontgomery (toMontgomery a) + , testProperty "addition with zero" $ \(ME a) -> + conjoin [ a === zero .+ a + , a === a .+ zero + ] + , testProperty "addition associative" $ \(ME a) (ME b) (ME c) -> + a .+ (b .+ c) === (a .+ b) .+ c + , testProperty "addition commutative" $ \(ME a) (ME b) -> + a .+ b === b .+ a + , testProperty "substraction with zero" $ \(ME a) -> + a === a .- zero + , testProperty "substraction non-associative" $ \(ME a) (ME b) (ME c) -> + a .- (b .- c) === (a .- b) .+ c + , testProperty "substraction anti-commutative" $ \(ME a) (ME b) -> + a .- b === neg (b .- a) + , testProperty "negation" $ \(ME a) -> + neg a === zero .- a + , testProperty "double negation" $ \(ME a) -> + a === neg (neg a) + , testProperty "multiplication with zero" $ \(ME a) -> + conjoin [ zero === zero .* a + , zero === a .* zero + ] + , testProperty "multiplication with one" $ \(ME a) -> + conjoin [ a === one .* a + , a === a .* one + ] + , testProperty "multiplication associative" $ \(ME a) (ME b) (ME c) -> + a .* (b .* c) === (a .* b) .* c + , testProperty "multiplication commutative" $ \(ME a) (ME b) -> + a .* b === b .* a + , testProperty "multiplication distributive" $ \(ME a) (ME b) (ME c) -> + conjoin [ (a .* b) .+ (a .* c) === a .* (b .+ c) + , (b .* a) .+ (c .* a) === (b .+ c) .* a + ] + ] + , testGroup "Rq" + [ testProperty "fromCoeffs . toCoeffs == id " $ \(Poly a) -> + Just a === fromCoeffs (toCoeffs a) + , testProperty "addition with zero" $ \(Poly a) -> + conjoin [ a === zero .+ a + , a === a .+ zero + ] + , testProperty "addition associative" $ \(Poly a) (Poly b) (Poly c) -> + a .+ (b .+ c) === (a .+ b) .+ c + , testProperty "addition commutative" $ \(Poly a) (Poly b) -> + a .+ b === b .+ a + , testProperty "substraction with zero" $ \(Poly a) -> + a === a .- zero + , testProperty "substraction non-associative" $ \(Poly a) (Poly b) (Poly c) -> + a .- (b .- c) === (a .- b) .+ c + , testProperty "substraction anti-commutative" $ \(Poly a) (Poly b) -> + a .- b === neg (b .- a) + , testProperty "negation" $ \(Poly a) -> + neg a === zero .- a + , testProperty "double negation" $ \(Poly a) -> + a === neg (neg a) + , testCase "norm zero" $ norm (zero :: Rq Sec) @=? 0 + , testProperty "norm positiveness" $ \(Poly a) -> + (norm a == 0) === (a == zero) + , testProperty "norm even" $ \(Poly a) -> + norm (neg a) === norm a + , testProperty "norm sub-additivity" $ \(Poly a) (Poly b) -> + norm (a .+ b) <= norm a + norm b + ] + , testGroup "Tq" + [ testProperty "nttInv . ntt == id" $ \(Poly a) -> + a === nttInv (ntt a) + , testProperty "addition with zero" $ \(PolyNTT a) -> + conjoin [ a === zero .+ a + , a === a .+ zero + ] + , testProperty "addition associative" $ \(PolyNTT a) (PolyNTT b) (PolyNTT c) -> + a .+ (b .+ c) === (a .+ b) .+ c + , testProperty "addition commutative" $ \(PolyNTT a) (PolyNTT b) -> + a .+ b === b .+ a + , testProperty "substraction with zero" $ \(PolyNTT a) -> + a === a .- zero + , testProperty "substraction non-associative" $ \(PolyNTT a) (PolyNTT b) (PolyNTT c) -> + a .- (b .- c) === (a .- b) .+ c + , testProperty "substraction anti-commutative" $ \(PolyNTT a) (PolyNTT b) -> + a .- b === neg (b .- a) + , testProperty "negation" $ \(PolyNTT a) -> + neg a === zero .- a + , testProperty "double negation" $ \(PolyNTT a) -> + a === neg (neg a) + , testProperty "multiplication with zero" $ \(PolyNTT a) -> + conjoin [ zero === zero .* a + , zero === a .* zero + ] + , testProperty "multiplication with one" $ \(PolyNTT a) -> + conjoin [ a === one .* a + , a === a .* one + ] + , testProperty "multiplication associative" $ \(PolyNTT a) (PolyNTT b) (PolyNTT c) -> + a .* (b .* c) === (a .* b) .* c + , testProperty "multiplication commutative" $ \(PolyNTT a) (PolyNTT b) -> + a .* b === b .* a + , testProperty "multiplication distributive" $ \(PolyNTT a) (PolyNTT b) (PolyNTT c) -> + conjoin [ (a .* b) .+ (a .* c) === a .* (b .+ c) + , (b .* a) .+ (c .* a) === (b .+ c) .* a + ] + , testProperty "mulAdd" $ \(PolyNTT a) (PolyNTT b) (PolyNTT c) -> + a .* b .+ c === mulAdd a b c + ] + , testGroup "Vector" + [ testProperty "addition with zero" $ \(Dim n) -> do + a <- arbitraryVector n + return $ conjoin + [ a === zero .+ a + , a === a .+ zero + ] + , testProperty "addition associative" $ \(Dim n) -> do + (a, b, c) <- (,,) <$> arbitraryVector n <*> arbitraryVector n <*> arbitraryVector n + return (a .+ (b .+ c) === (a .+ b) .+ c) + , testProperty "addition commutative" $ \(Dim n) -> do + (a, b) <- (,) <$> arbitraryVector n <*> arbitraryVector n + return (a .+ b === b .+ a) + , testProperty "substraction with zero" $ \(Dim n) -> do + a <- arbitraryVector n + return (a === a .- zero) + , testProperty "substraction non-associative" $ \(Dim n) -> do + (a, b, c) <- (,,) <$> arbitraryVector n <*> arbitraryVector n <*> arbitraryVector n + return (a .- (b .- c) === (a .- b) .+ c) + , testProperty "substraction anti-commutative" $ \(Dim n) -> do + (a, b) <- (,) <$> arbitraryVector n <*> arbitraryVector n + return (a .- b === neg (b .- a)) + , testProperty "negation" $ \(Dim n) -> do + a <- arbitraryVector n + return (neg a === zero .- a) + , testProperty "double negation" $ \(Dim n) -> do + a <- arbitraryVector n + return (a === neg (neg a)) + , testProperty "dot product commutative" $ \(Dim n) -> do + (u, v) <- (,) <$> arbitraryVector n <*> arbitraryVector n + return (u `dot` v === v `dot` u) + , testProperty "dot product distributive" $ \(Dim n) -> do + (u, v, w) <- (,,) <$> arbitraryVector n <*> arbitraryVector n <*> arbitraryVector n + return $ conjoin + [ u `dot` (v .+ w) === (u `dot` v) .+ (u `dot` w) + , (u .+ v) `dot` w === (u `dot` w) .+ (v `dot` w) + ] + ] + , testGroup "Matrix" + [ testProperty "mmul distributive left" $ \(Dim n) (Dim m) -> do + (a, b, u) <- (,,) <$> arbitraryMatrix n m <*> arbitraryMatrix n m <*> arbitraryVector m + return ((a .+ b) `mmul` u === (a `mmul` u) .+ (b `mmul` u)) + , testProperty "mmul distributive right" $ \(Dim n) (Dim m) -> do + (a, u, v) <- (,,) <$> arbitraryMatrix n m <*> arbitraryVector m <*> arbitraryVector m + return (a `mmul` (u .+ v) === (a `mmul` u) .+ (a `mmul` v)) + , testProperty "mmulAdd definition" $ \(Dim n) (Dim m) -> do + (a, u, v) <- (,,) <$> arbitraryMatrix n m <*> arbitraryVector m <*> arbitraryVector n + return (mmulAdd a u v == (a `mmul` u) .+ v) + , testProperty "mmulAdd distributive left" $ \(Dim n) (Dim m) -> do + (a, b, u, v) <- (,,,) <$> arbitraryMatrix n m <*> arbitraryMatrix n m <*> arbitraryVector m <*> arbitraryVector n + return (mmulAdd (a .+ b) u v === mmulAdd a u (mmulAdd b u zero) .+ v) + , testProperty "mmulAdd distributive right" $ \(Dim n) (Dim m) -> do + (a, u, v, w) <- (,,,) <$> arbitraryMatrix n m <*> arbitraryVector m <*> arbitraryVector m <*> arbitraryVector n + return (mmulAdd a (u .+ v) w === mmulAdd a u (mmulAdd a v w)) + ] +#endif + ] + ] diff --git a/tests/Util.hs b/tests/Util.hs new file mode 100644 --- /dev/null +++ b/tests/Util.hs @@ -0,0 +1,47 @@ +-- | +-- Module : Util +-- License : BSD-3-Clause +-- Copyright : (c) 2025 Olivier Chรฉron +-- +-- Utility to read hexadecimal byte arrays with aeson +-- +{-# LANGUAGE CPP #-} +module Util + ( (.::), (.::?) + ) where + +import Data.Aeson +import Data.Aeson.Types +import Data.Char +import Data.ByteString (ByteString, cons, empty) +import Data.Text (Text, uncons) + +#if !(MIN_VERSION_aeson(2,0,0)) +type Key = Text +#endif + +(.::) :: Object -> Key -> Parser ByteString +o .:: name = (o .: name) >>= fromBase16 + +(.::?) :: Object -> Key -> Parser (Maybe ByteString) +o .::? name = (o .:? name) >>= opt fromBase16 + +opt :: Monad m => (a -> m b) -> Maybe a -> m (Maybe b) +opt f = maybe (return Nothing) (fmap Just . f) + +fromBase16 :: Text -> Parser ByteString +fromBase16 t = case uncons t of + Nothing -> return empty + Just (a, as) -> + case uncons as of + Nothing -> fail "incomplete Base16" + Just (b, bs) -> do + ia <- fromHexDigit a + ib <- fromHexDigit b + let w = fromIntegral (ia * 16 + ib) + cons w <$> fromBase16 bs + +fromHexDigit :: Char -> Parser Int +fromHexDigit c + | isHexDigit c = return (digitToInt c) + | otherwise = fail "invalid hex digit" diff --git a/tests/Vectors.hs b/tests/Vectors.hs new file mode 100644 --- /dev/null +++ b/tests/Vectors.hs @@ -0,0 +1,41 @@ +-- | +-- Module : Vectors +-- License : BSD-3-Clause +-- Copyright : (c) 2025 Olivier Chรฉron +-- +-- Common implementation of ML-DSA test vectors +-- +{-# LANGUAGE OverloadedStrings #-} +module Vectors + ( VectorFile(..), readJson + ) where + +import Data.Aeson +import qualified Data.ByteString.Lazy as L + +import qualified Codec.Compression.GZip as GZip + +data VectorFile tg = VectorFile + { vsId :: Int + , algorithm :: String + , mode :: String + , revision :: String + , isSample :: Bool + , testGroups :: [tg] + } deriving Show + +instance FromJSON tg => FromJSON (VectorFile tg) where + parseJSON = withObject "File" $ \o -> VectorFile + <$> o .: "vsId" + <*> o .: "algorithm" + <*> o .: "mode" + <*> o .: "revision" + <*> o .: "isSample" + <*> o .: "testGroups" + +readJson :: FromJSON tg => FilePath -> IO (VectorFile tg) +readJson path = do + bs <- L.readFile path + case decode (GZip.decompress bs) of + Just file -> return file + _ -> fail "could not parse" diff --git a/tests/get-vectors.sh b/tests/get-vectors.sh new file mode 100644 --- /dev/null +++ b/tests/get-vectors.sh @@ -0,0 +1,13 @@ +#!/bin/sh + +DESTDIR="`dirname "$0"`" + +REF=commit/79e78ba49d1605baaf9acdbf475304af6ae36a59 +CURL=curl + +for KEY in keyGen sigGen sigVer; do + FILENAME="$DESTDIR"/$KEY.json.gz + URL=https://codeberg.org/ocheron/hs-mldsa/raw/$REF/tests/$KEY.json.gz + + "$CURL" --silent --fail -o "$FILENAME" "$URL" || exit $? +done