diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,5 @@ +# Changelog for `mlkem` + +## 0.1.0.0 - 2025-11-02 + +* 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 2025 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,38 @@ +# ML-KEM + +_Module-Lattice-based Key-Encapsulation Mechanism_ implemented in Haskell. + +See [FIPS 203](https://csrc.nist.gov/pubs/fips/203/final). + +Example session: + +```haskell +> import Crypto.PubKey.ML_KEM +> import Data.Proxy +> let params = Proxy :: Proxy ML_KEM_768 +> (encapKey, decapKey) <- generate params +> (sharedKey, ciphertext) <- encapsulate encapKey +> let sharedKey' = decapsulate decapKey ciphertext +> sharedKey == sharedKey' +True +``` + +## 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. + +Best performance is obtained with the LLVM code generator. + +## Testing + +The test suite executes all NIST test vectors but necessary files are not +included in the package to limit its size. Instead, two 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,52 @@ +{-# LANGUAGE ExistentialQuantification #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE OverloadedStrings #-} +module Main (main) where + +import Criterion.Main + +import Crypto.Random +import Crypto.PubKey.ML_KEM + +import Data.ByteArray (Bytes) +import Data.Proxy + +data KEM = forall a . ParamSet a => KEM (Proxy a) + +kems :: [(String, KEM)] +kems = + [ ("ML-KEM-512", KEM (Proxy :: Proxy ML_KEM_512)) + , ("ML-KEM-768", KEM (Proxy :: Proxy ML_KEM_768)) + , ("ML-KEM-1024", KEM (Proxy :: Proxy ML_KEM_1024)) + ] + +doBench :: (String, KEM) -> Benchmark +doBench (name, KEM p) = bgroup name + [ bench "generate" $ perRunEnv setupGenerate (return . runGenerate) + , bench "encapsulate" $ perRunEnv setupEncap (return . runEncap) + , bench "encapsulate (batch)" $ perBatchEnv (const setupEncap) (return . runEncap) + , bench "decapsulate" $ perRunEnv setupDecap (return . runDecap) + , bench "decapsulate (batch)" $ perBatchEnv (const setupDecap) (return . runDecap) + ] + where + gen32 = getRandomBytes 32 :: IO Bytes + + setupGenerate = (,) <$> gen32 <*> gen32 + runGenerate = uncurry (generateWith p) + + setupEncap = do + (ek, _) <- generate p + m <- gen32 + return (ek, m) + runEncap (ek, m) = encapsulateWith ek m + + runDecap (dk, c) = decapsulate dk c + setupDecap = do + (ek, dk) <- generate p + (_, c) <- encapsulate ek + return (dk, c) + +main :: IO () +main = defaultMain + [ bgroup "mlkem" $ map doBench kems + ] diff --git a/mlkem.cabal b/mlkem.cabal new file mode 100644 --- /dev/null +++ b/mlkem.cabal @@ -0,0 +1,188 @@ +cabal-version: 2.2 + +-- This file has been generated from package.yaml by hpack version 0.38.0. +-- +-- see: https://github.com/sol/hpack + +name: mlkem +version: 0.1.0.0 +synopsis: Module-Lattice-based Key-Encapsulation Mechanism +description: Module-Lattice-based Key-Encapsulation Mechanism (ML-KEM) implemented in + Haskell. +category: Crypto +homepage: https://codeberg.org/ocheron/hs-mlkem#readme +bug-reports: https://codeberg.org/ocheron/hs-mlkem/issues +author: Olivier Chéron +maintainer: olivier.cheron@gmail.com +copyright: 2025 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-mlkem + +flag use_crypton + description: Use crypton instead of cryptonite + manual: True + default: False + +library + exposed-modules: + Crypto.PubKey.ML_KEM + other-modules: + Auxiliary + Block + BlockN + Builder + ByteArrayST + Crypto + Internal + K_PKE + Marking + Math + Matrix + ScrubbedBlock + SecureBlock + SecureBytes + Vector + Paths_mlkem + autogen-modules: + Paths_mlkem + 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 -O2 + build-depends: + base >=4.7 && <5 + , basement >=0.0.8 + , deepseq + , memory + default-language: Haskell2010 + if flag(use_crypton) + build-depends: + crypton + else + build-depends: + cryptonite >=0.26 + +test-suite mlkem-test + type: exitcode-stdio-1.0 + main-is: Tests.hs + other-modules: + EncapDecap + KeyGen + Util + Vectors + Paths_mlkem + autogen-modules: + Paths_mlkem + 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 -threaded -rtsopts -with-rtsopts=-N + build-depends: + aeson + , base >=4.7 && <5 + , basement >=0.0.8 + , bytestring + , deepseq + , directory + , memory + , mlkem + , process + , tasty + , tasty-hunit + , tasty-quickcheck + , text + , zlib + default-language: Haskell2010 + if flag(use_crypton) + build-depends: + crypton + else + build-depends: + cryptonite >=0.26 + +test-suite mlkem-test-full + type: exitcode-stdio-1.0 + main-is: Tests.hs + other-modules: + Auxiliary + Block + BlockN + Builder + ByteArrayST + Crypto + Crypto.PubKey.ML_KEM + Internal + K_PKE + Marking + Math + Matrix + ScrubbedBlock + SecureBlock + SecureBytes + Vector + EncapDecap + KeyGen + Util + Vectors + Paths_mlkem + autogen-modules: + Paths_mlkem + 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 -fno-ignore-asserts -threaded -rtsopts -with-rtsopts=-N + cpp-options: -DML_KEM_TESTING + build-depends: + aeson + , base >=4.7 && <5 + , basement >=0.0.8 + , bytestring + , deepseq + , directory + , memory + , process + , tasty + , tasty-hunit + , tasty-quickcheck + , text + , zlib + default-language: Haskell2010 + if flag(use_crypton) + build-depends: + crypton + else + build-depends: + cryptonite >=0.26 + +benchmark mlkem-bench + type: exitcode-stdio-1.0 + main-is: Bench.hs + other-modules: + Paths_mlkem + autogen-modules: + Paths_mlkem + 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 -threaded -rtsopts -with-rtsopts=-N + build-depends: + base >=4.7 && <5 + , basement >=0.0.8 + , criterion + , deepseq + , memory + , mlkem + default-language: Haskell2010 + if flag(use_crypton) + build-depends: + crypton + else + build-depends: + cryptonite >=0.26 diff --git a/src/Auxiliary.hs b/src/Auxiliary.hs new file mode 100644 --- /dev/null +++ b/src/Auxiliary.hs @@ -0,0 +1,585 @@ +-- | +-- Module : Auxiliary +-- License : BSD-3-Clause +-- Copyright : (c) 2025 Olivier Chéron +-- +-- ML-KEM auxiliary functions +-- +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +module Auxiliary + ( Zq, Rq, Tq, (..+), (..-) + , ntt, nttInv, rcompress, rdecompress + , byteEncode, byteDecode, byteEncode12, byteDecode12 + , byteEncode1, byteDecode1, sampleNTT, samplePolyCBD +#ifdef ML_KEM_TESTING + , compress, decompress + , bitRev7, fromZq, toZq, fromCoeffs, toCoeffs +#endif + ) where + +import Basement.NormalForm +import Basement.PrimType +import Basement.Types.OffsetSize + +import Crypto.Hash.Algorithms + +import Data.ByteArray (ByteArrayAccess, Bytes, View) +import qualified Data.ByteArray as B +import qualified Data.Memory.Endian as B + +import Control.DeepSeq (NFData(..)) +import Control.Monad +import Control.Monad.ST + +import Data.Bits +import Data.Proxy +import Data.Word + +import GHC.TypeNats + +import Foreign.Ptr (Ptr, plusPtr) +import Foreign.Storable (pokeByteOff) + +import Block (blockIndex) +import BlockN (BlockN, MutableBlockN) +import Builder (Builder) +import Crypto (BlockDigest) +import Marking (Classified, SecurityMarking(..), Leak(..)) +import SecureBlock (SecureBlock) +import SecureBytes (SecureBytes) +import qualified BlockN +import qualified Builder +import qualified ByteArrayST as ST +import qualified Crypto +import Math + +type N = 256 + +n :: Int +n = 256 + +q :: Integer +q = 3329 + +q16 :: Word16 +q16 = fromInteger q + +q32 :: Word32 +q32 = fromInteger q + +q64 :: Word64 +q64 = fromInteger q + +bitRev7 :: Word8 -> Word8 +bitRev7 b = + (b `unsafeShiftR` 6 .&. 1) .|. + (b `unsafeShiftR` 5 .&. 1) `unsafeShiftL` 1 .|. + (b `unsafeShiftR` 4 .&. 1) `unsafeShiftL` 2 .|. + (b `unsafeShiftR` 3 .&. 1) `unsafeShiftL` 3 .|. + (b `unsafeShiftR` 2 .&. 1) `unsafeShiftL` 4 .|. + (b `unsafeShiftR` 1 .&. 1) `unsafeShiftL` 5 .|. + (b .&. 1) `unsafeShiftL` 6 + +-- Reduction 𝑥 mod 𝑞 for 0 ≤ 𝑥 < 2𝑞 +reduceSimple :: Word16 -> Word16 +reduceSimple x = (mask .&. x) .|. (complement mask .&. subtracted) + where + subtracted = x - q16 + mask = negate (subtracted `unsafeShiftR` 15) +{-# INLINE reduceSimple #-} + +-- Reduction 𝑥 mod 𝑞 for 0 ≤ 𝑥 < 2𝑞² + 𝑞 +reduce :: Word32 -> Word16 +reduce x = reduceSimple (fromIntegral remainder) + where + p = fromIntegral x * ((1 `unsafeShiftL` 24) `div` q64) + quotient = fromIntegral (p `unsafeShiftR` 24) + remainder = x - quotient * q32 +{-# INLINE reduce #-} + +newtype Zq = Zq Word16 +#ifdef ML_KEM_TESTING + deriving (Eq, Show) +#else + deriving Eq +#endif + +instance PrimType Zq where + type PrimSize Zq = 2 + primSizeInBytes _ = primSizeInBytes (Proxy :: Proxy Word16) + {-# INLINE primSizeInBytes #-} + primShiftToBytes _ = primShiftToBytes (Proxy :: Proxy Word16) + {-# INLINE primShiftToBytes #-} + primBaUIndex ba (Offset i) = Zq (primBaUIndex ba (Offset i)) + {-# INLINE primBaUIndex #-} + primMbaURead mba (Offset i) = Zq <$> primMbaURead mba (Offset i) + {-# INLINE primMbaURead #-} + primMbaUWrite mba (Offset i) (Zq a) = primMbaUWrite mba (Offset i) a + {-# INLINE primMbaUWrite #-} + primAddrIndex addr (Offset i) = Zq (primAddrIndex addr (Offset i)) + {-# INLINE primAddrIndex #-} + primAddrRead addr (Offset i) = Zq <$> primAddrRead addr (Offset i) + {-# INLINE primAddrRead #-} + primAddrWrite addr (Offset i) (Zq a) = primAddrWrite addr (Offset i) a + {-# INLINE primAddrWrite #-} + +instance Add Zq where + zero = Zq 0 + Zq a .+ Zq b = Zq $ reduceSimple (a + b) + Zq a .- Zq b = Zq $ reduceSimple (a + q16 - b) + neg (Zq a) = Zq $ reduceSimple (q16 - a) + +instance Mul Zq where + one = Zq 1 + Zq a .* Zq b = Zq $ reduce (fromIntegral a * fromIntegral b) + +#ifdef ML_KEM_TESTING +instance MulAdd Zq where + mulAdd (Zq a) (Zq b) (Zq c) = Zq $ reduce $ + fromIntegral a * fromIntegral b + fromIntegral c + +instance BiMul Zq Zq where + (..*) = (.*) + +instance BiMulAdd Zq Zq where + biMulAdd = mulAdd + +fromZq :: Zq -> Word16 +fromZq (Zq a) = a +#endif + +toZq :: Word16 -> Zq +toZq = Zq . reduce . fromIntegral + +newtype Rq marking = Rq (BlockN marking N Zq) +#ifdef ML_KEM_TESTING + deriving (Eq, Show) +#endif + +instance Classified marking => Add (Rq marking) where + zero = Rq zero + Rq a .+ Rq b = Rq (a .+ b) + Rq a .- Rq b = Rq (a .- b) + neg (Rq a) = Rq (neg a) + {-# SPECIALIZE instance Add (Rq Sec) #-} + {-# SPECIALIZE instance Add (Rq Pub) #-} + +infixl 6 ..+, ..- + +-- Transformation called only at expected location in the LWE problem, after +-- adding noise to secret information. +(..+) :: Rq Sec -> Rq Sec -> Rq Pub +Rq a ..+ Rq b = Rq $ BlockN.zipWith (.+) a b + +(..-) :: Rq Pub -> Rq Sec -> Rq Sec +Rq a ..- Rq b = Rq $ BlockN.zipWith (.-) a b + +instance Leak Rq + +#ifdef ML_KEM_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 Zq) +#ifdef ML_KEM_TESTING + deriving (Eq, Show) +#endif + +instance Classified marking => Add (Tq marking) where + zero = Tq zero + Tq a .+ Tq b = Tq (a .+ b) + Tq a .- Tq b = Tq (a .- b) + neg (Tq a) = Tq (neg a) + {-# SPECIALIZE instance Add (Tq Sec) #-} + {-# SPECIALIZE instance Add (Tq Pub) #-} + +instance Leak Tq + +instance Classified marking => NFData (Tq marking) where + rnf (Tq a) = toNormalForm a + +instance BiMul (Tq Pub) (Tq Sec) where + (..*) = multiplyNTTs + +instance BiMulAdd (Tq Pub) (Tq Sec) where + biMulAdd = multiplyNTTsAdd + +#ifdef ML_KEM_TESTING +instance Mul (Tq Sec) where + one = Tq $ BlockN.create $ \(Offset i) -> if even i then one else zero + (.*) = (..*) . leak + +instance MulAdd (Tq Sec) where + mulAdd = biMulAdd . leak +#endif + +instance Crypto.ConstEqW (Tq Sec) where + constEqW (Tq a) (Tq b) = Crypto.constEqW + (BlockN.unsafeCast a :: SecureBlock Sec Word) + (BlockN.unsafeCast b :: SecureBlock Sec Word) + +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) + +-- Computes the NTT representation of the given polynomial +ntt :: Classified marking => Rq marking -> Tq marking +ntt (Rq !a) = runST $ do + b <- BlockN.thaw a + outer b 1 128 + Tq <$> BlockN.unsafeFreeze b + where + outer !b !i len = when (len >= 2) $ inner b i len 0 + + inner !b !i !len start + | start < 256 = do + let zeta = BlockN.index zetaPowBitRev i -- 17 ^ bitRev7 i + loop b zeta (start + len) len start + inner b (i + 1) len (start + offsetShiftL 1 len) + | otherwise = outer b i (offsetShiftR 1 len) + + loop !b !zeta end len j = + when (j < end) $ do + t <- (zeta .*) <$> BlockN.read b (j + len) + x <- BlockN.read b j + BlockN.write b (j + len) (x .- t) + BlockN.write b j (x .+ t) + loop b zeta end len (j + 1) +{-# SPECIALIZE ntt :: Rq Sec -> Tq Sec #-} +{-# SPECIALIZE ntt :: Rq Pub -> Tq Pub #-} + +-- Computes the polynomial that corresponds to the given NTT representation +nttInv :: Tq Sec -> Rq Sec +nttInv (Tq !a) = runST $ do + b <- BlockN.thaw a + outer b 127 2 + BlockN.iterModify (\x -> x .* Zq 3303) b + Rq <$> BlockN.unsafeFreeze b + where + outer !b !i len = when (len <= 128) $ inner b i len 0 + + inner !b !i !len start + | start < 256 = do + let zeta = BlockN.index zetaPowBitRev i -- 17 ^ bitRev7 i + loop b zeta (start + len) len start + inner b (i - 1) len (start + offsetShiftL 1 len) + | otherwise = outer b i (offsetShiftL 1 len) + + loop !b !zeta 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) (zeta .* (x .- t)) + loop b zeta end len (j + 1) + +-- Computes the product of two NTT representations +multiplyNTTs :: Tq Pub -> Tq Sec -> Tq Sec +multiplyNTTs (Tq !f) (Tq !g) = runST $ do + b <- BlockN.new (Proxy :: Proxy Sec) + loop b 0 + Tq <$> BlockN.unsafeFreeze b + where + loop :: MutableBlockN Sec N Zq s -> Offset Zq -> ST s () + loop !b i = when (i < 128) $ do + let ii = offsetShiftL 1 i + a0 = BlockN.index f ii + a1 = BlockN.index f (ii + 1) + b0 = BlockN.index g ii + b1 = BlockN.index g (ii + 1) + (c0, c1) = baseCaseMultiply a0 a1 b0 b1 (BlockN.index gamma i) + BlockN.write b ii c0 + BlockN.write b (ii + 1) c1 + loop b (i + 1) + +-- Computes the product of two degree-one polynomials with respect to a quadratic modulus +baseCaseMultiply :: Zq -> Zq -> Zq -> Zq -> Zq -> (Zq, Zq) +baseCaseMultiply (Zq a0) (Zq a1) (Zq b0) (Zq b1) (Zq g) = (Zq c0, Zq c1) + where + x `mul` y = fromIntegral x * fromIntegral y + b1g = reduce (b1 `mul` g) + !c0 = reduce (a0 `mul` b0 + a1 `mul` b1g) + !c1 = reduce (a0 `mul` b1 + a1 `mul` b0) + +-- Multiply then add a third term +multiplyNTTsAdd :: Tq Pub -> Tq Sec -> Tq Sec -> Tq Sec +multiplyNTTsAdd (Tq !f) (Tq !g) (Tq !h) = runST $ do + b <- BlockN.new (Proxy :: Proxy Sec) + loop b 0 + Tq <$> BlockN.unsafeFreeze b + where + loop :: MutableBlockN Sec N Zq s -> Offset Zq -> ST s () + loop !b i = when (i < 128) $ do + let ii = offsetShiftL 1 i + a0 = BlockN.index f ii + a1 = BlockN.index f (ii + 1) + b0 = BlockN.index g ii + b1 = BlockN.index g (ii + 1) + c0 = BlockN.index h ii + c1 = BlockN.index h (ii + 1) + (d0, d1) = baseCaseMultiplyAdd a0 a1 b0 b1 c0 c1 (BlockN.index gamma i) + BlockN.write b ii d0 + BlockN.write b (ii + 1) d1 + loop b (i + 1) + +-- baseCaseMultiply then add a third term +baseCaseMultiplyAdd :: Zq -> Zq -> Zq -> Zq -> Zq -> Zq -> Zq -> (Zq, Zq) +baseCaseMultiplyAdd (Zq a0) (Zq a1) (Zq b0) (Zq b1) (Zq c0) (Zq c1) (Zq g) = (Zq d0, Zq d1) + where + x `mul` y = fromIntegral x * fromIntegral y + b1g = reduce (b1 `mul` g) + !d0 = reduce (fromIntegral c0 + a0 `mul` b0 + a1 `mul` b1g) + !d1 = reduce (fromIntegral c1 + a0 `mul` b1 + a1 `mul` b0) + +-- Values of 17 ^ BitRev7(𝑖) mod 𝑞 for 𝑖 ∈ {0, … , 127} +zetaPowBitRev :: BlockN Pub 128 Zq +zetaPowBitRev = runST $ do + out <- BlockN.new (Proxy :: Proxy Pub) + foldM_ (loop out) one offsets + BlockN.unsafeFreeze out + where + offsets = Prelude.map (fromIntegral . bitRev7) [0 .. 127] + loop b acc i = BlockN.write b i acc >> return (Zq 17 .* acc) + +-- Values of 17 ^ 2.BitRev7(𝑖)+1 mod 𝑞 for 𝑖 ∈ {0, … , 127} +gamma :: BlockN Pub 128 Zq +gamma = BlockN.map (\z -> z .* z .* Zq 17) zetaPowBitRev + +-- Compress a field element with 𝑑 < 12 +compress :: Int -> Zq -> Word16 +compress d (Zq x) = fromIntegral $ + ((fromIntegral x `unsafeShiftL` d + qHalf) * factor) `unsafeShiftR` 34 + where + qHalf = (q64 + 1) `unsafeShiftR` 1 + factor = (1 `unsafeShiftL` 34) `div` q64 +{-# INLINE compress #-} + +-- Decompress a field element with 𝑑 < 12 +decompress :: Int -> Word16 -> Zq +decompress d y = Zq $ fromIntegral (x2d `unsafeShiftR` d) + where x2d = fromIntegral y * q32 + (1 `unsafeShiftL` (d - 1)) +{-# INLINE decompress #-} + +-- Compress a polynomial with 𝑑 < 12 +rcompress :: Classified marking => Int -> Rq marking -> BlockN marking N Word16 +rcompress !d (Rq a) = BlockN.map (compress d) a +{-# SPECIALIZE rcompress :: Int -> Rq Sec -> BlockN Sec N Word16 #-} +{-# SPECIALIZE rcompress :: Int -> Rq Pub -> BlockN Pub N Word16 #-} + +-- Decompress a polynomial with 𝑑 < 12 +rdecompress :: Classified marking => Int -> BlockN marking N Word16 -> Rq marking +rdecompress !d = Rq . BlockN.map (decompress d) +{-# SPECIALIZE rdecompress :: Int -> BlockN Sec N Word16 -> Rq Sec #-} +{-# SPECIALIZE rdecompress :: Int -> BlockN Pub N Word16 -> Rq Pub #-} + +-- Generates a pseudorandom element of T𝑞 from a seed and two indices +sampleNTT :: SecureBytes Pub -> Word8 -> Word8 -> Tq Pub +sampleNTT seed !x !y = runST $ do + b <- BlockN.new (Proxy :: Proxy Pub) + runXof b (280 * 3) 0 0 + Tq <$> BlockN.unsafeFreeze b + where + runXof !b !xofLen !pos !j = case someNatVal (8 * fromIntegral 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 c0 = fromIntegral $ blockIndex bytes pos + c1 = fromIntegral $ blockIndex bytes (pos + 1) + c2 = fromIntegral $ blockIndex bytes (pos + 2) + d1 = c0 + (c1 .&. 0xF) `unsafeShiftL` 8 + d2 = (c1 `unsafeShiftR` 4) + (c2 `unsafeShiftL` 4) + j2 <- poke b j d1 + when (j2 < 256) $ poke b j2 d2 >>= loop b xofLen bytes (pos + 3) + + poke b j d + | d < q16 = BlockN.write b j (Zq d) >> 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 seed d + pokeByteOff d len x + pokeByteOff d (len + 1) y + len = B.length seed + +wordBits :: Int +wordBits = finiteBitSize (0 :: Word) + +wordBytes :: Int +wordBytes = div wordBits 8 + +type WordLE = B.LE Word + +peekWord :: Ptr WordLE -> ST s Word +peekWord p = B.unLE <$> ST.peek p + +peekWordPos :: Ptr WordLE -> BitPos -> ST s Word +peekWordPos a bp = B.unLE <$> ST.peekElemOff a (wordOff bp) + +pokeWordPos :: Ptr WordLE -> BitPos -> Word -> ST s () +pokeWordPos a bp = ST.pokeElemOff a (wordOff bp) . B.LE + +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 -> Word +getMask howMany + | howMany >= wordBits = maxBound + | otherwise = (1 `unsafeShiftL` howMany) - 1 + +-- Takes a seed as input and outputs a pseudorandom sample from the +-- distribution D_eta +samplePolyCBD :: Word -> SecureBytes Sec -> Rq Sec +samplePolyCBD !eta !input = runST $ ST.withByteArray input $ \p -> do + f <- BlockN.new (Proxy :: Proxy Sec) + loop p f 0 zeroPos + Rq <$> BlockN.unsafeFreeze f + where + loop :: Ptr WordLE -> MutableBlockN Sec N Zq s -> Offset Zq -> BitPos -> ST s () + loop !p !f !i !bp = when (i < Offset n) $ do + (xs, bp') <- getBits p bp 0 (fromIntegral eta) + (ys, bp'') <- getBits p bp' 0 (fromIntegral eta) + BlockN.write f i (Zq xs .- Zq ys) + loop p f (i + 1) bp'' + + getBits :: Ptr WordLE -> BitPos -> Word16 -> Int -> ST s (Word16, BitPos) + getBits !p !bp !acc !j + | j == 0 = return (acc, bp) + | otherwise = do + x <- (`unsafeShiftR` bitPos bp) <$> peekWordPos p bp + let (howMany, bp') = nextPos j bp + bits = x .&. getMask howMany + getBits p bp' (acc + fromIntegral (popCount bits)) (j - howMany) + +-- Encodes an array of 𝑑-bit integers into a byte array for 1 ≤ 𝑑 ≤ 12 +byteEncode :: Classified marking => Int -> BlockN marking N Word16 -> Builder marking +byteEncode !d !f = Builder.create (32 * d) $ \b -> + outer b zeroPos 0 0 + where + outer :: Ptr WordLE -> BitPos -> Word -> Int -> ST s () + outer !b !bp !o pos = when (pos < n) $ + inner b pos bp o (BlockN.index f (Offset pos)) d + + inner :: Ptr WordLE -> Int -> BitPos -> Word -> Word16 -> Int -> ST s () + inner !b !pos !bp !o !a j + | j == 0 = outer b bp o (pos + 1) + | bitPos bp + howMany < wordBits = inner b pos bp' o' a' j' + | otherwise = pokeWordPos b bp o' >> inner b pos bp' 0 a' j' + where + (howMany, bp') = nextPos j bp + x = fromIntegral a .&. getMask howMany + o' = o .|. (x `unsafeShiftL` bitPos bp) + a' = a `unsafeShiftR` howMany + j' = j - howMany +{-# SPECIALIZE byteEncode :: Int -> BlockN Sec N Word16 -> Builder Sec #-} +{-# SPECIALIZE byteEncode :: Int -> BlockN Pub N Word16 -> Builder Pub #-} + +-- Optimization of byteEncode when 𝑑=1 +byteEncode1 :: BlockN Sec N Word16 -> Builder Sec +byteEncode1 !f = Builder.create 32 $ \b -> + loop b 0 0 + where + loop :: Ptr WordLE -> Word -> Int -> ST s () + loop !b !o pos + | pos == n = return () + | bitPos bp + 1 < wordBits = loop b o' (pos + 1) + | otherwise = pokeWordPos b bp o' >> loop b 0 (pos + 1) + where + bp = BitPos pos + x = fromIntegral (a .&. 1) + o' = o .|. (x `unsafeShiftL` bitPos bp) + a = BlockN.index f (Offset pos) + +-- byteEncode with 𝑑=12 after conversion from the field +byteEncode12 :: Classified marking => Tq marking -> Builder marking +byteEncode12 (Tq f) = byteEncode 12 $ BlockN.map (\(Zq x) -> x) f + +-- Decodes a byte array into an array of 𝑑-bit integers for 1 ≤ 𝑑 ≤ 12 +byteDecode :: forall marking ba. Classified marking => ByteArrayAccess ba => Int -> ba -> BlockN marking N Word16 +byteDecode !d !b = runST $ + ST.withByteArray b $ \p -> do + f <- BlockN.new (Proxy :: Proxy marking) + outer f p zeroPos 0 + BlockN.unsafeFreeze f + where + outer :: MutableBlockN marking N Word16 s -> Ptr WordLE -> BitPos -> Offset Word16 -> ST s () + outer !f !p !bp i = when (i < Offset n) $ inner f p i bp 0 0 + + inner :: MutableBlockN marking N Word16 s -> Ptr WordLE -> Offset Word16 -> BitPos -> Word16 -> Int -> ST s () + inner !f !p !i !bp !v j + | j == d = BlockN.write f i v >> outer f 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 f p i bp' v' j' + + get :: Ptr WordLE -> BitPos -> Int -> ST s Word + get p bp howMany = do + x <- (`unsafeShiftR` bitPos bp) <$> peekWordPos p bp + return (x .&. getMask howMany) +{-# SPECIALIZE byteDecode :: forall ba. ByteArrayAccess ba => Int -> ba -> BlockN Sec N Word16 #-} +{-# SPECIALIZE byteDecode :: forall ba. ByteArrayAccess ba => Int -> ba -> BlockN Pub N Word16 #-} +{-# SPECIALIZE byteDecode :: Int -> View Bytes -> BlockN Sec N Word16 #-} +{-# SPECIALIZE byteDecode :: Int -> View Bytes -> BlockN Pub N Word16 #-} + +-- Optimization of byteDecode when 𝑑=1 +byteDecode1 :: ByteArrayAccess ba => ba -> BlockN Sec N Word16 +byteDecode1 !b = runST $ + ST.withByteArray b $ \p -> do + f <- BlockN.new (Proxy :: Proxy Sec) + outer f p 0 + BlockN.unsafeFreeze f + where + outer :: MutableBlockN Sec N Word16 s -> Ptr WordLE -> Int -> ST s () + outer !f !p i = when (i < n) $ do + x <- peekWord p + inner f (p `plusPtr` wordBytes) x i 0 + + inner :: MutableBlockN Sec N Word16 s -> Ptr WordLE -> Word -> Int -> Int -> ST s () + inner !f !p !acc !i j + | j == wordBits = outer f p i + | otherwise = do + let v = fromIntegral (acc .&. 1) + BlockN.write f (Offset i) v + inner f p (acc `unsafeShiftR` 1) (i + 1) (j + 1) + +-- byteDecode with 𝑑=12 and conversion to the field +byteDecode12 :: Classified marking => ByteArrayAccess ba => ba -> Tq marking +byteDecode12 = Tq . BlockN.map toZq . byteDecode 12 diff --git a/src/Block.hs b/src/Block.hs new file mode 100644 --- /dev/null +++ b/src/Block.hs @@ -0,0 +1,93 @@ +-- | +-- Module : Block +-- License : BSD-3-Clause +-- Copyright : (c) 2025 Olivier Chéron +-- +-- An array of primitive (unlifted) elements. This module currently exposes +-- the implementation from basement and fixes a lack of inlining in the +-- @create@ function. +-- +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +module Block + ( Block, MutableBlock, blockIndex, blockRead, blockWrite + , create, foldZipWith, iterModify, Block.length, Block.map + , Block.new, Block.newPinned, Block.thaw, Block.unsafeCast + , Block.unsafeFreeze, Block.unsafeThaw, Block.withMutablePtr +#ifdef ML_KEM_TESTING + , Block.replicate +#endif + ) where + +import Basement.Block (Block) +import Basement.Block.Mutable (MutableBlock) +import qualified Basement.Block as Block hiding (create, map) +import qualified Basement.Block.Mutable as Block +import Basement.Monad +import Basement.PrimType +import Basement.Types.OffsetSize + +import Control.Exception (assert) +import Control.Monad.ST + +import Data.Proxy + +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_KEM_TESTING +blockIndex = Block.index +blockRead = Block.read +blockWrite = Block.write +#else +blockIndex = Block.unsafeIndex +blockRead = Block.unsafeRead +blockWrite = Block.unsafeWrite +#endif + +create :: PrimType ty + => CountOf ty + -> (Offset ty -> ty) + -> Block ty +create n initializer = runST $ do + mb <- Block.new n + loop mb 0 + Block.unsafeFreeze mb + where + loop !mb i + | i .==# n = pure () + | otherwise = Block.unsafeWrite mb i (initializer i) >> loop mb (i + 1) +{-# INLINE create #-} + +map :: (PrimType a, PrimType b) => (a -> b) -> Block a -> Block b +map f a = Block.create lenB (f . Block.unsafeIndex a . offsetCast Proxy) + where !lenB = sizeCast (Proxy :: Proxy (a -> b)) (Block.length a) +{-# INLINE map #-} + +iterModify :: (PrimType ty, PrimMonad prim) + => (ty -> ty) + -> MutableBlock ty (PrimState prim) + -> prim () +iterModify f ma = loop 0 + where + !sz = Block.mutableLength ma + loop i + | i .==# sz = pure () + | otherwise = Block.unsafeRead ma i >>= \x -> Block.unsafeWrite ma i (f x) >> loop (i+1) +{-# INLINE iterModify #-} + +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 + CountOf sa = Block.length a + CountOf sb = Block.length b + + loop !acc i + | i == sa = acc + | otherwise = do + let va = Block.unsafeIndex a (Offset i) + let vb = Block.unsafeIndex b (Offset i) + loop (f acc va vb) (i + 1) +{-# INLINE foldZipWith #-} diff --git a/src/BlockN.hs b/src/BlockN.hs new file mode 100644 --- /dev/null +++ b/src/BlockN.hs @@ -0,0 +1,123 @@ +-- | +-- 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 KindSignatures #-} +{-# LANGUAGE ScopedTypeVariables #-} +module BlockN + ( BlockN, MutableBlockN, create, index, iterModify, BlockN.map + , BlockN.new, BlockN.read, BlockN.thaw, BlockN.unsafeCast + , BlockN.unsafeFreeze, BlockN.write, BlockN.zipWith +#ifdef ML_KEM_TESTING + , BlockN.fromList, BlockN.replicate, BlockN.toList +#endif + ) where + +import Basement.Monad +import Basement.Nat +import Basement.NormalForm +import Basement.PrimType +import Basement.Types.OffsetSize + +import Data.Proxy + +import Block (MutableBlock, blockRead, blockWrite) +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_KEM_TESTING +instance (Classified marking, 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 Classified marking => NormalForm (BlockN marking n a) where + toNormalForm = SecureBlock.toNormalForm . unBlockN + +instance (Classified marking, KnownNat n, PrimType a, Add a) => Add (BlockN marking n a) where + zero = create (const zero) + {-# INLINE zero #-} + (.+) = BlockN.zipWith (.+) + {-# INLINE (.+) #-} + (.-) = BlockN.zipWith (.-) + {-# INLINE (.-) #-} + neg = BlockN.map neg + {-# INLINE neg #-} + +newtype MutableBlockN (marking :: SecurityMarking) (n :: Nat) a m = MutableBlockN { unMutableBlockN :: MutableBlock a m } + +index :: (Classified marking, PrimType a) => BlockN marking n a -> Offset a -> a +index = SecureBlock.index . unBlockN + +#ifdef ML_KEM_TESTING +replicate :: forall marking n a. (Classified marking, KnownNat n, PrimType a) => a -> BlockN marking n a +replicate = BlockN . SecureBlock.replicate sz + where !sz = fromIntegral $ natVal (Proxy :: Proxy n) + +fromList :: forall marking n a. (Classified marking, KnownNat n, PrimType a) => [a] -> Maybe (BlockN marking n a) +fromList elems + | SecureBlock.length a == CountOf sz = Just (BlockN a) + | otherwise = Nothing + where + a = SecureBlock.fromList elems + !sz = fromIntegral $ natVal (Proxy :: Proxy n) + +toList :: (Classified marking, PrimType a) => BlockN marking n a -> [a] +toList = SecureBlock.toList . unBlockN +#endif + +create :: forall marking n ty. (Classified marking, KnownNat n, PrimType ty) + => (Offset ty -> ty) + -> BlockN marking n ty +create initializer = BlockN $ SecureBlock.create (CountOf sz) initializer + where !sz = fromIntegral $ natVal (Proxy :: Proxy n) +{-# INLINE create #-} + +map :: (Classified marking, PrimType a, PrimType b) => (a -> b) -> BlockN marking n a -> BlockN marking n b +map f = BlockN . SecureBlock.map f . unBlockN +{-# INLINE map #-} + +iterModify :: (PrimType ty, PrimMonad prim) + => (ty -> ty) + -> MutableBlockN marking n ty (PrimState prim) + -> prim () +iterModify f = SecureBlock.iterModify f . unMutableBlockN +{-# INLINE iterModify #-} + +zipWith :: (Classified ma, Classified mb, 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) = + create $ \(Offset i) -> + f (SecureBlock.index a (Offset i)) (SecureBlock.index b (Offset i)) +{-# INLINE zipWith #-} + +unsafeCast :: (Classified marking, PrimType b) => 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) + +thaw :: (Classified marking, PrimMonad prim, PrimType a) => 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 diff --git a/src/Builder.hs b/src/Builder.hs new file mode 100644 --- /dev/null +++ b/src/Builder.hs @@ -0,0 +1,83 @@ +-- | +-- 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, run, runRelaxed + , runToBlock, unsafeCreate + ) where + +import Basement.Types.OffsetSize + +import Data.ByteArray (ByteArray) + +import Control.Monad.ST +import Control.Monad.ST.Unsafe + +import Data.Semigroup +import Data.Word + +import Foreign.Ptr (Ptr, castPtr, plusPtr) + +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 a s. Ptr a -> 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 n (`loop` builders) + where + n = getSum $ Prelude.mconcat $ map (Sum . builderLength) builders + loop !_ [] = return () + loop !p (b : bs) = + copyBuilderToPtr b p >> loop (p `plusPtr` builderLength b) bs + +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 () + +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) + Block.withMutablePtr mb (copyBuilderToPtr b) + Block.unsafeFreeze mb + +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,42 @@ +-- | +-- 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, pokeByteOff + ) where + +import Data.ByteArray (ByteArray, ByteArrayAccess) +import qualified Data.ByteArray as B + +import Control.Monad.ST +import Control.Monad.ST.Unsafe + +import Foreign.Ptr (Ptr) +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 -> (forall s. Ptr p -> ST s a) -> ST t a +withByteArray b f = unsafeIOToST $ B.withByteArray b (stToIO . 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 + +pokeByteOff :: Storable a => Ptr a -> Int -> a -> ST s () +pokeByteOff a off = unsafeIOToST . S.pokeByteOff a off 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) 2025 Olivier Chéron +-- +-- Crypto-related utilities like the ML-KEM hash and PRF functions, or more +-- general concerns like constant-time equality and selection. +-- +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE ScopedTypeVariables #-} +module Crypto + ( ConstEqW(..), BoolW, andW, toBool, constSelectBytes, snoc, append + , prf, h, j, g, BlockDigest, unBlockDigest, hashToBlock + ) where + +import Crypto.Hash (Context) +import Crypto.Hash.Algorithms +import Crypto.Hash.IO + +import Control.Exception (assert) +import Control.Monad +import Control.Monad.ST + +import Data.ByteArray (ByteArray, ByteArrayAccess, Bytes, ScrubbedBytes) +import qualified Data.ByteArray as B + +import Data.Bits +import Data.Word + +import GHC.TypeNats + +import Foreign.Ptr (Ptr, castPtr, plusPtr) +import Foreign.Storable (pokeByteOff) + +import Block (Block) +import Builder (Builder) +import ScrubbedBlock (ScrubbedBlock) +import Vector (Vector) +import qualified Block +import qualified Builder +import qualified ByteArrayST as ST +import qualified ScrubbedBlock +import qualified Vector + +newtype BoolW = BoolW Word + +#ifdef ML_KEM_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 + +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 + +eqW :: Word -> Word -> BoolW +eqW a b = isZeroW (a `xor` b) + where + isZeroW x = BoolW $ msbW (complement x .&. (x - 1)) + msbW x = negate (x `unsafeShiftR` (bitsW - 1)) + +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 = + assertMultW (B.length a) $ + assertMultW (B.length b) $ + 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 #-} + +zipWith :: (Word -> Word -> Word) -> ScrubbedBytes -> ScrubbedBytes -> ScrubbedBytes +zipWith f a b = assert (sa == sb) $ assertMultW sa $ assertMultW sb $ + ST.unsafeCreate sa $ \out -> + ST.withByteArray a $ \pa -> ST.withByteArray b $ \pb -> + loop out pa pb 0 + where + !sa = B.length a + !sb = B.length b + + loop :: Ptr Word -> Ptr Word -> Ptr Word -> Int -> ST s () + loop !out !pa !pb i = when (i < sa) $ do + va <- ST.peek pa + vb <- ST.peek pb + ST.pokeByteOff out i $ f va vb + loop out (pa `plusPtr` bytesW) (pb `plusPtr` bytesW) (i + bytesW) +{-# INLINE zipWith #-} + +constSelectBytes :: BoolW -> ScrubbedBytes -> ScrubbedBytes -> ScrubbedBytes +constSelectBytes (BoolW !mask) = Crypto.zipWith f + where f yes no = (mask .&. yes) .|. (complement mask .&. no) + +-- This version of snoc accepts a more general input and uses internally a call +-- to copyByteArrayToPtr, so it does not need a trampoline when the input is +-- backed by Block Word8 +snoc :: ByteArrayAccess a => a -> Word8 -> ScrubbedBytes +snoc a b = + B.allocAndFreeze (na + 1) $ \p -> do + B.copyByteArrayToPtr a p + pokeByteOff p na b + where na = B.length a +{-# INLINE snoc #-} + +-- This version of append is more polymorphic and requires no trampoline when +-- fed with an input backed by Block Word8. +append :: (ByteArrayAccess a, ByteArrayAccess b) => a -> b -> ScrubbedBytes +append a 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 +{-# INLINE append #-} + +prf :: ByteArrayAccess s => Word -> s -> Word8 -> ScrubbedBytes +prf !eta s !b = case someNatVal (8 * 64 * fromIntegral eta) of + SomeNat proxy -> unDigest (doHash proxy) + where + doHash :: KnownNat bitlen => proxy bitlen -> Digest (SHAKE256 bitlen) + doHash _ = hash (snoc s b) + +h :: ByteArrayAccess s => s -> Bytes +h = Builder.run . hashWith SHA3_256 + +j :: ScrubbedBytes -> ScrubbedBytes +j = Builder.run . hashWith (SHAKE256 :: SHAKE256 256) + +g :: ByteArray ba => ScrubbedBytes -> (ba, B.View ScrubbedBytes) +g c = (B.convert $ B.takeView ab 32, B.dropView ab 32) + where ab = Builder.run $ hashWith SHA3_512 c + +-- Override cryptonite 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 -> do + ctx <- hashMutableInit + hashMutableUpdate (ctx :: MutableContext a) ba + B.withByteArray ctx $ \pctx -> do + hashInternalFinalize (castPtr pctx :: Ptr (Context a)) dig + ScrubbedBlock.erasePtr (B.length ctx) pctx diff --git a/src/Crypto/PubKey/ML_KEM.hs b/src/Crypto/PubKey/ML_KEM.hs new file mode 100644 --- /dev/null +++ b/src/Crypto/PubKey/ML_KEM.hs @@ -0,0 +1,94 @@ +-- | +-- Module : Crypto.PubKey.ML_KEM +-- License : BSD-3-Clause +-- Maintainer : Olivier Chéron +-- Stability : provisional +-- Portability : unknown +-- +-- Module-Lattice-based Key-Encapsulation Mechanism (ML-KEM), defined +-- in . +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE TypeFamilies #-} +module Crypto.PubKey.ML_KEM + ( EncapsulationKey, DecapsulationKey, Ciphertext, SharedSecret + -- * Operations + , generate, generateWith, encapsulate, encapsulateWith, decapsulate + -- * Parameter sets + , ParamSet, ML_KEM_512, ML_KEM_768, ML_KEM_1024 + -- * Conversions and checks + , Decode(..), Encode(..) + , toPublic, checkKeyPair + ) where + +import Crypto.Random + +import Data.ByteArray (ByteArrayAccess, ScrubbedBytes) +import qualified Data.ByteArray as B + +import Internal + +-- | ML-KEM-512 (security category 1) +data ML_KEM_512 = ML_KEM_512 deriving Show +-- | ML-KEM-768 (security category 3) +data ML_KEM_768 = ML_KEM_768 deriving Show +-- | ML-KEM-1024 (security category 5) +data ML_KEM_1024 = ML_KEM_1024 deriving Show + +instance ParamSet ML_KEM_512 where + type K ML_KEM_512 = 2 + getParams _ = Params 3 2 10 4 +instance ParamSet ML_KEM_768 where + type K ML_KEM_768 = 3 + getParams _ = Params 2 2 10 4 +instance ParamSet ML_KEM_1024 where + type K ML_KEM_1024 = 4 + getParams _ = Params 2 2 11 5 + +-- | Generate an ML-KEM key pair from a random seed. +generate :: (ParamSet a, MonadRandom m) + => proxy a -> m (EncapsulationKey a, DecapsulationKey a) +generate p = do + d <- getRandomBytes 32 + z <- getRandomBytes 32 + return (Internal.keyGen p (d :: ScrubbedBytes) z) + +-- | Generate an ML-KEM key pair from the specified seed (d, z). Length of +-- inputs must be 32 bytes. +generateWith :: (ParamSet a, ByteArrayAccess d, ByteArrayAccess z) + => proxy a -> d -> z -> Maybe (EncapsulationKey a, DecapsulationKey a) +generateWith p d z + | B.length d /= 32 = Nothing + | B.length z /= 32 = Nothing + | otherwise = Just $ Internal.keyGen p d (B.convert z) + +-- | Generate a shared secret key and an associated ciphertext using randomness. +encapsulate :: (ParamSet a, MonadRandom m) + => EncapsulationKey a -> m (SharedSecret a, Ciphertext a) +encapsulate ek = do + m <- getRandomBytes 32 + return (Internal.encaps ek (m :: ScrubbedBytes)) + +-- | Generate a shared secret key and an associated ciphertext using a +-- specified random input. This byte array must be 32 bytes and not repeated +-- with other encapsulations. For testing purposes. +encapsulateWith :: (ParamSet a, ByteArrayAccess m) + => EncapsulationKey a -> m -> Maybe (SharedSecret a, Ciphertext a) +encapsulateWith ek m + | B.length m /= 32 = Nothing + | otherwise = Just $ Internal.encaps ek m + +-- | Return the shared secret for a given ciphertext. Does implicit rejection +-- in the event the ciphertext or encapsulation key have been tampered with. +decapsulate :: ParamSet a => DecapsulationKey a -> Ciphertext a -> SharedSecret a +decapsulate = Internal.decaps + +-- | Try to detect corruptions in a pair of keys. Note that this does not +-- fully guarantee that the key pair was properly generated. Returns @True@ +-- when the key pair is found valid. +checkKeyPair :: (ParamSet a, MonadRandom m) + => (EncapsulationKey a, DecapsulationKey a) -> m Bool +checkKeyPair (ek, dk) = do + m <- getRandomBytes 32 + let (kk, ct) = Internal.encaps ek (m :: ScrubbedBytes) + kk' = Internal.decaps dk ct + return (kk' == kk) diff --git a/src/Internal.hs b/src/Internal.hs new file mode 100644 --- /dev/null +++ b/src/Internal.hs @@ -0,0 +1,188 @@ +-- | +-- Module : Internal +-- License : BSD-3-Clause +-- Copyright : (c) 2025 Olivier Chéron +-- +-- ML-KEM main internal algorithms +-- +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE TypeFamilies #-} +module Internal + ( ParamSet(..), Params(..), Encode(..), Decode(..) + , DecapsulationKey, EncapsulationKey, Ciphertext, SharedSecret + , keyGen, toPublic, encaps, decaps + ) where + +import Basement.Nat + +import Control.DeepSeq (NFData(..)) +import Control.Monad + +import Data.ByteArray (ByteArray, ByteArrayAccess, Bytes, ScrubbedBytes) +import qualified Data.ByteArray as B + +import qualified Builder +import qualified Crypto +import qualified K_PKE as K +import K_PKE (Params(..)) +import Marking (Leak(..)) + +-- | The class of ML-KEM parameter sets. +class KnownNat (K a) => ParamSet a where + type K a :: Nat + getParams :: proxy a -> Params (K a) + +-- | Utility class to serialize ML-KEM objects to byte arrays. +class Encode obj where + -- | Serializes an object to a sequence of bytes. + encode :: ByteArray ba => obj a -> ba + +-- | Utility class to deserialize ML-KEM 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-KEM decapsulation key, aka private key. +data DecapsulationKey a = DK (K.DecryptionKey (K a)) (K.EncryptionKey (K a)) Bytes ScrubbedBytes + +-- | An ML-KEM encapsulation key, aka public key. +data EncapsulationKey a = EK Bytes (K.EncryptionKey (K a)) + +-- | The ciphertext produced by the encapsulation function and consumed by the +-- decapsulation function. +newtype Ciphertext a = C Bytes deriving (Eq, ByteArrayAccess) + +-- | A shared secret returned by the encapsulation and decapsulation functions. +-- Length is 32 bytes for all defined parameter sets. +newtype SharedSecret a = S ScrubbedBytes deriving ByteArrayAccess + +instance Eq (DecapsulationKey a) where + DK dk1 ek1 h1 z1 == DK dk2 ek2 h2 z2 = Crypto.toBool $ + Crypto.constEqW dk1 dk2 `Crypto.andW` + Crypto.constEqW ek1 ek2 `Crypto.andW` + Crypto.constEqW h1 h2 `Crypto.andW` + Crypto.constEqW z1 z2 + +instance Eq (EncapsulationKey a) where + EK _ ek1 == EK _ ek2 = Crypto.toBool $ Crypto.constEqW ek1 ek2 + +instance Eq (SharedSecret a) where + S a == S b = Crypto.toBool $ Crypto.constEqW a b + +instance Show (DecapsulationKey a) where +#ifdef ML_KEM_TESTING + showsPrec d dk = showParen (d > 10) $ + showString "DecapsulationKey " . showsPrec 11 (encode dk :: Bytes) +#else + showsPrec _ _ = showString "DecapsulationKey" +#endif + +instance Show (EncapsulationKey a) where + showsPrec d ek = showParen (d > 10) $ + showString "EncapsulationKey " . showsPrec 11 (encode ek :: Bytes) + +instance Show (Ciphertext a) where + showsPrec d (C ct) = showParen (d > 10) $ + showString "Ciphertext " . showsPrec 11 ct + +instance Show (SharedSecret a) where +#ifdef ML_KEM_TESTING + showsPrec d (S kk) = showParen (d > 10) $ + showString "SharedSecret " . showsPrec 11 kk +#else + showsPrec _ _ = showString "SharedSecret" +#endif + +instance NFData (DecapsulationKey a) where + rnf (DK dk ek h z) = rnf dk `seq` rnf ek `seq` rnf h `seq` rnf z + +instance NFData (EncapsulationKey a) where + rnf (EK _ ek) = rnf ek -- h omitted because just for caching + +instance NFData (Ciphertext a) where + rnf (C c) = rnf c + +instance NFData (SharedSecret a) where + rnf (S kk) = rnf kk + +instance Encode EncapsulationKey where + encode (EK _ ek) = Builder.runRelaxed $ K.ekEncode ek + +instance Decode EncapsulationKey where + decode p input = EK (Crypto.h input) <$> K.ekDecode params input + where params = getParams p + +instance Encode DecapsulationKey where + encode (DK dk ek h z) = Builder.runRelaxed $ + leak (K.dkEncode dk) <> K.ekEncode ek <> Builder.bytes h <> leak (Builder.bytes z) + +instance Decode DecapsulationKey where + decode p input = do + -- decapsulation key type check: + guard (B.length input == 768 * k + 96) + let dks = B.view input 0 (384 * k) + eks = B.view input (384 * k) (384 * k + 32) + !h = B.convert $ B.view input (768 * k + 32) 32 + -- hash check: + guard (Crypto.toBool $ Crypto.constEqW h (Crypto.h eks)) + let !dk = K.dkDecode dks + !ek <- K.ekDecode params eks + let !z = B.convert $ B.view input (768 * k + 64) 32 + return (DK dk ek h z) + where + params = getParams p + k = K.dimension params + +instance Decode Ciphertext where + decode p input + -- ciphertext type check: + | B.length input == 32 * (du * k + dv) = Just (C $ B.convert input) + | otherwise = Nothing + where + params@Params{..} = getParams p + k = K.dimension params + +instance Decode SharedSecret where + decode _ input + | B.length input == 32 = Just (S $ B.convert input) + | otherwise = Nothing + +-- Uses randomness to generate an encapsulation key and a corresponding decapsulation key +keyGen :: (ParamSet a, ByteArrayAccess d) => proxy a -> d -> ScrubbedBytes -> (EncapsulationKey a, DecapsulationKey a) +keyGen p d z = (EK h ek, DK dk ek h z) + where + params = getParams p + (ek, dk) = K.keyGen params d + h = Crypto.h $ Builder.run (K.ekEncode ek) + +-- | Returns the encapsulation key embedded in the given decapsulation key. +-- Note that they may not necessarily match when the decapsulation key was +-- decoded from an untrusted source. +toPublic :: DecapsulationKey a -> EncapsulationKey a +toPublic (DK _ ek h _) = EK h ek + +-- Uses the encapsulation key and randomness to generate a key and an associated ciphertext +encaps :: (ParamSet a, ByteArrayAccess m) => EncapsulationKey a -> m -> (SharedSecret a, Ciphertext a) +encaps p@(EK h ek) m = (S kk, C c) + where + params = getParams p + (kk, r) = Crypto.g (m `Crypto.append` h) + c = K.encrypt params ek m r + +-- Uses the decapsulation key to produce a shared secret key from a ciphertext +decaps :: ParamSet a => DecapsulationKey a -> Ciphertext a -> SharedSecret a +decaps p@(DK dk ek h z) (C c) = S $ + Crypto.constSelectBytes + (Crypto.constEqW c c') -- condition + kk' -- when equal + (Crypto.j (z `Crypto.append` c)) -- when different + where + params = getParams p + m' = K.decrypt params dk c + (kk', r') = Crypto.g (m' `Crypto.append` h) + c' = K.encrypt params ek m' r' diff --git a/src/K_PKE.hs b/src/K_PKE.hs new file mode 100644 --- /dev/null +++ b/src/K_PKE.hs @@ -0,0 +1,143 @@ +-- | +-- Module : K_PKE +-- License : BSD-3-Clause +-- Copyright : (c) 2025 Olivier Chéron +-- +-- The K-PKE component scheme +-- +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE KindSignatures #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE RecordWildCards #-} +module K_PKE + ( Params(..), dimension, keyGen, encrypt, decrypt + , DecryptionKey, dkEncode, dkDecode + , EncryptionKey, ekEncode, ekDecode + ) where + +import Basement.Nat +import Basement.Types.OffsetSize + +import Control.DeepSeq (NFData(..)) +import Control.Monad + +import Data.ByteArray (ByteArrayAccess, Bytes, ScrubbedBytes) +import qualified Data.ByteArray as B + +import Unsafe.Coerce + +import Auxiliary (Rq, Tq, (..+), (..-)) +import Builder (Builder) +import Marking (SecurityMarking(..), Leak(..)) +import Vector (Vector) +import qualified Auxiliary as Aux +import qualified Crypto +import qualified Builder +import Math +import qualified Matrix +import qualified Vector + +data Params (k :: Nat) = Params + { eta1 :: {-# UNPACK #-} !Word + , eta2 :: {-# UNPACK #-} !Word + , du :: {-# UNPACK #-} !Int + , dv :: {-# UNPACK #-} !Int + } + +dimension :: KnownNat k => Params k -> Int +dimension = fromIntegral . natVal + +class Leak t => LeakVec vec t where + leakVec :: vec (t Sec) -> vec (t Pub) + leakVec = unsafeCoerce + +instance LeakVec (Vector k) Tq +instance LeakVec (Vector k) Rq + +newtype DecryptionKey (k :: Nat) = DecryptionKey { dkS :: Vector k (Tq Sec) } +data EncryptionKey (k :: Nat) = EncryptionKey { ekT :: Vector k (Tq Pub), ekRho :: Bytes, ekA :: Vector k (Vector k (Tq Pub)) } + +instance Crypto.ConstEqW (DecryptionKey k) where + constEqW a b = Crypto.constEqW (dkS a) (dkS b) + +instance Crypto.ConstEqW (EncryptionKey k) where + constEqW a b = Crypto.constEqW (ekT a) (ekT b) `Crypto.andW` Crypto.constEqW (ekRho a) (ekRho b) + +instance NFData (DecryptionKey k) where + rnf = Vector.toNormalForm . dkS + +instance NFData (EncryptionKey k) where + rnf ek = Vector.toNormalForm (ekT ek) `seq` rnf (ekRho ek) + -- ekA omitted because just for caching + +ekEncode :: EncryptionKey k -> Builder Pub +ekEncode ek = Vector.concatMap Aux.byteEncode12 (ekT ek) <> Builder.bytes (ekRho ek) + +ekDecode :: (KnownNat k, ByteArrayAccess ba) => Params k -> ba -> Maybe (EncryptionKey k) +ekDecode params input = do + -- type check: + guard (B.length input == 384 * k + 32) + let !tt = Vector.create $ \i -> Aux.byteDecode12 (view384 i) + !rho = B.convert $ B.view input (384 * k) 32 + elem384 off = Builder.run (Aux.byteEncode12 (Vector.index tt off)) + -- modulus check: + forM_ [0 .. k - 1] $ \i -> guard (elem384 (Offset i) `B.eq` view384 (Offset i)) + let aa = Matrix.create $ \(Offset i) (Offset j) -> Aux.sampleNTT rho (fromIntegral j) (fromIntegral i) + Just EncryptionKey { ekT = tt, ekRho = rho, ekA = aa } + where + k = dimension params + view384 (Offset i) = B.view input (384 * i) 384 + +dkEncode :: DecryptionKey k -> Builder Sec +dkEncode = Vector.concatMap Aux.byteEncode12 . dkS + +dkDecode :: (KnownNat k, ByteArrayAccess ba) => ba -> DecryptionKey k +dkDecode input = do + let !dk = Vector.create $ \i -> Aux.byteDecode12 (view384 i) + in DecryptionKey { dkS = dk } + where + view384 (Offset i) = B.view input (384 * i) 384 + +-- Uses randomness to generate an encryption key and a corresponding decryption key +keyGen :: (KnownNat k, ByteArrayAccess d) => Params k -> d -> (EncryptionKey k, DecryptionKey k) +keyGen params@Params{..} d = (ek, dk) + where + k = dimension params + (rho, sigma) = Crypto.g (Crypto.snoc d (fromIntegral k)) + aa = Matrix.create $ \(Offset i) (Offset j) -> Aux.sampleNTT rho (fromIntegral j) (fromIntegral i) + s = Vector.create $ \(Offset i) -> Aux.samplePolyCBD eta1 (Crypto.prf eta1 sigma $ fromIntegral i) + e = Vector.create $ \(Offset i) -> Aux.samplePolyCBD eta1 (Crypto.prf eta1 sigma $ fromIntegral (k + i)) + !ss = Aux.ntt <$> s + ee = Aux.ntt <$> e + !tt = leakVec $ Matrix.mulw aa ss ee + ek = EncryptionKey { ekT = tt, ekRho = rho, ekA = aa } + dk = DecryptionKey { dkS = ss } + +-- Uses the encryption key to encrypt a plaintext message using the randomness 𝑟 +encrypt :: (KnownNat k, ByteArrayAccess m, ByteArrayAccess r) => Params k -> EncryptionKey k -> m -> r -> Bytes +encrypt params@Params{..} ek m r = Builder.run (c1 <> c2) + where + k = dimension params + tt = ekT ek + aa = ekA ek + y = Vector.create $ \(Offset i) -> Aux.samplePolyCBD eta1 (Crypto.prf eta1 r $ fromIntegral i) + e1 = Vector.create $ \(Offset i) -> Aux.samplePolyCBD eta2 (Crypto.prf eta2 r $ fromIntegral (k + i)) + e2 = Aux.samplePolyCBD eta2 (Crypto.prf eta2 r $ fromIntegral (2 * k)) + yy = Aux.ntt <$> y + u = leakVec $ (Aux.nttInv <$> Matrix.muly aa yy) .+ e1 + mu = Aux.rdecompress 1 (Aux.byteDecode1 m) + v = Aux.nttInv (tt `Matrix.mulz` yy) .+ e2 ..+ mu + c1 = Vector.concatMap (Aux.byteEncode du . Aux.rcompress du) u + c2 = Aux.byteEncode dv (Aux.rcompress dv v) + +-- Uses the decryption key to decrypt a ciphertext +decrypt :: KnownNat k => Params k -> DecryptionKey k -> Bytes -> ScrubbedBytes +decrypt params@Params{..} dk c = Builder.run m + where + k = dimension params + c2 = B.view c (32 * du * k) (32 * dv) + u' = Vector.create $ \(Offset i) -> Aux.rdecompress du . Aux.byteDecode du $ B.view c (32 * du * i) (32 * du) :: Rq Pub + v' = Aux.rdecompress dv (Aux.byteDecode dv c2) + w = v' ..- Aux.nttInv ((Aux.ntt <$> u') `Matrix.mulz` dkS dk) + m = Aux.byteEncode1 (Aux.rcompress 1 w) diff --git a/src/Marking.hs b/src/Marking.hs new file mode 100644 --- /dev/null +++ b/src/Marking.hs @@ -0,0 +1,139 @@ +-- | +-- 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(..) + ) where + +import Basement.Monad +import Basement.NormalForm +import Basement.PrimType +import Basement.Types.OffsetSize +#ifdef ML_KEM_TESTING +import qualified Basement.Compat.IsList +#endif + +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 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 + + create :: PrimType ty => CountOf ty -> (Offset ty -> ty) -> SecureBlock marking ty + index :: PrimType ty => SecureBlock marking ty -> Offset ty -> ty + map :: (PrimType a, PrimType b) => (a -> b) -> SecureBlock marking a -> SecureBlock marking b + new :: (PrimType ty, PrimMonad prim) => proxy marking -> CountOf ty -> prim (MutableBlock ty (PrimState prim)) + thaw :: (PrimType ty, PrimMonad m) => SecureBlock marking ty -> m (MutableBlock ty (PrimState m)) + unsafeCast :: PrimType b => SecureBlock marking a -> SecureBlock marking b + unsafeFreeze :: PrimMonad prim => MutableBlock ty (PrimState prim) -> prim (SecureBlock marking ty) + + toNormalForm :: SecureBlock marking ty -> () +#ifdef ML_KEM_TESTING + eq :: PrimType ty => SecureBlock marking ty -> SecureBlock marking ty -> Bool + showsPrec :: (PrimType ty, Show ty) => Int -> SecureBlock marking ty -> ShowS + fromList :: PrimType ty => [ty] -> SecureBlock marking ty + replicate :: PrimType ty => CountOf ty -> ty -> SecureBlock marking ty + toList :: PrimType ty => SecureBlock marking ty -> [ty] + 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 + + create = Block.create + {-# INLINE create #-} + index = blockIndex + map = Block.map + {-# INLINE map #-} + new _ = Block.new + thaw = Block.thaw + unsafeCast = Block.unsafeCast + unsafeFreeze = Block.unsafeFreeze + + toNormalForm = Basement.NormalForm.toNormalForm +#ifdef ML_KEM_TESTING + eq = (==) + showsPrec = Prelude.showsPrec + fromList = Basement.Compat.IsList.fromList + replicate = Block.replicate + toList = Basement.Compat.IsList.toList + 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 + + create = ScrubbedBlock.create + {-# INLINE create #-} + index = ScrubbedBlock.index + map = ScrubbedBlock.map + {-# INLINE map #-} + new _ = ScrubbedBlock.new + thaw = ScrubbedBlock.thaw + unsafeCast = ScrubbedBlock.unsafeCast + unsafeFreeze = ScrubbedBlock.unsafeFreeze + + toNormalForm = Basement.NormalForm.toNormalForm +#ifdef ML_KEM_TESTING + eq = (==) + showsPrec = Prelude.showsPrec + fromList = ScrubbedBlock.fromList + replicate = ScrubbedBlock.replicate + toList = ScrubbedBlock.toList + lengthBlock = ScrubbedBlock.length +#endif + + type SecureBytes Sec = ScrubbedBytes + unsafeCreate = ST.unsafeCreate + {-# INLINE unsafeCreate #-} + lengthBytes = B.length + copyByteArrayToPtr = B.copyByteArrayToPtr diff --git a/src/Math.hs b/src/Math.hs new file mode 100644 --- /dev/null +++ b/src/Math.hs @@ -0,0 +1,40 @@ +-- | +-- 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 MultiParamTypeClasses #-} +module Math + ( Add(..), Mul(..), MulAdd(..), BiMul(..), BiMulAdd(..) + ) where + +infixl 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 + -- invariant: biMulAdd a b c == a ..* b .+ c + biMulAdd :: b -> a -> a -> a diff --git a/src/Matrix.hs b/src/Matrix.hs new file mode 100644 --- /dev/null +++ b/src/Matrix.hs @@ -0,0 +1,46 @@ +-- | +-- Module : Matrix +-- License : BSD-3-Clause +-- Copyright : (c) 2025 Olivier Chéron +-- +-- A matrix here is simply a vector of vectors. The module also implements +-- 'mulz' as dot product and two utility functions 'mulw' and 'muly' that +-- multiply a matrix and a vector. +-- +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +module Matrix + ( create, mulw, muly, mulz +#ifdef ML_KEM_TESTING + , transpose +#endif + ) where + +import Basement.Nat +import Basement.Types.OffsetSize + +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 ty -> Offset (Vector n ty) -> ty +index a i j = Vector.index (Vector.index a j) i + +mulw :: (KnownNat n, BiMulAdd b a) => Vector m (Vector n b) -> Vector m a -> Vector n a -> Vector n a +mulw a !u !b = Vector.create $ \(Offset i) -> + Vector.foldIndexWith (\c (Offset j) vu -> biMulAdd (index a (Offset i) (Offset j)) vu c) (Vector.index b (Offset i)) u + +muly :: BiMulAdd b a => Vector m (Vector n b) -> Vector n a -> Vector m a +muly a !u = fmap (`mulz` u) a + +mulz :: BiMulAdd b a => Vector n b -> Vector n a -> a +mulz = Vector.fold1ZipWith (\c a b -> biMulAdd a b c) (..*) + +#ifdef ML_KEM_TESTING +transpose :: (KnownNat m, KnownNat n) => Vector m (Vector n ty) -> Vector n (Vector m ty) +transpose a = create $ \(Offset j) (Offset i) -> index a (Offset i) (Offset j) +#endif diff --git a/src/ScrubbedBlock.hs b/src/ScrubbedBlock.hs new file mode 100644 --- /dev/null +++ b/src/ScrubbedBlock.hs @@ -0,0 +1,158 @@ +-- | +-- 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. +-- +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE UnboxedTuples #-} +module ScrubbedBlock + ( ScrubbedBlock, create, foldZipWith, index + , ScrubbedBlock.length, ScrubbedBlock.map, new + , thaw, unsafeCast, unsafeFreeze, Block.withMutablePtr +#ifdef ML_KEM_TESTING + , ScrubbedBlock.fromList, ScrubbedBlock.replicate, ScrubbedBlock.toList +#endif + , erasePtr + ) where + +import Basement.Block (Block(..), MutableBlock(..), isPinned) +import Basement.Block.Mutable (mutableLengthBytes, unsafeCopyBytesRO) +import Basement.Compat.Primitive + +#ifdef ML_KEM_TESTING +import Basement.Compat.IsList +#endif +import Basement.Monad +import Basement.NormalForm +import Basement.PrimType +import Basement.Types.OffsetSize + +import Control.Exception (assert) +import Control.Monad.ST + +import Data.Word + +import Foreign.Ptr (Ptr) + +import Block (blockIndex, blockWrite) +import qualified Block + +#if MIN_VERSION_base(4,19,0) +import GHC.Base (Int(I#), setAddrRange#) +import GHC.Exts (Ptr(Ptr)) +#else +import Data.Memory.PtrMethods (memSet) +#endif +import GHC.Base (IO(IO), setByteArray#) +import GHC.Exts (getSizeofMutableByteArray#, mkWeak#) + +newtype ScrubbedBlock ty = ScrubbedBlock (Block ty) + deriving (Eq, Show, NormalForm) + +create :: PrimType ty + => CountOf ty + -> (Offset ty -> ty) + -> ScrubbedBlock ty +create n initializer = runST $ do + mb <- new n + loop mb 0 + unsafeFreeze mb + where + loop !mb i + | i .==# n = pure () + | otherwise = blockWrite mb i (initializer i) >> loop mb (i + 1) +{-# INLINE create #-} + +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 #-} + +index :: PrimType ty => ScrubbedBlock ty -> Offset ty -> ty +index (ScrubbedBlock b) = blockIndex b + +length :: PrimType ty => ScrubbedBlock ty -> CountOf ty +length (ScrubbedBlock b) = Block.length b + +map :: (PrimType a, PrimType b) => (a -> b) -> ScrubbedBlock a -> ScrubbedBlock b +map f (ScrubbedBlock b) = + create (CountOf n) $ \(Offset i) -> f (blockIndex b (Offset i)) + where + CountOf n = Block.length b +{-# INLINE map #-} + +new :: (PrimType ty, PrimMonad prim) => CountOf ty -> prim (MutableBlock ty (PrimState prim)) +new = Block.newPinned -- always pinned + +thaw :: (PrimType ty, PrimMonad m) => ScrubbedBlock ty -> m (MutableBlock ty (PrimState m)) +thaw (ScrubbedBlock b) = do + mb <- new (Block.length b) + unsafeCopyBytesRO mb 0 b 0 (mutableLengthBytes mb) + return mb + +unsafeCast :: PrimType b => ScrubbedBlock a -> ScrubbedBlock b +unsafeCast (ScrubbedBlock b) = ScrubbedBlock (Block.unsafeCast b) + +unsafeFreeze :: PrimMonad prim => MutableBlock ty (PrimState prim) -> prim (ScrubbedBlock ty) +unsafeFreeze mb = Block.unsafeFreeze mb >>= scrubbed + +#ifdef ML_KEM_TESTING +replicate :: PrimType ty => CountOf ty -> ty -> ScrubbedBlock ty +replicate n e = create n (const e) + +fromList :: PrimType ty => [ty] -> ScrubbedBlock ty +fromList elems = runST $ do + mb <- new (CountOf len) + go mb 0 elems + unsafeFreeze mb + where + !len = Prelude.length elems + + go !mb !i list = case list of + [] -> return () + (x:xs) -> blockWrite mb i x >> go mb (i + 1) xs + +toList :: PrimType ty => ScrubbedBlock ty -> [ty] +toList (ScrubbedBlock b) = Basement.Compat.IsList.toList b +#endif + + +{- internal -} + +assertPinned :: Block ty -> a -> a +assertPinned mb = assert (isPinned mb == Pinned) + +scrubbed :: PrimMonad prim => Block ty -> prim (ScrubbedBlock ty) +scrubbed b = assertPinned b $ unsafePrimFromIO $ do + addBlockFinalizer b (scrub $ Block.unsafeCast b) + return (ScrubbedBlock b) + +scrub :: Block Word8 -> IO () +scrub b = Block.unsafeThaw b >>= erase + +addBlockFinalizer :: Block ty -> IO () -> IO () +addBlockFinalizer (Block barr) (IO finalizer) = IO $ \s -> + case mkWeak# barr () finalizer s of { (# s1, _ #) -> (# s1, () #) } + +erase :: MutableBlock ty RealWorld -> IO () +erase (MutableBlock mbarr) = IO $ \s1 -> + case getSizeofMutableByteArray# mbarr s1 of + (# s2, len #) -> case setByteArray# mbarr 0# len 0# s2 of + s3 -> (# s3, () #) + +erasePtr :: Int -> Ptr Word8 -> IO () +#if MIN_VERSION_base(4,19,0) +erasePtr (I# n) (Ptr addr) = IO $ \s1 -> + case setAddrRange# addr n 0# s1 of + s2 -> (# s2, () #) +#else +erasePtr n ptr = memSet ptr 0 n +#endif diff --git a/src/SecureBlock.hs b/src/SecureBlock.hs new file mode 100644 --- /dev/null +++ b/src/SecureBlock.hs @@ -0,0 +1,35 @@ +-- | +-- 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, create, index, iterModify, Marking.map, new, thaw + , unsafeCast, unsafeFreeze, toNormalForm +#ifdef ML_KEM_TESTING + , eq, Marking.showsPrec, fromList, Marking.replicate, toList + , SecureBlock.length +#endif + ) where + +import Basement.Monad +import Basement.PrimType +#ifdef ML_KEM_TESTING +import Basement.Types.OffsetSize +#endif + +import Block (MutableBlock) +import Marking +import qualified Block + +iterModify :: (PrimType ty, PrimMonad prim) => (ty -> ty) -> MutableBlock ty (PrimState prim) -> prim () +iterModify = Block.iterModify + +#ifdef ML_KEM_TESTING +length :: (Classified marking, PrimType ty) => SecureBlock marking ty -> CountOf ty +length = lengthBlock +#endif 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,100 @@ +-- | +-- Module : Vector +-- License : BSD-3-Clause +-- Copyright : (c) 2025 Olivier Chéron +-- +-- A vector of lifted elements with the vector dimension at type level. +-- Currently backed by type t'Array' from basement. +-- +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE KindSignatures #-} +{-# LANGUAGE ScopedTypeVariables #-} +module Vector + ( Vector, Vector.concatMap + , Vector.fold1ZipWith, Vector.foldIndexWith, Vector.toNormalForm + , Vector.create, Vector.index +#ifdef ML_KEM_TESTING + , Vector.replicateM, Vector.zipWith +#endif + ) where + +import Basement.BoxedArray (Array) +import qualified Basement.BoxedArray as Array +import Basement.Compat.IsList +import Basement.Nat +import Basement.NormalForm +import Basement.Types.OffsetSize + +import Control.DeepSeq (NFData(..)) +#ifdef ML_KEM_TESTING +import Control.Monad +#endif + +#if !(MIN_VERSION_base(4,20,0)) +import Data.List (foldl') +#endif +import Data.Proxy + +import Math + +newtype Vector (n :: Nat) a = Vector { unVector :: Array a } + deriving (Eq, Show, Functor, NormalForm) + +instance (Add a, KnownNat n) => Add (Vector n a) where + zero = create (const zero) + (.+) = Vector.zipWith (.+) + (.-) = Vector.zipWith (.-) + neg (Vector a) = Vector (fmap neg a) + +create :: forall n a. KnownNat n => (Offset a -> a) -> Vector n a +create = Vector . Array.create (CountOf sz) + where !sz = fromIntegral $ natVal (Proxy :: Proxy n) +{-# INLINE create #-} + +arrayIndex :: Array a -> Offset a -> a +#ifdef ML_KEM_TESTING +arrayIndex = Array.index + +replicateM :: forall n m a. (KnownNat n, Applicative m) => m a -> m (Vector n a) +replicateM f = Vector . fromList <$> Control.Monad.replicateM sz f + where !sz = fromIntegral $ natVal (Proxy :: Proxy n) +#else +arrayIndex = Array.unsafeIndex +#endif + +index :: Vector n a -> Offset a -> a +index = arrayIndex . unVector + +concatMap :: Monoid b => (a -> b) -> Vector n a -> b +concatMap f = mconcat . Prelude.map f . toList . unVector +{-# INLINE concatMap #-} + +zipWith :: (a -> b -> c) -> Vector n a -> Vector n b -> Vector n c +zipWith f (Vector a) (Vector !b) = Vector $ + Array.create (CountOf sa) $ \(Offset i) -> + f (arrayIndex a (Offset i)) (arrayIndex b (Offset i)) + where + CountOf sa = Array.length a +{-# INLINE zipWith #-} + +fold1ZipWith :: (c -> a -> b -> c) -> (a -> b -> c) -> Vector n a -> Vector n b -> c +fold1ZipWith f g (Vector a) (Vector !b) = + foldl' ff gg [1 .. sa - 1] + where + ff x i = f x (arrayIndex a (Offset i)) (arrayIndex b (Offset i)) + gg = g (arrayIndex a 0) (arrayIndex b 0) + CountOf !sa = Array.length a +{-# INLINE fold1ZipWith #-} + +foldIndexWith :: (c -> Offset a -> a -> c) -> c -> Vector n a -> c +foldIndexWith f c (Vector a) = foldl' g c [0 .. sa - 1] + where + g x i = f x (Offset i) (arrayIndex a (Offset i)) + CountOf !sa = Array.length a +{-# INLINE foldIndexWith #-} + +toNormalForm :: NFData a => Vector n a -> () +toNormalForm = Array.foldl' (\acc x -> acc `seq` rnf x) () . unVector diff --git a/tests/EncapDecap.hs b/tests/EncapDecap.hs new file mode 100644 --- /dev/null +++ b/tests/EncapDecap.hs @@ -0,0 +1,127 @@ +-- | +-- Module : EncapDecap +-- License : BSD-3-Clause +-- Copyright : (c) 2025 Olivier Chéron +-- +-- Types that are specific to encapsulation/decapsulation test vectors +-- +{-# LANGUAGE OverloadedStrings #-} +module EncapDecap + ( TestGroup(..), Test(..), TestGroupPayload(..) + , EncapsulationExt(..), DecapsulationExt(..) + , EncapsulationKeyCheckExt(..), DecapsulationKeyCheckExt(..) + ) where + +import Data.Aeson +import Data.Aeson.Types +import Data.ByteString (ByteString) + +import Util + +data TestGroup = TestGroup + { tgId :: Int + , testType :: String + , parameterSet :: String + , function :: String + , payload :: TestGroupPayload + } deriving Show + +data TestGroupPayload + = FunctionEncapsulation [Test EncapsulationExt] + | FunctionDecapsulation [Test DecapsulationExt] + | FunctionEncapsulationKeyCheck [Test EncapsulationKeyCheckExt] + | FunctionDecapsulationKeyCheck [Test DecapsulationKeyCheckExt] + deriving Show + +parsePayload :: Object -> Parser TestGroupPayload +parsePayload o = do + fn <- o .: "function" + case fn of + "encapsulation" -> FunctionEncapsulation <$> (o .: "tests") + "decapsulation" -> FunctionDecapsulation <$> (o .: "tests") + "encapsulationKeyCheck" -> FunctionEncapsulationKeyCheck <$> (o .: "tests") + "decapsulationKeyCheck" -> FunctionDecapsulationKeyCheck <$> (o .: "tests") + unknown -> fail ("parsePayload: unknown function " ++ unknown) + +instance FromJSON TestGroup where + parseJSON = withObject "TestGroup" $ \o -> TestGroup + <$> o .: "tgId" + <*> o .: "testType" + <*> o .: "parameterSet" + <*> o .: "function" + <*> parsePayload o + +data Test ext = Test + { tcId :: Int + , deferred :: Bool + , tcExt :: ext + } deriving Show + +class TestExt ext where + parseExt :: Object -> Parser ext + +instance TestExt ext => FromJSON (Test ext) where + parseJSON = withObject "Test" $ \o -> Test + <$> o .: "tcId" + <*> o .: "deferred" + <*> parseExt o + +data EncapsulationExt = EncapsulationExt + { ekEnc :: ByteString + , dkEnc :: ByteString + , cEnc :: ByteString + , kEnc :: ByteString + , mEnc :: ByteString + } deriving Show + +instance TestExt EncapsulationExt where + parseExt o = EncapsulationExt + <$> o .:: "ek" + <*> o .:: "dk" + <*> o .:: "c" + <*> o .:: "k" + <*> o .:: "m" + +data DecapsulationExt = DecapsulationExt + { ekDec :: ByteString + , dkDec :: ByteString + , cDec :: ByteString + , kDec :: ByteString + , reasonDec :: String + } deriving Show + +instance TestExt DecapsulationExt where + parseExt o = DecapsulationExt + <$> o .:: "ek" + <*> o .:: "dk" + <*> o .:: "c" + <*> o .:: "k" + <*> o .: "reason" + +data DecapsulationKeyCheckExt = DecapsulationKeyCheckExt + { passedDkc :: Bool + , ekDkc :: ByteString + , dkDkc :: ByteString + , reasonDkc :: String + } deriving Show + +instance TestExt DecapsulationKeyCheckExt where + parseExt o = DecapsulationKeyCheckExt + <$> o .: "testPassed" + <*> o .:: "ek" + <*> o .:: "dk" + <*> o .: "reason" + +data EncapsulationKeyCheckExt = EncapsulationKeyCheckExt + { passedEkc :: Bool + , ekEkc :: ByteString + , dkEkc :: ByteString + , reasonEkc :: String + } deriving Show + +instance TestExt EncapsulationKeyCheckExt where + parseExt o = EncapsulationKeyCheckExt + <$> o .: "testPassed" + <*> o .:: "ek" + <*> o .:: "dk" + <*> o .: "reason" diff --git a/tests/KeyGen.hs b/tests/KeyGen.hs new file mode 100644 --- /dev/null +++ b/tests/KeyGen.hs @@ -0,0 +1,48 @@ +-- | +-- 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 + , d :: ByteString + , z :: ByteString + , ek :: ByteString + , dk :: ByteString + } deriving Show + +instance FromJSON Test where + parseJSON = withObject "Test" $ \o -> Test + <$> o .: "tcId" + <*> o .: "deferred" + <*> o .:: "d" + <*> o .:: "z" + <*> o .:: "ek" + <*> o .:: "dk" diff --git a/tests/Tests.hs b/tests/Tests.hs new file mode 100644 --- /dev/null +++ b/tests/Tests.hs @@ -0,0 +1,468 @@ +-- | +-- Module : Main +-- License : BSD-3-Clause +-- Copyright : (c) 2025 Olivier Chéron +-- +-- The ML-KEM test suite. Can be instanciated twice, with and without the +-- @ML_KEM_TESTING@ macro to run property testing with assertions enabled in +-- the internal modules. +-- +{-# LANGUAGE CPP #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE ExistentialQuantification #-} +{-# LANGUAGE KindSignatures #-} +module Main (main) where + +import Data.ByteArray (Bytes) +import qualified Data.ByteArray as B + +import Control.Monad + +import Data.Maybe (isJust, fromJust) +import Data.Proxy + +import GHC.IO.Exception (IOErrorType(..)) + +import System.Directory (doesFileExist) +import System.IO.Error (catchIOError, mkIOError) +import System.Process (readProcess) + +#ifdef ML_KEM_TESTING + +import Data.Bits +import Data.Word + +import GHC.TypeNats + +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_KEM as Lib + +import Test.Tasty +import Test.Tasty.HUnit +import Test.Tasty.QuickCheck + +import qualified EncapDecap +import qualified KeyGen +import qualified Vectors + +#ifdef ML_KEM_TESTING + +newtype Bit7 = Bit7 Word8 deriving Show + +instance Arbitrary Bit7 where +#if (MIN_VERSION_tasty_quickcheck(0,10,2)) + arbitrary = Bit7 <$> chooseBoundedIntegral (0, 127) +#else + arbitrary = Bit7 <$> choose (0, 127) +#endif + +newtype FE = FE { unFE :: Zq} deriving Show + +instance Arbitrary FE where +#if (MIN_VERSION_tasty_quickcheck(0,10,2)) + arbitrary = FE . toZq <$> chooseBoundedIntegral (0, 3328) +#else + arbitrary = FE . toZq <$> choose (0, 3328) +#endif + +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 + +newtype D = D Int deriving Show + +instance Arbitrary D where + arbitrary = D <$> choose (0, 12) + +data Dim = forall (n :: Nat). KnownNat n => Dim (Proxy n) + +instance Show Dim where + show (Dim n) = show n + +instance Arbitrary Dim where + arbitrary = toDim <$> choose (1, 9) + +toDim :: Int -> Dim +toDim n = case someNatVal (fromIntegral n) of SomeNat p -> Dim p + +type VElem = Zq -- test with any ring but Tq would also work here + +arbitraryVector :: KnownNat n => proxy n -> Gen (Vector n VElem) +arbitraryVector _ = Vector.replicateM (unFE <$> arbitrary) + +arbitraryMatrix :: (KnownNat m, KnownNat n) => proxy n -> proxy m -> Gen (Vector n (Vector m VElem)) +arbitraryMatrix _ m = Vector.replicateM (arbitraryVector m) + +arbitraryBytes :: Int -> Gen Bytes +arbitraryBytes n = B.pack <$> vectorOf n arbitrary + +byteDecodeBytes :: Int -> Bytes -> BlockN Sec 256 Word16 +byteDecodeBytes = byteDecode + +byteEncodeBytes :: Int -> BlockN Sec 256 Word16 -> Bytes +byteEncodeBytes d = runBytes . byteEncode d + +byteEncodeBytes1 :: BlockN Sec 256 Word16 -> Bytes +byteEncodeBytes1 = runBytes . byteEncode1 + +byteEncodeBytes12 :: Tq Sec -> Bytes +byteEncodeBytes12 = runBytes . byteEncode12 + +runBytes :: Builder Sec -> Bytes +runBytes = Builder.run . leak + +#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_KEM_512) + , P (Proxy :: Proxy ML_KEM_768) + , P (Proxy :: Proxy ML_KEM_1024) + ] + +toP :: String -> P +toP "ML-KEM-512" = P (Proxy :: Proxy ML_KEM_512) +toP "ML-KEM-768" = P (Proxy :: Proxy ML_KEM_768) +toP "ML-KEM-1024" = P (Proxy :: Proxy ML_KEM_1024) +toP paramSet = error ("unknown parameter set " ++ paramSet) + +withVectors :: (IO () -> TestTree) -> TestTree +withVectors = withResource alloc free + where + scriptPath = "tests/get-vectors.sh" + free _ = return () + alloc = do + keyGenExists <- doesFileExist "tests/keyGen.json.gz" + encapDecapExists <- doesFileExist "tests/encapDecap.json.gz" + unless (keyGenExists && encapDecapExists) $ 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) + ) + +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 + eks = Lib.encode ek + dks = Lib.encode dk + (ek, dk) = fromJust $ Lib.generateWith p (KeyGen.d t) (KeyGen.z t) + assertEqual ("ek mismatch for tcId=" ++ show tcId) (KeyGen.ek t) eks + assertEqual ("dk mismatch for tcId=" ++ show tcId) (KeyGen.dk t) dks + +encapDecapVectors :: (String -> IO ()) -> Assertion +encapDecapVectors step = do + step "Reading test vectors ..." + file <- Vectors.readJson "tests/encapDecap.json.gz" + forM_ (Vectors.testGroups file) $ \group -> do + let paramSet = EncapDecap.parameterSet group + step (paramSet ++ " (" ++ EncapDecap.function group ++ ")") + case toP paramSet of + P p -> case EncapDecap.payload group of + EncapDecap.FunctionEncapsulation tests -> + forM_ tests (testEncapsulation p) + EncapDecap.FunctionDecapsulation tests -> + forM_ tests (testDecapsulation p) + EncapDecap.FunctionEncapsulationKeyCheck tests -> + forM_ tests (testEncapsulationKeyCheck p) + EncapDecap.FunctionDecapsulationKeyCheck tests -> + forM_ tests (testDecapsulationKeyCheck p) + where + ensureEk = id :: f (EncapsulationKey a) -> f (EncapsulationKey a) + ensureDk = id :: f (DecapsulationKey a) -> f (DecapsulationKey a) + testEncapsulation p test = do + let tcId = EncapDecap.tcId test + ext = EncapDecap.tcExt test + ek = fromJust $ Lib.decode p (EncapDecap.ekEnc ext) + k' = fromJust $ Lib.decode p (EncapDecap.kEnc ext) + c' = fromJust $ Lib.decode p (EncapDecap.cEnc ext) + (k, c) = fromJust $ Lib.encapsulateWith ek (EncapDecap.mEnc ext) + assertEqual ("k mismatch for tcId=" ++ show tcId) k' k + assertEqual ("c mismatch for tcId=" ++ show tcId) c' c + testDecapsulation p test = do + let tcId = EncapDecap.tcId test + ext = EncapDecap.tcExt test + dk = fromJust $ Lib.decode p (EncapDecap.dkDec ext) + c = fromJust $ Lib.decode p (EncapDecap.cDec ext) + k' = fromJust $ Lib.decode p (EncapDecap.kDec ext) + k = Lib.decapsulate dk c + assertEqual ("k mismatch for tcId=" ++ show tcId) k' k + testEncapsulationKeyCheck p test = do + let tcId = EncapDecap.tcId test + ext = EncapDecap.tcExt test + mek = ensureEk $ Lib.decode p (EncapDecap.ekEkc ext) + assertBool ("opposite outcome for tcId=" ++ show tcId) + (EncapDecap.passedEkc ext == isJust mek) + testDecapsulationKeyCheck p test = do + let tcId = EncapDecap.tcId test + ext = EncapDecap.tcExt test + mdk = ensureDk $ Lib.decode p (EncapDecap.dkDkc ext) + assertBool ("opposite outcome for tcId=" ++ show tcId) + (EncapDecap.passedDkc ext == isJust mdk) + +main :: IO () +main = defaultMain $ testGroup "mlkem" + [ withVectors $ \_ -> testGroup "vectors" + [ testCaseSteps "keyGen" keyGenVectors + , testCaseSteps "encapDecap" encapDecapVectors + ] + , testGroup "properties" + [ testGroup "ML-KEM" + [ testProperty "encapsulate/decapsulate" $ \(P p) -> ioProperty $ do + (ek, dk) <- Lib.generate p + (kk, c) <- Lib.encapsulate ek + let kk' = Lib.decapsulate dk c + return (kk === kk') + , testProperty "encode/decode keys" $ \(P p) -> ioProperty $ do + (ek, dk) <- Lib.generate p + return $ conjoin + [ Just ek === Lib.decode p (Lib.encode ek :: Bytes) + , Just dk === Lib.decode p (Lib.encode dk :: Bytes) + ] + , testProperty "convert/decode ciphertext and shared secret" $ \(P p) -> ioProperty $ do + (ek, _) <- Lib.generate p + (kk, c) <- Lib.encapsulate ek + return $ conjoin + [ Just c === Lib.decode p (B.convert c :: Bytes) + , Just kk === Lib.decode p (B.convert kk :: Bytes) + ] + , testProperty "toPublic" $ \(P p) -> ioProperty $ do + (ek, dk) <- Lib.generate p + return (ek === toPublic dk) + , testProperty "checkKeyPair" $ \(P p) -> ioProperty $ + Lib.generate p >>= checkKeyPair + ] +#ifdef ML_KEM_TESTING + , testGroup "bitRev7" + [ testCase "powers of two" $ + let powers = [1, 2, 4, 8, 16, 32, 64] + in reverse powers @=? map bitRev7 powers + , testProperty "or" $ \(Bit7 a) (Bit7 b) -> + bitRev7 (a .|. b) === bitRev7 a .|. bitRev7 b + , testProperty "not" $ \(Bit7 a) -> + let comp = xor 127 + in bitRev7 (comp a) === comp (bitRev7 a) + , testProperty "involutive" $ \(Bit7 a) -> + a === bitRev7 (bitRev7 a) + , testProperty "preserves bit count" $ \(Bit7 a) -> + popCount a === popCount (bitRev7 a) + ] + , testGroup "compression" + [ testProperty "compress . decompress == id" $ \(D d) -> do + y <- choose (0, 2^d - 1) + return (d < 12 ==> y === compress d (decompress d y)) + ] + , testGroup "encoding" + [ testProperty "byteEncode . byteDecode == id" $ \(D d) -> do + b <- arbitraryBytes (32 * d) + return (b === byteEncodeBytes d (byteDecode d b)) + , testProperty "byteEncode1 . byteDecode1 == id" $ do + b <- arbitraryBytes 32 + return (b === byteEncodeBytes1 (byteDecode1 b)) + , testProperty "byteDecode12 . byteEncode12 == id" $ \(PolyNTT p) -> + p === byteDecode12 (byteEncodeBytes12 p) + , testProperty "byteEncode 8" $ \x -> + B.replicate 256 x === byteEncodeBytes 8 (BlockN.replicate $ fromIntegral x) + , testCase "byteEncode 1 (zeros)" $ + B.replicate 32 0 @=? byteEncodeBytes 1 (BlockN.replicate 0) + , testCase "byteEncode 1 (ones)" $ + B.replicate 32 255 @=? byteEncodeBytes 1 (BlockN.replicate 1) + , testProperty "byteDecode1 == byteDecode 1" $ do + b <- arbitraryBytes 32 + return (byteDecodeBytes 1 b === byteDecode1 b) + ] + , testGroup "Zq" + [ testProperty "toZq . fromZq == id " $ \(FE a) -> + a === toZq (fromZq a) + , testProperty "fromZq . toZq == id " $ \a -> + mod a 3329 === fromZq (toZq a) + , testCase "field order" $ zero @=? toZq 3329 + , 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 + ] + , testProperty "mulAdd" $ \(FE a) (FE b) (FE c) -> + a .* b .+ c === mulAdd a b c + ] + , 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) + ] + , 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)) + ] + , testGroup "Matrix" + [ testProperty "mulw distributive left" $ \(Dim n) (Dim m) -> do + (a, b, u, v) <- (,,,) <$> arbitraryMatrix m n <*> arbitraryMatrix m n <*> arbitraryVector m <*> arbitraryVector n + return (mulw (a .+ b) u v === mulw a u (mulw b u zero) .+ v) + , testProperty "mulw distributive right" $ \(Dim n) (Dim m) -> do + (a, u, v, w) <- (,,,) <$> arbitraryMatrix m n <*> arbitraryVector m <*> arbitraryVector m <*> arbitraryVector n + return (mulw a (u .+ v) w === mulw a u (mulw a v w)) + , testProperty "muly definition" $ \(Dim n) (Dim m) -> do + (a, u, v) <- (,,) <$> arbitraryMatrix n m <*> arbitraryVector m <*> arbitraryVector n + return ((a `muly` u) .+ v == mulw (transpose a) u v) + , testProperty "muly distributive left" $ \(Dim n) (Dim m) -> do + (a, b, u) <- (,,) <$> arbitraryMatrix n m <*> arbitraryMatrix n m <*> arbitraryVector m + return ((a .+ b) `muly` u === (a `muly` u) .+ (b `muly` u)) + , testProperty "muly distributive right" $ \(Dim n) (Dim m) -> do + (a, u, v) <- (,,) <$> arbitraryMatrix n m <*> arbitraryVector m <*> arbitraryVector m + return (a `muly` (u .+ v) === (a `muly` u) .+ (a `muly` v)) + , testProperty "mulz commutative" $ \(Dim n) -> do + (u, v) <- (,) <$> arbitraryVector n <*> arbitraryVector n + return (u `mulz` v === v `mulz` u) + , testProperty "mulz distributive" $ \(Dim n) -> do + (u, v, w) <- (,,) <$> arbitraryVector n <*> arbitraryVector n <*> arbitraryVector n + return $ conjoin + [ u `mulz` (v .+ w) === (u `mulz` v) .+ (u `mulz` w) + , (u .+ v) `mulz` w === (u `mulz` w) .+ (v `mulz` 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,41 @@ +-- | +-- 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 + +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-KEM 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/cbb097019deb25eb966932ee293d89417482d8bc +CURL=curl + +for KEY in keyGen encapDecap; do + FILENAME="$DESTDIR"/$KEY.json.gz + URL=https://codeberg.org/ocheron/hs-mlkem/raw/$REF/tests/$KEY.json.gz + + "$CURL" --silent -o "$FILENAME" "$URL" || exit $? +done