crypto-rng-0.3.0.2: test/Main.hs

{-# LANGUAGE CPP #-}
module Main (main) where

import Control.Exception
import Control.Monad
import Data.ByteString qualified as BS
import Data.Set qualified as S
import Test.Tasty
import Test.Tasty.HUnit

import Crypto.RNG
import Crypto.RNG.Utils

#if MIN_VERSION_random(1,3,0)
import Data.ByteString.Short qualified as SBS
import Data.Primitive.ByteArray
import System.Random.Stateful qualified as R
#endif

main :: IO ()
main = defaultMain $ testGroup "crypto-rng"
  [ testGroup "randomBytesIO" $ map bufferRefill configurations
  , testGroup "randomString"
    [ testCase "draws from the allowed chars" $ do
        rng <- newCryptoRNGState
        s <- runCryptoRNGT rng $ randomString 1000 alphabet
        assertEqual "length" 1000 (length s)
        assertBool "every char is allowed" $ all (`elem` alphabet) s
        assertBool "the whole alphabet shows up" $ all (`elem` s) alphabet
    , testCase "rejects an empty list of allowed chars" $ do
        rng <- newCryptoRNGState
        r <- try @ErrorCall $
          evaluate . length =<< runCryptoRNGT rng (randomString 8 "")
        case r of
          Left _ -> pure ()
          Right len -> assertFailure $ "returned a string of length " ++ show len
    ]
#if MIN_VERSION_random(1,3,0)
  , testGroup "uniformByteArrayM" [byteArrays]
#endif
  ]
  where
    alphabet :: [Char]
    alphabet = ['a' .. 'z'] ++ ['0' .. '9']

#if MIN_VERSION_random(1,3,0)
-- | The conversion from the generated bytes used to prepend a serialized
-- length, so the results were longer than requested and started with bytes
-- that were not random.
byteArrays :: TestTree
byteArrays = testCase "results have the requested length" $ do
  rng <- newCryptoRNGState
  forM_ [0, 1, 8, 16, 100, 1000] $ \n -> do
    forM_ [False, True] $ \isPinned -> do
      ba <- R.uniformByteArrayM isPinned n rng
      assertEqual ("byte array of " ++ show n) n (sizeofByteArray ba)
      when (isPinned && n > 0) $ do
        assertBool ("byte array of " ++ show n ++ " is pinned") (isByteArrayPinned ba)
    sbs <- R.uniformShortByteStringM n rng
    assertEqual ("short byte string of " ++ show n) n (SBS.length sbs)
#endif

-- | Buffer size paired with the request sizes to cycle through.
--
-- A request that is larger than the bytes left in the buffer is what triggers a
-- refill, so in each configuration the request sizes do not divide the buffer
-- size evenly.
configurations :: [(Int, [Int])]
configurations =
  [ (16, [10])
  , (16, [1, 7, 13, 40])
  , (32, [100])
  , (64, [20])
  , (32 * 1024, [100])
  , (1024, [3000, 10])
  , (1, [1, 2, 3])
  ]

-- | A refill used to hand out the bytes of the drained buffer a second time, so
-- a returned chunk repeated its own prefix and the repeat showed up again in
-- the following chunk.
bufferRefill :: (Int, [Int]) -> TestTree
bufferRefill (bufSize, sizes) = testCase name $ do
  rng <- newCryptoRNGStateSized bufSize
  chunks <- forM requestSizes $ \n -> do
    chunk <- randomBytesIO n rng
    assertEqual ("length of a " ++ show n ++ " byte request") n (BS.length chunk)
    assertBool ("a " ++ show n ++ " byte request repeats its own prefix") $
      not (repeatsPrefix 4 chunk)
    pure chunk
  let ws = windows 8 $ BS.concat chunks
  assertEqual "repeated windows" 0 (length ws - S.size (S.fromList ws))
  where
    name :: String
    name = "buffer of " ++ show bufSize ++ " bytes, requests of " ++ show sizes

    -- Enough requests to drain and refill the buffer several times.
    requestSizes :: [Int]
    requestSizes = takeUntilTotal (max 30000 (4 * bufSize)) (cycle sizes)

    takeUntilTotal :: Int -> [Int] -> [Int]
    takeUntilTotal _ [] = []
    takeUntilTotal remaining (n : ns)
      | remaining <= 0 = []
      | otherwise = n : takeUntilTotal (remaining - n) ns

    -- A shift by less than minLen bytes is left out, because a short match
    -- happens by chance often enough.
    repeatsPrefix :: Int -> BS.ByteString -> Bool
    repeatsPrefix minLen chunk = any matches [1 .. BS.length chunk - minLen]
      where
        matches :: Int -> Bool
        matches p = BS.drop p chunk == BS.take (BS.length chunk - p) chunk

    windows :: Int -> BS.ByteString -> [BS.ByteString]
    windows k bs
      | BS.length bs < k = []
      | otherwise = BS.take k bs : windows k (BS.drop 1 bs)