<commit_before>{-# LANGUAGE DeriveDataTypeable #-}
module System.IO.Streams.SHA where

import Control.Exception (Exception, throwIO)
import Data.Binary.Get
import Data.Typeable (Typeable)
import System.IO.Streams.Internal (InputStream (..))
import qualified System.IO.Streams as S

import Data.ByteString (ByteString)
import qualified Data.ByteString.Char8 as C
import Data.IORef (IORef, newIORef, readIORef, writeIORef)

import Data.Digest.Pure.SHA

sha1Input :: InputStream ByteString -> IO (InputStream ByteString, IO (Digest SHA1State))
sha1Input = shaInput sha1Incremental completeSha1Incremental

sha224Input :: InputStream ByteString -> IO (InputStream ByteString, IO (Digest SHA256State))
sha224Input = shaInput sha224Incremental completeSha224Incremental

sha256Input :: InputStream ByteString -> IO (InputStream ByteString, IO (Digest SHA256State))
sha256Input = shaInput sha256Incremental completeSha256Incremental

sha384Input :: InputStream ByteString -> IO (InputStream ByteString, IO (Digest SHA512State))
sha384Input = shaInput sha384Incremental completeSha384Incremental

sha512Input :: InputStream ByteString -> IO (InputStream ByteString, IO (Digest SHA512State))
sha512Input = shaInput sha512Incremental completeSha512Incremental

checkedSha1Input :: Digest SHA1State -> InputStream ByteString -> IO (InputStream ByteString)
checkedSha1Input = checkedShaInput sha1Incremental completeSha1Incremental

checkedSha224Input :: Digest SHA256State -> InputStream ByteString -> IO (InputStream ByteString)
checkedSha224Input = checkedShaInput sha224Incremental completeSha224Incremental

checkedSha256Input :: Digest SHA256State -> InputStream ByteString -> IO (InputStream ByteString)
checkedSha256Input = checkedShaInput sha256Incremental completeSha256Incremental

checkedSha384Input :: Digest SHA512State -> InputStream ByteString -> IO (InputStream ByteString)
checkedSha384Input = checkedShaInput sha384Incremental completeSha384Incremental

checkedSha512Input :: Digest SHA512State -> InputStream ByteString -> IO (InputStream ByteString)
checkedSha512Input = checkedShaInput sha512Incremental completeSha512Incremental

-- | Inspired by `S.countInput`. The returned IO action can be run only
-- when the input stream is exhausted, otherwise an error occurs.
shaInput :: Decoder a -> (Decoder a -> Int -> Digest a)
  -> InputStream ByteString -> IO (InputStream ByteString, IO (Digest a))
shaInput increment end is = do
  ref <- newIORef (increment, 0)
  is' <- S.makeInputStream $ prod ref
  return $! (is', readIORef ref >>= uncurry complete)

  where

  prod ref = do
    mbs <- S.read is
    maybe
      (return Nothing)
      (\bs -> (modifyRef ref (uncurry $ modify bs)) >> (return $! Just bs))
      mbs

  complete decoder c = return $! end decoder c
  modify bs decoder c = (pushChunk decoder bs, c + (fromIntegral $ C.length bs))

-- | This returns an input stream exactly as the one being wrapped, but throws
-- an error if the computed SHA hash does not match the one given.
checkedShaInput :: Decoder a -> (Decoder a -> Int -> Digest a)
  -> Digest a -> InputStream ByteString -> IO (InputStream ByteString)
checkedShaInput increment end digest is = do
  ref <- newIORef (increment, 0)
  is' <- S.makeInputStream $ prod ref
  return $! is'

  where

  prod ref = do
    mbs <- S.read is
    maybe
      (do r <- readIORef ref
          digest' <- uncurry complete r
          if digest == digest'
            then return Nothing
            else throwIO UnmatchedSHAException)
      (\bs -> (modifyRef ref (uncurry $ modify bs)) >> (return $! Just bs))
      mbs

  complete decoder c = return $! end decoder c
  modify bs decoder c = (pushChunk decoder bs, c + (fromIntegral $ C.length bs))

-- | Taken from System.IO.Streams.ByteString.
{-# INLINE modifyRef #-}
modifyRef :: IORef a -> (a -> a) -> IO ()
modifyRef ref f = do
    x <- readIORef ref
    writeIORef ref $! f x

-- | Exception raised by `checkedShaInput`.
data UnmatchedSHAException = UnmatchedSHAException
  deriving (Typeable)

instance Show UnmatchedSHAException where
    show _ = "Unmatched SHA exception."

instance Exception UnmatchedSHAException
<commit_msg>Checked SHA streams use a String instead of Digest<commit_after>{-# LANGUAGE DeriveDataTypeable #-}
module System.IO.Streams.SHA where

import Control.Exception (Exception, throwIO)
import Data.Binary.Get
import Data.Typeable (Typeable)
import System.IO.Streams.Internal (InputStream (..))
import qualified System.IO.Streams as S

import Data.ByteString (ByteString)
import qualified Data.ByteString.Char8 as C
import Data.IORef (IORef, newIORef, readIORef, writeIORef)

import Data.Digest.Pure.SHA

sha1Input :: InputStream ByteString -> IO (InputStream ByteString, IO (Digest SHA1State))
sha1Input = shaInput sha1Incremental completeSha1Incremental

sha224Input :: InputStream ByteString -> IO (InputStream ByteString, IO (Digest SHA256State))
sha224Input = shaInput sha224Incremental completeSha224Incremental

sha256Input :: InputStream ByteString -> IO (InputStream ByteString, IO (Digest SHA256State))
sha256Input = shaInput sha256Incremental completeSha256Incremental

sha384Input :: InputStream ByteString -> IO (InputStream ByteString, IO (Digest SHA512State))
sha384Input = shaInput sha384Incremental completeSha384Incremental

sha512Input :: InputStream ByteString -> IO (InputStream ByteString, IO (Digest SHA512State))
sha512Input = shaInput sha512Incremental completeSha512Incremental

checkedSha1Input :: Digest SHA1State -> InputStream ByteString -> IO (InputStream ByteString)
checkedSha1Input = checkedShaInput sha1Incremental completeSha1Incremental . showDigest

checkedSha224Input :: Digest SHA256State -> InputStream ByteString -> IO (InputStream ByteString)
checkedSha224Input = checkedShaInput sha224Incremental completeSha224Incremental . showDigest

checkedSha256Input :: Digest SHA256State -> InputStream ByteString -> IO (InputStream ByteString)
checkedSha256Input = checkedShaInput sha256Incremental completeSha256Incremental . showDigest

checkedSha384Input :: Digest SHA512State -> InputStream ByteString -> IO (InputStream ByteString)
checkedSha384Input = checkedShaInput sha384Incremental completeSha384Incremental . showDigest

checkedSha512Input :: Digest SHA512State -> InputStream ByteString -> IO (InputStream ByteString)
checkedSha512Input = checkedShaInput sha512Incremental completeSha512Incremental . showDigest

checkedSha1Input' :: String -> InputStream ByteString -> IO (InputStream ByteString)
checkedSha1Input' = checkedShaInput sha1Incremental completeSha1Incremental

checkedSha224Input' :: String -> InputStream ByteString -> IO (InputStream ByteString)
checkedSha224Input' = checkedShaInput sha224Incremental completeSha224Incremental

checkedSha256Input' :: String -> InputStream ByteString -> IO (InputStream ByteString)
checkedSha256Input' = checkedShaInput sha256Incremental completeSha256Incremental

checkedSha384Input' :: String -> InputStream ByteString -> IO (InputStream ByteString)
checkedSha384Input' = checkedShaInput sha384Incremental completeSha384Incremental

checkedSha512Input' :: String -> InputStream ByteString -> IO (InputStream ByteString)
checkedSha512Input' = checkedShaInput sha512Incremental completeSha512Incremental

-- | Inspired by `S.countInput`. The returned IO action can be run only
-- when the input stream is exhausted, otherwise an error occurs.
shaInput :: Decoder a -> (Decoder a -> Int -> Digest a)
  -> InputStream ByteString -> IO (InputStream ByteString, IO (Digest a))
shaInput increment end is = do
  ref <- newIORef (increment, 0)
  is' <- S.makeInputStream $ prod ref
  return $! (is', readIORef ref >>= uncurry complete)

  where

  prod ref = do
    mbs <- S.read is
    maybe
      (return Nothing)
      (\bs -> (modifyRef ref (uncurry $ modify bs)) >> (return $! Just bs))
      mbs

  complete decoder c = return $! end decoder c
  modify bs decoder c = (pushChunk decoder bs, c + (fromIntegral $ C.length bs))

-- | This returns an input stream exactly as the one being wrapped, but throws
-- an error if the computed SHA hash does not match the one given.
checkedShaInput :: Decoder a -> (Decoder a -> Int -> Digest a)
  -> String -> InputStream ByteString -> IO (InputStream ByteString)
checkedShaInput increment end digest is = do
  ref <- newIORef (increment, 0)
  is' <- S.makeInputStream $ prod ref
  return $! is'

  where

  prod ref = do
    mbs <- S.read is
    maybe
      (do r <- readIORef ref
          digest' <- uncurry complete r
          if digest == showDigest digest'
            then return Nothing
            else throwIO UnmatchedSHAException)
      (\bs -> (modifyRef ref (uncurry $ modify bs)) >> (return $! Just bs))
      mbs

  complete decoder c = return $! end decoder c
  modify bs decoder c = (pushChunk decoder bs, c + (fromIntegral $ C.length bs))

-- | Taken from System.IO.Streams.ByteString.
{-# INLINE modifyRef #-}
modifyRef :: IORef a -> (a -> a) -> IO ()
modifyRef ref f = do
    x <- readIORef ref
    writeIORef ref $! f x

-- | Exception raised by `checkedShaInput`.
data UnmatchedSHAException = UnmatchedSHAException
  deriving (Typeable)

instance Show UnmatchedSHAException where
    show _ = "Unmatched SHA exception."

instance Exception UnmatchedSHAException

{'fraction_non_alphanumeric': 0.08167223359551773, 'fraction_numerical': 0.034263549186510076, 'mean_word_length': 4.5381861575179, 'pattern_counts': {'":': 0, '<': 21, '<?xml version=': 0, '>': 87, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2815655', 'n_tokens_mistral': 2945, 'n_tokens_neox': 2614, 'n_words': 888}