-- | A primitive interface to SHA-2
module Codec.Digest.SHA.IO(
  Length(..), SHACtx, newCtx, unsafeUpdateCtx, unsafeFinalizeCtx, showBSasHex
) where

import Foreign
import Foreign.C.Types
import qualified Data.ByteString as B
import qualified Data.ByteString.Internal as BI
import Control.Applicative
import Codec.Digest.SHA.Misc

#include "sha2.h"

data Length = SHA256 | SHA384 | SHA512

data SHACtx = Ctx256 (ForeignPtr Ctx256)
            | Ctx384 (ForeignPtr Ctx384)
            | Ctx512 (ForeignPtr Ctx512)

newCtx :: Length -> IO SHACtx
newCtx SHA256 = Ctx256 <$> initialize _sha256Init
newCtx SHA384 = Ctx384 <$> initialize _sha384Init
newCtx SHA512 = Ctx512 <$> initialize _sha512Init

initialize :: Storable a => (Ptr a -> IO ()) -> IO (ForeignPtr a)
initialize f = do
  ctx <- mallocForeignPtr
  withForeignPtr ctx f
  return ctx

-- | Heads up: Calling this function after calling finalizeCtx is
-- likely to result in crashes, therefore unsafe.
unsafeUpdateCtx :: SHACtx -> B.ByteString -> IO ()
unsafeUpdateCtx (Ctx256 p) = update _sha256Update p
unsafeUpdateCtx (Ctx384 p) = update _sha384Update p
unsafeUpdateCtx (Ctx512 p) = update _sha512Update p

update :: (Ptr a -> Ptr Word8 -> CSize -> IO ()) -> ForeignPtr a -> B.ByteString -> IO ()
update f ctx (BI.toForeignPtr -> (bs,offset,(fromIntegral -> len))) =
  withForeignPtr ctx $ \ctxp ->
  withForeignPtr bs $ \bsp ->
  f ctxp (bsp `plusPtr` offset) len

-- | After calling this, you must not call either unsafeUpdateCtx or
-- unsafeFinalizeCtx again. (On the same context)
unsafeFinalizeCtx :: SHACtx -> IO B.ByteString
unsafeFinalizeCtx (Ctx256 p) = finalize _sha256End p 32
unsafeFinalizeCtx (Ctx384 p) = finalize _sha384End p 48
unsafeFinalizeCtx (Ctx512 p) = finalize _sha512End p 64

finalize :: (Ptr a -> Ptr Word8 -> IO (Ptr Word8)) -> ForeignPtr a -> Int -> IO B.ByteString
finalize f ctx len =
  withForeignPtr ctx $ \ctxp ->
  BI.create len $ \bsp ->
  f ctxp bsp >> return ()


data Ctx256
data Ctx384
data Ctx512

instance Storable Ctx256 where
  sizeOf _ = #size SHA256_CTX
  alignment _ = 16

instance Storable Ctx384 where
  sizeOf _ = #size SHA384_CTX
  alignment _ = 16

instance Storable Ctx512 where
  sizeOf _ = #size SHA512_CTX
  alignment _ = 16

foreign import ccall unsafe "SHA256_Init" _sha256Init :: Ptr Ctx256 -> IO ()
foreign import ccall unsafe "SHA256_Update" _sha256Update :: Ptr Ctx256 -> Ptr Word8 -> CSize -> IO ()
foreign import ccall unsafe "SHA256_End" _sha256End :: Ptr Ctx256 -> Ptr Word8 -> IO (Ptr Word8)
foreign import ccall unsafe "SHA384_Init" _sha384Init :: Ptr Ctx384 -> IO ()
foreign import ccall unsafe "SHA384_Update" _sha384Update :: Ptr Ctx384 -> Ptr Word8 -> CSize -> IO ()
foreign import ccall unsafe "SHA384_End" _sha384End :: Ptr Ctx384 -> Ptr Word8 -> IO (Ptr Word8)
foreign import ccall unsafe "SHA512_Init" _sha512Init :: Ptr Ctx512 -> IO ()
foreign import ccall unsafe "SHA512_Update" _sha512Update :: Ptr Ctx512 -> Ptr Word8 -> CSize -> IO ()
foreign import ccall unsafe "SHA512_End" _sha512End :: Ptr Ctx512 -> Ptr Word8 -> IO (Ptr Word8)