{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# OPTIONS_GHC -Wno-orphans #-}

module Test.Crypto.Instances (
  withMLSBFromPSB,
  withMLockedSeedFromPSB,
) where

import Cardano.Crypto.Libsodium
import Cardano.Crypto.Libsodium.MLockedSeed
import Cardano.Crypto.PinnedSizedBytes
import Control.Monad.Class.MonadST
import Control.Monad.Class.MonadThrow
import Data.Maybe (mapMaybe)
import Data.Proxy (Proxy (Proxy))
import GHC.Exts (fromList, fromListN, toList)
import GHC.TypeLits (KnownNat, natVal)
import Test.QuickCheck (Arbitrary (..))
import qualified Test.QuickCheck.Gen as Gen

-- We cannot allow this instance, because it doesn't guarantee timely
-- forgetting of the MLocked memory, and in a QuickCheck context, where
-- tens of thousands of these values may be generated, waiting for GC to clean
-- up after us could have us run over our mlock quota.
--
-- Instead, use 'arbitrary' to generate a suitably sized PinnedSizedBytes
-- value, and then mlsbFromPSB or withMLSBFromPSB to convert it to an
-- MLockedSizedBytes value.
--
-- instance KnownNat n => Arbitrary (MLockedSizedBytes n) where
--     arbitrary = unsafePerformIO . mlsbFromByteString . BS.pack <$> vectorOf size arbitrary
--       where
--         size :: Int
--         size = fromInteger (natVal (Proxy :: Proxy n))

mlsbFromPSB :: (MonadST m, KnownNat n) => PinnedSizedBytes n -> m (MLockedSizedBytes n)
mlsbFromPSB :: forall (m :: * -> *) (n :: Nat).
(MonadST m, KnownNat n) =>
PinnedSizedBytes n -> m (MLockedSizedBytes n)
mlsbFromPSB = ByteString -> m (MLockedSizedBytes n)
forall (n :: Nat) (m :: * -> *).
(KnownNat n, MonadST m) =>
ByteString -> m (MLockedSizedBytes n)
mlsbFromByteString (ByteString -> m (MLockedSizedBytes n))
-> (PinnedSizedBytes n -> ByteString)
-> PinnedSizedBytes n
-> m (MLockedSizedBytes n)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PinnedSizedBytes n -> ByteString
forall (n :: Nat). PinnedSizedBytes n -> ByteString
psbToByteString

withMLSBFromPSB ::
  (MonadST m, MonadThrow m, KnownNat n) => PinnedSizedBytes n -> (MLockedSizedBytes n -> m a) -> m a
withMLSBFromPSB :: forall (m :: * -> *) (n :: Nat) a.
(MonadST m, MonadThrow m, KnownNat n) =>
PinnedSizedBytes n -> (MLockedSizedBytes n -> m a) -> m a
withMLSBFromPSB PinnedSizedBytes n
psb =
  m (MLockedSizedBytes n)
-> (MLockedSizedBytes n -> m ())
-> (MLockedSizedBytes n -> m a)
-> m a
forall a b c. m a -> (a -> m b) -> (a -> m c) -> m c
forall (m :: * -> *) a b c.
MonadThrow m =>
m a -> (a -> m b) -> (a -> m c) -> m c
bracket
    (PinnedSizedBytes n -> m (MLockedSizedBytes n)
forall (m :: * -> *) (n :: Nat).
(MonadST m, KnownNat n) =>
PinnedSizedBytes n -> m (MLockedSizedBytes n)
mlsbFromPSB PinnedSizedBytes n
psb)
    MLockedSizedBytes n -> m ()
forall (m :: * -> *) (n :: Nat).
MonadST m =>
MLockedSizedBytes n -> m ()
mlsbFinalize

mlockedSeedFromPSB :: (MonadST m, KnownNat n) => PinnedSizedBytes n -> m (MLockedSeed n)
mlockedSeedFromPSB :: forall (m :: * -> *) (n :: Nat).
(MonadST m, KnownNat n) =>
PinnedSizedBytes n -> m (MLockedSeed n)
mlockedSeedFromPSB = (MLockedSizedBytes n -> MLockedSeed n)
-> m (MLockedSizedBytes n) -> m (MLockedSeed n)
forall a b. (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap MLockedSizedBytes n -> MLockedSeed n
forall (n :: Nat). MLockedSizedBytes n -> MLockedSeed n
MLockedSeed (m (MLockedSizedBytes n) -> m (MLockedSeed n))
-> (PinnedSizedBytes n -> m (MLockedSizedBytes n))
-> PinnedSizedBytes n
-> m (MLockedSeed n)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PinnedSizedBytes n -> m (MLockedSizedBytes n)
forall (m :: * -> *) (n :: Nat).
(MonadST m, KnownNat n) =>
PinnedSizedBytes n -> m (MLockedSizedBytes n)
mlsbFromPSB

withMLockedSeedFromPSB ::
  (MonadST m, MonadThrow m, KnownNat n) => PinnedSizedBytes n -> (MLockedSeed n -> m a) -> m a
withMLockedSeedFromPSB :: forall (m :: * -> *) (n :: Nat) a.
(MonadST m, MonadThrow m, KnownNat n) =>
PinnedSizedBytes n -> (MLockedSeed n -> m a) -> m a
withMLockedSeedFromPSB PinnedSizedBytes n
psb =
  m (MLockedSeed n)
-> (MLockedSeed n -> m ()) -> (MLockedSeed n -> m a) -> m a
forall a b c. m a -> (a -> m b) -> (a -> m c) -> m c
forall (m :: * -> *) a b c.
MonadThrow m =>
m a -> (a -> m b) -> (a -> m c) -> m c
bracket
    (PinnedSizedBytes n -> m (MLockedSeed n)
forall (m :: * -> *) (n :: Nat).
(MonadST m, KnownNat n) =>
PinnedSizedBytes n -> m (MLockedSeed n)
mlockedSeedFromPSB PinnedSizedBytes n
psb)
    MLockedSeed n -> m ()
forall (m :: * -> *) (n :: Nat). MonadST m => MLockedSeed n -> m ()
mlockedSeedFinalize

instance KnownNat n => Arbitrary (PinnedSizedBytes n) where
  arbitrary :: Gen (PinnedSizedBytes n)
arbitrary = do
    let Int
size :: Int = Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer -> Int) -> (Proxy n -> Integer) -> Proxy n -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Proxy n -> Integer
forall (n :: Nat) (proxy :: Nat -> *).
KnownNat n =>
proxy n -> Integer
natVal (Proxy n -> Int) -> Proxy n -> Int
forall a b. (a -> b) -> a -> b
$ forall (t :: Nat). Proxy t
forall {k} (t :: k). Proxy t
Proxy @n
    Gen ByteString
-> (ByteString -> Maybe (PinnedSizedBytes n))
-> Gen (PinnedSizedBytes n)
forall a b. Gen a -> (a -> Maybe b) -> Gen b
Gen.suchThatMap
      (Int -> [Item ByteString] -> ByteString
forall l. IsList l => Int -> [Item l] -> l
fromListN Int
size ([Word8] -> ByteString) -> Gen [Word8] -> Gen ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Gen Word8 -> Gen [Word8]
forall a. Int -> Gen a -> Gen [a]
Gen.vectorOf Int
size Gen Word8
forall a. Arbitrary a => Gen a
arbitrary)
      ByteString -> Maybe (PinnedSizedBytes n)
forall (n :: Nat).
KnownNat n =>
ByteString -> Maybe (PinnedSizedBytes n)
psbFromByteStringCheck
  shrink :: PinnedSizedBytes n -> [PinnedSizedBytes n]
shrink PinnedSizedBytes n
psb = case ByteString -> [Word8]
ByteString -> [Item ByteString]
forall l. IsList l => l -> [Item l]
toList (ByteString -> [Word8])
-> (PinnedSizedBytes n -> ByteString)
-> PinnedSizedBytes n
-> [Word8]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PinnedSizedBytes n -> ByteString
forall (n :: Nat). PinnedSizedBytes n -> ByteString
psbToByteString (PinnedSizedBytes n -> [Word8]) -> PinnedSizedBytes n -> [Word8]
forall a b. (a -> b) -> a -> b
$ PinnedSizedBytes n
psb of
    [Word8]
bytes -> ([Word8] -> Maybe (PinnedSizedBytes n))
-> [[Word8]] -> [PinnedSizedBytes n]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (ByteString -> Maybe (PinnedSizedBytes n)
forall (n :: Nat).
KnownNat n =>
ByteString -> Maybe (PinnedSizedBytes n)
psbFromByteStringCheck (ByteString -> Maybe (PinnedSizedBytes n))
-> ([Word8] -> ByteString) -> [Word8] -> Maybe (PinnedSizedBytes n)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Word8] -> ByteString
[Item ByteString] -> ByteString
forall l. IsList l => [Item l] -> l
fromList) ([[Word8]] -> [PinnedSizedBytes n])
-> ([Word8] -> [[Word8]]) -> [Word8] -> [PinnedSizedBytes n]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Word8] -> [[Word8]]
forall a. Arbitrary a => a -> [a]
shrink ([Word8] -> [PinnedSizedBytes n])
-> [Word8] -> [PinnedSizedBytes n]
forall a b. (a -> b) -> a -> b
$ [Word8]
bytes