mldsa-0.1.0.0: src/Crypto.hs

-- |
-- Module      : Crypto
-- License     : BSD-3-Clause
-- Copyright   : (c) 2026 Olivier Chéron
--
-- Crypto-related utilities like the ML-DSA hash functions, or more general
-- concerns like constant-time equality and selection.
--
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Crypto
    ( ConstEqW(..), BoolW, trueW, andW, ltW, lteW, toBool
    , h, h64, h128
    , BlockDigest, unBlockDigest, hashToBlock
    , Digest, unDigest, hash
    ) where

import Crypto.Hash (Context)
import Crypto.Hash.Algorithms
import Crypto.Hash.IO

import Control.Exception (assert, mask_)
import Control.Monad.ST

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

import Data.Bits
import Data.Word

import GHC.TypeNats

import Foreign.Marshal.Utils (fillBytes)
import Foreign.Ptr (Ptr, castPtr, plusPtr)
import Foreign.Storable (pokeByteOff)

import Block (Block)
import Builder (Builder)
import Marking (Classified)
import ScrubbedBlock (ScrubbedBlock)
import SecureBytes (SecureBytes)
import Vector (Vector)
import qualified Block
import qualified Builder
import qualified ByteArrayST as ST
import qualified ScrubbedBlock
import qualified Vector

