{-# OPTIONS_HADDOCK prune #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE UnliftedNewtypes #-}

-- |
-- Module: Crypto.Hash.SHA256
-- Copyright: (c) 2024 Jared Tobin
-- License: MIT
-- Maintainer: Jared Tobin <jared@ppad.tech>
--
-- SHA-256 and HMAC-SHA256 implementations for
-- strict and lazy ByteStrings, as specified by RFC's
-- [6234](https://datatracker.ietf.org/doc/html/rfc6234) and
-- [2104](https://datatracker.ietf.org/doc/html/rfc2104).
--
-- The 'hash' and 'hmac' functions will use primitive instructions from
-- the ARM cryptographic extensions via FFI if they're available, and
-- will otherwise use a pure Haskell implementation.

module Crypto.Hash.SHA256 (
  -- * SHA-256 message digest functions
    hash
  , Lazy.hash_lazy

  -- * SHA256-based MAC functions
  , MAC(..)
  , hmac
  , Lazy.hmac_lazy

  -- low-level specialized HMAC primitives
  , _hmac_rr
  , _hmac_rsb
  ) where

import qualified Data.ByteString as BS
import qualified Data.ByteString.Internal as BI
import qualified Data.ByteString.Unsafe as BU
import Data.Word (Word8, Word32, Word64)
import Foreign.Ptr (Ptr)
import qualified GHC.Exts as Exts
import qualified Crypto.Hash.SHA256.Arm as Arm
import Crypto.Hash.SHA256.Internal
import qualified Crypto.Hash.SHA256.Lazy as Lazy

-- utilities ------------------------------------------------------------------

fi :: (Integral a, Num b) => a -> b
fi :: forall a b. (Integral a, Num b) => a -> b
fi = a -> b
forall a b. (Integral a, Num b) => a -> b
fromIntegral
{-# INLINE fi #-}

-- hash -----------------------------------------------------------------------

-- | Compute a condensed representation of a strict bytestring via
--   SHA-256.
--
--   The 256-bit output digest is returned as a strict bytestring.
--
--   >>> hash "strict bytestring input"
--   "<strict 256-bit message digest>"
hash :: BS.ByteString -> BS.ByteString
hash :: ByteString -> ByteString
hash ByteString
m
  | Bool
Arm.sha256_arm_available = ByteString -> ByteString
Arm.hash ByteString
m
  | Bool
otherwise = Registers -> ByteString
cat (Word64 -> Registers -> ByteString -> Registers
_hash Word64
0 (() -> Registers
iv ()) ByteString
m)
{-# INLINABLE hash #-}

_hash
  :: Word64        -- ^ extra prefix length for padding calculations
  -> Registers     -- ^ register state
  -> BS.ByteString -- ^ input
  -> Registers
_hash :: Word64 -> Registers -> ByteString -> Registers
_hash Word64
el Registers
rs m :: ByteString
m@(BI.PS ForeignPtr Word8
_ Int
_ Int
l) = do
  let !state :: Registers
state = Registers -> ByteString -> Registers
_hash_blocks Registers
rs ByteString
m
      !fin :: ByteString
fin@(BI.PS ForeignPtr Word8
_ Int
_ Int
ll) = Int -> ByteString -> ByteString
BU.unsafeDrop (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l Int -> Int -> Int
forall a. Integral a => a -> a -> a
`rem` Int
64) ByteString
m
      !total :: Word64
total = Word64
el Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fi Int
l
  if   Int
ll Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
56
  then
    let !ult :: Block
ult = ByteString -> Word64 -> Block
parse_pad1 ByteString
fin Word64
total
    in  Registers -> Block -> Registers
update Registers
state Block
ult
  else
    let !(# Block
pen, Block
ult #) = ByteString -> Word64 -> (# Block, Block #)
parse_pad2 ByteString
fin Word64
total
    in  Registers -> Block -> Registers
update (Registers -> Block -> Registers
update Registers
state Block
pen) Block
ult
{-# INLINABLE _hash #-}

_hash_blocks
  :: Registers     -- ^ state
  -> BS.ByteString -- ^ input
  -> Registers
_hash_blocks :: Registers -> ByteString -> Registers
_hash_blocks Registers
rs m :: ByteString
m@(BI.PS ForeignPtr Word8
_ Int
_ Int
l) = Registers -> Int -> Registers
loop Registers
rs Int
0 where
  loop :: Registers -> Int -> Registers
loop !Registers
acc !Int
j
    | Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
64 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
l = Registers
acc
    | Bool
otherwise  =
        let !nacc :: Registers
nacc = Registers -> Block -> Registers
update Registers
acc (ByteString -> Int -> Block
parse ByteString
m Int
j)
        in  Registers -> Int -> Registers
loop Registers
nacc (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
64)
{-# INLINABLE _hash_blocks #-}

-- hmac ----------------------------------------------------------------------

-- | Compute a condensed representation of a strict bytestring via
--   SHA-256.
--
--   The 256-bit output digest is returned as a strict bytestring.
--
--   >>> hash "strict bytestring input"
--   "<strict 256-bit message digest>"
hmac :: BS.ByteString -> BS.ByteString -> MAC
hmac :: ByteString -> ByteString -> MAC
hmac ByteString
k ByteString
m
  | Bool
Arm.sha256_arm_available = ByteString -> MAC
MAC (ByteString -> ByteString -> ByteString
Arm.hmac ByteString
k ByteString
m)
  | Bool
otherwise = ByteString -> MAC
MAC (Registers -> ByteString
cat (Block -> ByteString -> Registers
_hmac (ByteString -> Block
prep_key ByteString
k) ByteString
m))
{-# INLINABLE hmac #-}

prep_key :: BS.ByteString -> Block
prep_key :: ByteString -> Block
prep_key k :: ByteString
k@(BI.PS ForeignPtr Word8
_ Int
_ Int
l)
    | Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
64    = ByteString -> Block
parse_key (ByteString -> ByteString
hash ByteString
k)
    | Bool
otherwise = ByteString -> Block
parse_key ByteString
k
{-# INLINABLE prep_key #-}

_hmac
  :: Block          -- ^ padded key
  -> BS.ByteString  -- ^ message
  -> Registers
_hmac :: Block -> ByteString -> Registers
_hmac Block
k ByteString
m =
  let !rs0 :: Registers
rs0   = Registers -> Block -> Registers
update (() -> Registers
iv ()) (Block -> Word32# -> Block
xor Block
k (Word# -> Word32#
Exts.wordToWord32# Word#
0x36363636##))
      !block :: Block
block = Registers -> Block
pad_registers_with_length (Word64 -> Registers -> ByteString -> Registers
_hash Word64
64 Registers
rs0 ByteString
m)
      !rs1 :: Registers
rs1   = Registers -> Block -> Registers
update (() -> Registers
iv ()) (Block -> Word32# -> Block
xor Block
k (Word# -> Word32#
Exts.wordToWord32# Word#
0x5C5C5C5C##))
  in  Registers -> Block -> Registers
update Registers
rs1 Block
block
{-# INLINABLE _hmac #-}

-- the following functions are useful when we want to avoid allocating certain
-- components of the HMAC key and message on the heap.

-- Computes hmac(k, v) when k and v are Registers.
--
-- The 32-byte result is written to the destination pointer.
_hmac_rr
  :: Ptr Word32    -- ^ destination (8 Word32s)
  -> Ptr Word32    -- ^ scratch block buffer (16 Word32s)
  -> Registers     -- ^ key
  -> Registers     -- ^ message
  -> IO ()
_hmac_rr :: Ptr Word32 -> Ptr Word32 -> Registers -> Registers -> IO ()
_hmac_rr Ptr Word32
rp Ptr Word32
bp Registers
k Registers
m
  | Bool
Arm.sha256_arm_available = Ptr Word32 -> Ptr Word32 -> Registers -> Registers -> IO ()
Arm._hmac_rr Ptr Word32
rp Ptr Word32
bp Registers
k Registers
m
  | Bool
otherwise = do
      let !key :: Block
key   = Registers -> Block
pad_registers Registers
k
          !block :: Block
block = Registers -> Block
pad_registers_with_length Registers
m
          !rs :: Registers
rs    = Block -> Block -> Registers
_hmac_bb Block
key Block
block
      Ptr Word32 -> Registers -> IO ()
poke_registers Ptr Word32
rp Registers
rs
{-# INLINABLE _hmac_rr #-}

_hmac_bb
  :: Block     -- ^ key
  -> Block     -- ^ message
  -> Registers
_hmac_bb :: Block -> Block -> Registers
_hmac_bb Block
k Block
m =
  let !rs0 :: Registers
rs0   = Registers -> Block -> Registers
update (() -> Registers
iv ()) (Block -> Word32# -> Block
xor Block
k (Word# -> Word32#
Exts.wordToWord32# Word#
0x36363636##))
      !rs1 :: Registers
rs1   = Registers -> Block -> Registers
update Registers
rs0 Block
m
      !inner :: Block
inner = Registers -> Block
pad_registers_with_length Registers
rs1
      !rs2 :: Registers
rs2   = Registers -> Block -> Registers
update (() -> Registers
iv ()) (Block -> Word32# -> Block
xor Block
k (Word# -> Word32#
Exts.wordToWord32# Word#
0x5C5C5C5C##))
  in  Registers -> Block -> Registers
update Registers
rs2 Block
inner
{-# INLINABLE _hmac_bb #-}

-- Calculate hmac(k, m) where m is the concatenation of v (registers), a
-- separator byte, and a ByteString. This avoids allocating 'v' on the
-- heap.
--
-- The 32-byte result is written to the destination pointer.
_hmac_rsb
  :: Ptr Word32    -- ^ destination pointer (8 x Word32)
  -> Ptr Word32    -- ^ scratch block pointer (16 x Word32)
  -> Registers     -- ^ k
  -> Registers     -- ^ v
  -> Word8         -- ^ separator byte
  -> BS.ByteString -- ^ data
  -> IO ()
_hmac_rsb :: Ptr Word32
-> Ptr Word32
-> Registers
-> Registers
-> Word8
-> ByteString
-> IO ()
_hmac_rsb Ptr Word32
rp Ptr Word32
bp Registers
k Registers
v Word8
sep ByteString
dat
  | Bool
Arm.sha256_arm_available = Ptr Word32
-> Ptr Word32
-> Registers
-> Registers
-> Word8
-> ByteString
-> IO ()
Arm._hmac_rsb Ptr Word32
rp Ptr Word32
bp Registers
k Registers
v Word8
sep ByteString
dat
  | Bool
otherwise = do
      let !key :: Block
key   = Registers -> Block
pad_registers Registers
k
          !rs0 :: Registers
rs0   = Registers -> Block -> Registers
update (() -> Registers
iv ()) (Block -> Word32# -> Block
xor Block
key (Word# -> Word32#
Exts.wordToWord32# Word#
0x36363636##))
          !inner :: Registers
inner = Word64
-> Registers -> Registers -> Word8 -> ByteString -> Registers
_hash_vsb Word64
64 Registers
rs0 Registers
v Word8
sep ByteString
dat
          !block :: Block
block = Registers -> Block
pad_registers_with_length Registers
inner
          !rs1 :: Registers
rs1   = Registers -> Block -> Registers
update (() -> Registers
iv ()) (Block -> Word32# -> Block
xor Block
key (Word# -> Word32#
Exts.wordToWord32# Word#
0x5C5C5C5C##))
          !rs :: Registers
rs    = Registers -> Block -> Registers
update Registers
rs1 Block
block
      Ptr Word32 -> Registers -> IO ()
poke_registers Ptr Word32
rp Registers
rs
{-# INLINABLE _hmac_rsb #-}

-- hash(v || sep || dat) with a custom initial state and extra
-- prefix length. used for producing a more specialized hmac.
_hash_vsb
  :: Word64        -- ^ extra prefix length
  -> Registers     -- ^ initial state
  -> Registers     -- ^ v
  -> Word8         -- ^ sep
  -> BS.ByteString -- ^ dat
  -> Registers
_hash_vsb :: Word64
-> Registers -> Registers -> Word8 -> ByteString -> Registers
_hash_vsb Word64
el Registers
rs0 Registers
v Word8
sep dat :: ByteString
dat@(BI.PS ForeignPtr Word8
_ Int
_ Int
l)
  | Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
31 =
      -- first block is complete
      let !b0 :: Block
b0    = Registers -> Word8 -> ByteString -> Block
parse_vsb Registers
v Word8
sep ByteString
dat
          !rs1 :: Registers
rs1   = Registers -> Block -> Registers
update Registers
rs0 Block
b0
          !rest :: ByteString
rest  = Int -> ByteString -> ByteString
BU.unsafeDrop Int
31 ByteString
dat
          !rlen :: Int
rlen  = Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
31
          !rs2 :: Registers
rs2   = Registers -> ByteString -> Registers
_hash_blocks Registers
rs1 ByteString
rest
          !flen :: Int
flen  = Int
rlen Int -> Int -> Int
forall a. Integral a => a -> a -> a
`rem` Int
64
          !fin :: ByteString
fin   = Int -> ByteString -> ByteString
BU.unsafeDrop (Int
rlen Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
flen) ByteString
rest
          !total :: Word64
total = Word64
el Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Word64
33 Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fi Int
l
      in  if   Int
flen Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
56
          then Registers -> Block -> Registers
update Registers
rs2 (ByteString -> Word64 -> Block
parse_pad1 ByteString
fin Word64
total)
          else let !(# Block
pen, Block
ult #) = ByteString -> Word64 -> (# Block, Block #)
parse_pad2 ByteString
fin Word64
total
               in  Registers -> Block -> Registers
update (Registers -> Block -> Registers
update Registers
rs2 Block
pen) Block
ult
  | Bool
otherwise =
      -- message < 64 bytes, goes straight to padding
      let !total :: Word64
total = Word64
el Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Word64
33 Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fi Int
l
      in  if   Int
33 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
56
          then Registers -> Block -> Registers
update Registers
rs0 (Registers -> Word8 -> ByteString -> Word64 -> Block
parse_pad1_vsb Registers
v Word8
sep ByteString
dat Word64
total)
          else let !(# Block
pen, Block
ult #) = Registers -> Word8 -> ByteString -> Word64 -> (# Block, Block #)
parse_pad2_vsb Registers
v Word8
sep ByteString
dat Word64
total
               in  Registers -> Block -> Registers
update (Registers -> Block -> Registers
update Registers
rs0 Block
pen) Block
ult
{-# INLINABLE _hash_vsb #-}