simplify

src/Simplex/Messaging/Crypto/SNTRUP761/Bindings.hs

@@ -1,11 +1,14 @@
 {-# LANGUAGE CPP #-}
 {-# LANGUAGE TypeApplications #-}
+{-# LANGUAGE PatternSynonyms #-}
 
 module Simplex.Messaging.Crypto.SNTRUP761.Bindings
-  ( KEMPublicKey (..),
+  ( KEMPublicKey,
     KEMSecretKey,
-    KEMCiphertext (..),
-    KEMSharedKey (..),
+    KEMCiphertext,
+    KEMSharedKey,
+    pattern KEMPublicKey,
+    pattern KEMSharedKey,
     KEMKeyPair,
     sntrup761Keypair,
     sntrup761Enc,
@@ -13,7 +16,6 @@ module Simplex.Messaging.Crypto.SNTRUP761.Bindings
   ) where
 
 import Control.Concurrent.STM
-import Control.Exception (throwIO)
 import Crypto.Random (ChaChaDRG)
 import Data.Aeson (FromJSON (..), ToJSON (..))
 import Data.Bifunctor (bimap)
@@ -26,8 +28,9 @@ import Simplex.Messaging.Crypto.SNTRUP761.Bindings.FFI
 import Simplex.Messaging.Crypto.SNTRUP761.Bindings.RNG (rngFuncPtr, withDRG)
 import Simplex.Messaging.Encoding
 import Simplex.Messaging.Encoding.String
+import Simplex.Messaging.Util ((<$?>))
 
-newtype KEMPublicKey = KEMPublicKey ByteString
+newtype KEMPublicKey = KEMPublicKey_ ByteString
   deriving (Eq, Show)
 
 newtype KEMSecretKey = KEMSecretKey ScrubbedBytes
@@ -36,18 +39,24 @@ newtype KEMSecretKey = KEMSecretKey ScrubbedBytes
 newtype KEMCiphertext = KEMCiphertext ByteString
   deriving (Eq, Show)
 
-newtype KEMSharedKey = KEMSharedKey ScrubbedBytes
+newtype KEMSharedKey = KEMSharedKey_ ScrubbedBytes
   deriving (Eq, Show)
 
-unsafeRevealKEMSharedKey :: KEMSharedKey -> String
-unsafeRevealKEMSharedKey (KEMSharedKey scrubbed) = show (BA.convert scrubbed :: ByteString)
-{-# DEPRECATED unsafeRevealKEMSharedKey "unsafeRevealKEMSharedKey left in code" #-}
+pattern KEMPublicKey :: ByteString -> KEMPublicKey
+pattern KEMPublicKey s <- KEMPublicKey_ s
+
+pattern KEMSharedKey :: ScrubbedBytes -> KEMSharedKey
+pattern KEMSharedKey s <- KEMSharedKey_ s
+
+{-# COMPLETE KEMPublicKey #-}
+
+{-# COMPLETE KEMSharedKey #-}
 
 type KEMKeyPair = (KEMPublicKey, KEMSecretKey)
 
 sntrup761Keypair :: TVar ChaChaDRG -> IO KEMKeyPair
 sntrup761Keypair drg =
-  bimap KEMPublicKey KEMSecretKey
+  bimap KEMPublicKey_ KEMSecretKey
     <$> BA.allocRet
       c_SNTRUP761_SECRETKEY_SIZE
       ( \skPtr ->
@@ -56,10 +65,9 @@ sntrup761Keypair drg =
       )
 
 sntrup761Enc :: TVar ChaChaDRG -> KEMPublicKey -> IO (KEMCiphertext, KEMSharedKey)
-sntrup761Enc drg (KEMPublicKey pk) = do
-  requireByteArrayLength "SNTRUP761 public key" c_SNTRUP761_PUBLICKEY_SIZE pk
+sntrup761Enc drg (KEMPublicKey pk) =
   BA.withByteArray pk $ \pkPtr ->
-    bimap KEMCiphertext KEMSharedKey
+    bimap KEMCiphertext KEMSharedKey_
       <$> BA.allocRet
         c_SNTRUP761_SIZE
         ( \kPtr ->
@@ -68,75 +76,50 @@ sntrup761Enc drg (KEMPublicKey pk) = do
         )
 
 sntrup761Dec :: KEMCiphertext -> KEMSecretKey -> IO KEMSharedKey
-sntrup761Dec (KEMCiphertext c) (KEMSecretKey sk) = do
-  requireByteArrayLength "SNTRUP761 ciphertext" c_SNTRUP761_CIPHERTEXT_SIZE c
-  requireByteArrayLength "SNTRUP761 secret key" c_SNTRUP761_SECRETKEY_SIZE sk
+sntrup761Dec (KEMCiphertext c) (KEMSecretKey sk) =
   BA.withByteArray sk $ \skPtr ->
     BA.withByteArray c $ \cPtr ->
-      KEMSharedKey
+      KEMSharedKey_
         <$> BA.alloc c_SNTRUP761_SIZE (\kPtr -> c_sntrup761_dec kPtr cPtr skPtr)
 
-requireByteArrayLength :: BA.ByteArrayAccess bytes => String -> Int -> bytes -> IO ()
-requireByteArrayLength valueName expected bytes =
-  either (throwIO . userError) (const $ pure ()) $
-    validateByteArrayLength valueName expected bytes
-
-validateByteArrayLength :: BA.ByteArrayAccess bytes => String -> Int -> bytes -> Either String bytes
-validateByteArrayLength valueName expected bytes
-  | actual == expected = Right bytes
-  | otherwise = Left $ valueName <> " must be " <> show expected <> " bytes, got " <> show actual
+parseKey :: BA.ByteArrayAccess bs => (bs -> key) -> String -> Int -> bs -> Either String key
+parseKey kCon name expected s
+  | len == expected = Right $ kCon s
+  | otherwise = Left $ name <> " must be " <> show expected <> " bytes, got " <> show len
   where
-    actual = BA.length bytes
-
-parseKEMPublicKey :: ByteString -> Either String KEMPublicKey
-parseKEMPublicKey =
-  fmap KEMPublicKey . validateByteArrayLength "SNTRUP761 public key" c_SNTRUP761_PUBLICKEY_SIZE
-
-parseKEMSecretKey :: ScrubbedBytes -> Either String KEMSecretKey
-parseKEMSecretKey =
-  fmap KEMSecretKey . validateByteArrayLength "SNTRUP761 secret key" c_SNTRUP761_SECRETKEY_SIZE
-
-parseKEMCiphertext :: ByteString -> Either String KEMCiphertext
-parseKEMCiphertext =
-  fmap KEMCiphertext . validateByteArrayLength "SNTRUP761 ciphertext" c_SNTRUP761_CIPHERTEXT_SIZE
+    len = BA.length s    
 
 instance Encoding KEMSecretKey where
   smpEncode (KEMSecretKey c) = smpEncode . Large $ BA.convert c
-  smpP = do
-    Large bytes <- smpP
-    either fail pure $ parseKEMSecretKey (BA.convert bytes)
+  smpP = parseKey KEMSecretKey "SNTRUP761 secret key" c_SNTRUP761_SECRETKEY_SIZE . BA.convert . unLarge <$?> smpP
 
 instance StrEncoding KEMSecretKey where
   strEncode (KEMSecretKey pk) = strEncode (BA.convert pk :: ByteString)
-  strP = either fail pure . parseKEMSecretKey . BA.convert =<< strP @ByteString
+  strP = parseKey KEMSecretKey "SNTRUP761 secret key" c_SNTRUP761_SECRETKEY_SIZE . BA.convert <$?> strP @ByteString
 
 instance Encoding KEMPublicKey where
   smpEncode (KEMPublicKey pk) = smpEncode . Large $ BA.convert pk
-  smpP = do
-    Large bytes <- smpP
-    either fail pure $ parseKEMPublicKey bytes
+  smpP = parseKey KEMPublicKey_ "SNTRUP761 public key" c_SNTRUP761_PUBLICKEY_SIZE . unLarge <$?> smpP
 
 instance StrEncoding KEMPublicKey where
   strEncode (KEMPublicKey pk) = strEncode (BA.convert pk :: ByteString)
-  strP = either fail pure . parseKEMPublicKey =<< strP @ByteString
+  strP = parseKey KEMPublicKey_ "SNTRUP761 public key" c_SNTRUP761_PUBLICKEY_SIZE <$?> strP @ByteString
 
 instance Encoding KEMCiphertext where
   smpEncode (KEMCiphertext c) = smpEncode . Large $ BA.convert c
-  smpP = do
-    Large bytes <- smpP
-    either fail pure $ parseKEMCiphertext bytes
+  smpP = parseKey KEMCiphertext "SNTRUP761 ciphertext" c_SNTRUP761_CIPHERTEXT_SIZE . unLarge <$?> smpP
 
 instance Encoding KEMSharedKey where
   smpEncode (KEMSharedKey c) = smpEncode (BA.convert c :: ByteString)
-  smpP = KEMSharedKey . BA.convert <$> smpP @ByteString
+  smpP = KEMSharedKey_ . BA.convert <$> smpP @ByteString
 
 instance StrEncoding KEMCiphertext where
   strEncode (KEMCiphertext pk) = strEncode (BA.convert pk :: ByteString)
-  strP = either fail pure . parseKEMCiphertext =<< strP @ByteString
+  strP = parseKey KEMCiphertext "SNTRUP761 ciphertext" c_SNTRUP761_CIPHERTEXT_SIZE <$?> strP @ByteString
 
 instance StrEncoding KEMSharedKey where
   strEncode (KEMSharedKey pk) = strEncode (BA.convert pk :: ByteString)
-  strP = KEMSharedKey . BA.convert <$> strP @ByteString
+  strP = KEMSharedKey_ . BA.convert <$> strP @ByteString
 
 instance ToJSON KEMSecretKey where
   toJSON = strToJSON
@@ -164,9 +147,9 @@ instance ToField KEMSharedKey where
 
 instance FromField KEMSharedKey where
 #if defined(dbPostgres)
-  fromField f dat = KEMSharedKey . BA.convert @ByteString <$> fromField f dat
+  fromField f dat = KEMSharedKey_ . BA.convert @ByteString <$> fromField f dat
 #else
-  fromField f = KEMSharedKey . BA.convert @ByteString <$> fromField f
+  fromField f = KEMSharedKey_ . BA.convert @ByteString <$> fromField f
 #endif
 
 instance ToJSON KEMSharedKey where

tests/CoreTests/CryptoTests.hs

@@ -7,14 +7,12 @@
 module CoreTests.CryptoTests (cryptoTests) where
 
 import Control.Concurrent.STM
-import Control.Exception (SomeException)
 import Control.Monad.Except
 import qualified Data.Aeson as J
 import qualified Data.ByteString.Char8 as B
 import qualified Data.ByteString.Lazy.Char8 as LB
 import Data.Either (isLeft, isRight)
 import Data.Int (Int64)
-import Data.List (isInfixOf)
 import qualified Data.Text as T
 import Data.Text.Encoding (encodeUtf8)
 import qualified Data.Text.Lazy as LT
@@ -110,7 +108,6 @@ cryptoTests = do
   describe "sntrup761" $ do
     it "should enc/dec key" testSNTRUP761
     it "should reject malformed KEM encodings" testSNTRUP761RejectsMalformedEncodings
-    it "should reject malformed KEM FFI inputs" testSNTRUP761RejectsMalformedFFIInputs
   describe "BBS+" $ do
     it "should sign and verify" testBBSSignVerify
     it "should derive public key from secret key" testBBSPublicKeyDerivation
@@ -305,18 +302,6 @@ testSNTRUP761RejectsMalformedEncodings = do
   smpDecode @KEMSecretKey (smpEncode $ Large shortSecretKey) `shouldSatisfy` isLeft
   strDecode @KEMSecretKey (strEncode shortSecretKey) `shouldSatisfy` isLeft
 
-testSNTRUP761RejectsMalformedFFIInputs :: IO ()
-testSNTRUP761RejectsMalformedFFIInputs = do
-  drg <- C.newRandom
-  (_, sk) <- sntrup761Keypair drg
-  sntrup761Enc drg (KEMPublicKey shortPublicKey)
-    `shouldThrow` kemLengthException "public key"
-  sntrup761Dec (KEMCiphertext shortCiphertext) sk
-    `shouldThrow` kemLengthException "ciphertext"
-
-kemLengthException :: String -> SomeException -> Bool
-kemLengthException valueName e = valueName `isInfixOf` show e
-
 shortPublicKey :: B.ByteString
 shortPublicKey = B.replicate (c_SNTRUP761_PUBLICKEY_SIZE - 1) 'p'