unsafeShiftIR :: Word -> Int -> Word
unsafeShiftIR x s = fromIntegral ((fromIntegral x :: Int) `unsafeShiftR` s)
{-# INLINE unsafeShiftIR #-}

newtype BoolW = BoolW Word

#ifdef ML_DSA_TESTING
instance Show BoolW where
    showsPrec d = showsPrec d . toBool
#endif

toBool :: BoolW -> Bool
toBool (BoolW mask) = mask /= 0

falseW, trueW :: BoolW
falseW = BoolW 0
trueW = BoolW maxBound

notW :: BoolW -> BoolW
notW (BoolW a) = BoolW (complement a)

andW :: BoolW -> BoolW -> BoolW
andW (BoolW a) (BoolW b) = BoolW (a .&. b)

bitsW :: Int
bitsW = let BoolW x = falseW in finiteBitSize x

bytesW :: Int
bytesW = div bitsW 8

msbW :: Word -> Word
msbW x = x `unsafeShiftIR` (bitsW - 1)

eqW :: Word -> Word -> BoolW
eqW a b = isZeroW (a `xor` b)
  where isZeroW x = BoolW $ msbW (complement x .&. (x - 1))

ltW :: Word -> Word -> BoolW
ltW a b = BoolW $ msbW (a - b)

lteW :: Word -> Word -> BoolW
lteW a b = notW (ltW b a)

assertMultW :: Int -> a -> a
assertMultW n = assert (n .&. mask == 0)
  where mask = bytesW - 1

class ConstEqW a where
    constEqW :: a -> a -> BoolW

instance ConstEqW a => ConstEqW (Vector n a) where
    constEqW =
        Vector.fold1ZipWith (\mask x y -> mask `andW` constEqW x y) constEqW

instance ConstEqW (Block Word) where
    constEqW a b
        | Block.length a /= Block.length b = falseW
        | otherwise = Block.foldZipWith (\mask x y -> mask `andW` eqW x y) trueW a b

instance ConstEqW (ScrubbedBlock Word) where
    constEqW a b
        | ScrubbedBlock.length a /= ScrubbedBlock.length b = falseW
        | otherwise = ScrubbedBlock.foldZipWith (\mask x y -> mask `andW` eqW x y) trueW a b

instance ConstEqW Bytes where
    constEqW = bytesConstEqW

instance ConstEqW ScrubbedBytes where
    constEqW = bytesConstEqW

bytesConstEqW :: (ByteArrayAccess bs1, ByteArrayAccess bs2) => bs1 -> bs2 -> BoolW
bytesConstEqW a b
    | B.length a /= B.length b = falseW
    | otherwise = foldZipWith (\mask x y -> mask `andW` eqW x y) trueW a b

foldZipWith :: (ByteArrayAccess bs1, ByteArrayAccess bs2)
            => (c -> Word -> Word -> c) -> c -> bs1 -> bs2 -> c
foldZipWith f c a b = assert (sa == sb) $ assertMultW sa $ assertMultW sb $
    runST $ ST.withByteArray a $ \pa -> ST.withByteArray b $ \pb ->
        loop (pa :: Ptr Word) (pb :: Ptr Word) c 0
  where
    !sa = B.length a
    !sb = B.length b

    loop !pa !pb !acc i
        | i == sa = return acc
        | otherwise = do
            va <- ST.peek pa
            vb <- ST.peek pb
            loop (pa `plusPtr` bytesW) (pb `plusPtr` bytesW) (f acc va vb) (i + bytesW)
{-# INLINE foldZipWith #-}

h :: (Classified marking, ByteArrayAccess (SecureBytes marking))
  => Int -> SecureBytes marking -> SecureBytes marking
h !len input = case someNatVal (fromIntegral (8 * len)) of
    SomeNat proxy -> Builder.run $ hashWith (alg proxy) input
  where
    alg :: proxy bitlen -> SHAKE256 bitlen
    alg _ = SHAKE256

h64 :: ByteArrayAccess (SecureBytes marking)
    => SecureBytes marking -> Builder marking
h64 = hashWith (SHAKE256 :: SHAKE256 512)

h128 :: ByteArrayAccess a
     => a -> Word8 -> Word8 -> (Bytes, ScrubbedBytes, ScrubbedBytes)
h128 xi !k !l = (a, b, c)
  where
    rho  = B.takeView bs 32
    rho' = B.view bs 32 64
    kk   = B.dropView bs 96

    !a = B.convert rho
    !b = B.convert rho'
    !c = B.convert kk

    bs :: ScrubbedBytes
    bs = Builder.run $ hashWith (SHAKE256 :: SHAKE256 1024) input

    input :: ScrubbedBytes
    input = B.unsafeCreate (len + 2) $ \p -> do
        B.copyByteArrayToPtr xi p
        pokeByteOff p len k
        pokeByteOff p (len + 1) l
    len = B.length xi

-- Override cryptonite/crypton types and hashing functions.
--
-- Standard type Digest is a newtype over an unpinned Block Word8, which
-- requires a trampoline to implement most Ptr access to the underlying byte
-- array.  Instead we re-implement here the Digest type over ScrubbedBytes as
-- well as pinned Block backends, to avoid trampoline costs.  Additionnally
-- we use the mutable API to avoid copying the hashing Context in between
-- steps init/update/finalize and then clear the content.

newtype Digest a = Digest { unDigest :: ScrubbedBytes }
newtype BlockDigest a = BlockDigest { unBlockDigest :: Block Word8 }

hash :: forall a ba. (HashAlgorithm a, ByteArrayAccess ba) => ba -> Digest a
hash = Digest . Builder.run . hashWith (undefined :: a)

hashToBlock :: forall a. HashAlgorithm a => Bytes -> BlockDigest a
hashToBlock = BlockDigest . Builder.runToBlock . hashWith (undefined :: a)

hashWith :: forall marking a ba. (HashAlgorithm a, ByteArrayAccess ba) => a -> ba -> Builder marking
hashWith a ba = Builder.unsafeCreate (hashDigestSize a) $ \dig ->
    hashMutableInit >>= \ctx -> mask_ $ do
        hashUpdateChunked (ctx :: MutableContext a) ba
        B.withByteArray ctx $ \pctx -> do
            hashInternalFinalize (castPtr pctx :: Ptr (Context a)) dig
            fillBytes pctx 0 (B.length ctx)

hashUpdateChunked :: (HashAlgorithm a, ByteArrayAccess ba) => MutableContext a -> ba -> IO ()
hashUpdateChunked ctx ba = B.withByteArray ba $ goChunked (B.length ba)
  where
    chunkSize = 1073741824  -- 1 GB

    goChunked remaining p
        | remaining > chunkSize = do
            hashMutableUpdate ctx (MemView p chunkSize)
            goChunked (remaining - chunkSize) (p `plusPtr` chunkSize)
        | otherwise = hashMutableUpdate ctx (MemView p remaining)