Arm.hs (8414B)
1 {-# OPTIONS_HADDOCK hide #-} 2 {-# LANGUAGE BangPatterns #-} 3 {-# LANGUAGE MagicHash #-} 4 {-# LANGUAGE PatternSynonyms #-} 5 {-# LANGUAGE UnboxedTuples #-} 6 7 -- | 8 -- Module: Crypto.Hash.SHA256.Arm 9 -- Copyright: (c) 2024 Jared Tobin 10 -- License: MIT 11 -- Maintainer: Jared Tobin <jared@ppad.tech> 12 -- 13 -- ARM crypto extension support for SHA-256. 14 15 module Crypto.Hash.SHA256.Arm ( 16 sha256_arm_available 17 , hash 18 , hmac 19 , _hmac_rr 20 , _hmac_rsb 21 ) where 22 23 import qualified Data.ByteString as BS 24 import qualified Data.ByteString.Internal as BI 25 import qualified Data.ByteString.Unsafe as BU 26 import Data.Word (Word8, Word32, Word64) 27 import Foreign.Marshal.Alloc (allocaBytes) 28 import Foreign.Ptr (Ptr) 29 import qualified GHC.Exts as Exts 30 import qualified GHC.IO (IO(..)) 31 import qualified GHC.Ptr 32 import Crypto.Hash.SHA256.Internal hiding (update) 33 import System.IO.Unsafe (unsafeDupablePerformIO) 34 35 -- ffi ------------------------------------------------------------------------ 36 37 foreign import ccall unsafe "sha256_block_arm" 38 c_sha256_block :: Ptr Word32 -> Ptr Word32 -> IO () 39 40 foreign import ccall unsafe "sha256_arm_available" 41 c_sha256_arm_available :: IO Int 42 43 -- utilities ------------------------------------------------------------------ 44 45 fi :: (Integral a, Num b) => a -> b 46 fi = fromIntegral 47 {-# INLINE fi #-} 48 49 50 peek_registers 51 :: Ptr Word32 52 -> Registers 53 peek_registers (GHC.Ptr.Ptr addr) = R 54 (Exts.indexWord32OffAddr# addr 0#) 55 (Exts.indexWord32OffAddr# addr 1#) 56 (Exts.indexWord32OffAddr# addr 2#) 57 (Exts.indexWord32OffAddr# addr 3#) 58 (Exts.indexWord32OffAddr# addr 4#) 59 (Exts.indexWord32OffAddr# addr 5#) 60 (Exts.indexWord32OffAddr# addr 6#) 61 (Exts.indexWord32OffAddr# addr 7#) 62 {-# INLINE peek_registers #-} 63 64 poke_block :: Ptr Word32 -> Block -> IO () 65 poke_block 66 (GHC.Ptr.Ptr addr) 67 (B w00 w01 w02 w03 w04 w05 w06 w07 w08 w09 w10 w11 w12 w13 w14 w15) 68 = GHC.IO.IO $ \s00 -> 69 case Exts.writeWord32OffAddr# addr 00# w00 s00 of { s01 -> 70 case Exts.writeWord32OffAddr# addr 01# w01 s01 of { s02 -> 71 case Exts.writeWord32OffAddr# addr 02# w02 s02 of { s03 -> 72 case Exts.writeWord32OffAddr# addr 03# w03 s03 of { s04 -> 73 case Exts.writeWord32OffAddr# addr 04# w04 s04 of { s05 -> 74 case Exts.writeWord32OffAddr# addr 05# w05 s05 of { s06 -> 75 case Exts.writeWord32OffAddr# addr 06# w06 s06 of { s07 -> 76 case Exts.writeWord32OffAddr# addr 07# w07 s07 of { s08 -> 77 case Exts.writeWord32OffAddr# addr 08# w08 s08 of { s09 -> 78 case Exts.writeWord32OffAddr# addr 09# w09 s09 of { s10 -> 79 case Exts.writeWord32OffAddr# addr 10# w10 s10 of { s11 -> 80 case Exts.writeWord32OffAddr# addr 11# w11 s11 of { s12 -> 81 case Exts.writeWord32OffAddr# addr 12# w12 s12 of { s13 -> 82 case Exts.writeWord32OffAddr# addr 13# w13 s13 of { s14 -> 83 case Exts.writeWord32OffAddr# addr 14# w14 s14 of { s15 -> 84 case Exts.writeWord32OffAddr# addr 15# w15 s15 of { s16 -> 85 (# s16, () #) }}}}}}}}}}}}}}}} 86 {-# INLINE poke_block #-} 87 88 -- update --------------------------------------------------------------------- 89 90 update :: Ptr Word32 -> Ptr Word32 -> Block -> IO () 91 update rp bp block = do 92 poke_block bp block 93 c_sha256_block rp bp 94 {-# INLINE update #-} 95 96 -- api ----------------------------------------------------------------------- 97 98 -- | Are ARM +sha2 extensions available? 99 sha256_arm_available :: Bool 100 sha256_arm_available = unsafeDupablePerformIO c_sha256_arm_available /= 0 101 {-# NOINLINE sha256_arm_available #-} 102 103 hash 104 :: BS.ByteString 105 -> BS.ByteString 106 hash m = unsafeDupablePerformIO $ 107 allocaBytes 32 $ \rp -> 108 allocaBytes 64 $ \bp -> do 109 poke_registers rp (iv ()) 110 _hash rp bp 0 m 111 let !rs = peek_registers rp 112 pure (cat rs) 113 114 _hash 115 :: Ptr Word32 -- ^ register state 116 -> Ptr Word32 -- ^ block state 117 -> Word64 -- ^ extra prefix length, for padding calculation 118 -> BS.ByteString -- ^ input 119 -> IO () 120 _hash rp bp el m@(BI.PS _ _ l) = do 121 hash_blocks rp bp m 122 let !fin@(BI.PS _ _ ll) = BU.unsafeDrop (l - l `rem` 64) m 123 !total = el + fi l 124 if ll < 56 125 then do 126 let !ult = parse_pad1 fin total 127 update rp bp ult 128 else do 129 let !(# pen, ult #) = parse_pad2 fin total 130 update rp bp pen 131 update rp bp ult 132 {-# INLINABLE _hash #-} 133 134 hash_blocks 135 :: Ptr Word32 -- ^ register state 136 -> Ptr Word32 -- ^ block state 137 -> BS.ByteString -- ^ input 138 -> IO () 139 hash_blocks rp bp m@(BI.PS _ _ l) = loop 0 where 140 loop !j 141 | j + 64 > l = pure () 142 | otherwise = do 143 let !block = parse m j 144 update rp bp block 145 loop (j + 64) 146 {-# INLINE hash_blocks #-} 147 148 -- hmac ------------------------------------------------------------------------ 149 150 hmac :: BS.ByteString -> BS.ByteString -> BS.ByteString 151 hmac k m = unsafeDupablePerformIO $ 152 allocaBytes 32 $ \rp -> 153 allocaBytes 64 $ \bp -> do 154 _hmac rp bp (prep_key k) m 155 pure (cat (peek_registers rp)) 156 157 prep_key :: BS.ByteString -> Block 158 prep_key k@(BI.PS _ _ l) 159 | l > 64 = parse_key (hash k) 160 | otherwise = parse_key k 161 {-# INLINABLE prep_key #-} 162 163 -- assume padded key as block. 164 _hmac 165 :: Ptr Word32 -- ^ register state 166 -> Ptr Word32 -- ^ block state 167 -> Block -- ^ padded key 168 -> BS.ByteString -- ^ message 169 -> IO () 170 _hmac rp bp k m = do 171 poke_registers rp (iv ()) 172 update rp bp (xor k (Exts.wordToWord32# 0x36363636##)) 173 _hash rp bp 64 m 174 let !block = pad_registers_with_length (peek_registers rp) 175 poke_registers rp (iv ()) 176 update rp bp (xor k (Exts.wordToWord32# 0x5C5C5C5C##)) 177 update rp bp block 178 {-# NOINLINE _hmac #-} 179 180 _hmac_rr 181 :: Ptr Word32 -- ^ register state 182 -> Ptr Word32 -- ^ block state 183 -> Registers -- ^ key 184 -> Registers -- ^ message 185 -> IO () 186 _hmac_rr rp bp k m = do 187 let !key = pad_registers k 188 !block = pad_registers_with_length m 189 _hmac_bb rp bp key block 190 {-# INLINABLE _hmac_rr #-} 191 192 _hmac_bb 193 :: Ptr Word32 -- ^ register state 194 -> Ptr Word32 -- ^ block state 195 -> Block -- ^ padded key 196 -> Block -- ^ padded message 197 -> IO () 198 _hmac_bb rp bp k m = do 199 poke_registers rp (iv ()) 200 update rp bp (xor k (Exts.wordToWord32# 0x36363636##)) 201 update rp bp m 202 let !inner = pad_registers_with_length (peek_registers rp) 203 poke_registers rp (iv ()) 204 update rp bp (xor k (Exts.wordToWord32# 0x5C5C5C5C##)) 205 update rp bp inner 206 {-# INLINABLE _hmac_bb #-} 207 208 -- | HMAC(key, v || sep || data) using ARM crypto extensions. 209 -- Writes result to destination pointer. 210 _hmac_rsb 211 :: Ptr Word32 -- ^ destination (8 Word32s) 212 -> Ptr Word32 -- ^ scratch block buffer (16 Word32s) 213 -> Registers -- ^ key 214 -> Registers -- ^ v 215 -> Word8 -- ^ separator byte 216 -> BS.ByteString -- ^ data 217 -> IO () 218 _hmac_rsb rp bp k v sep dat = do 219 poke_registers rp (iv ()) 220 let !key = pad_registers k 221 update rp bp (xor key (Exts.wordToWord32# 0x36363636##)) 222 _hash_vsb rp bp 64 v sep dat 223 let !inner = pad_registers_with_length (peek_registers rp) 224 poke_registers rp (iv ()) 225 update rp bp (xor key (Exts.wordToWord32# 0x5C5C5C5C##)) 226 update rp bp inner 227 {-# INLINABLE _hmac_rsb #-} 228 229 -- | Hash (v || sep || dat) with ARM crypto extensions. 230 -- Assumes register state already initialized at rp. 231 _hash_vsb 232 :: Ptr Word32 -- ^ register state 233 -> Ptr Word32 -- ^ block buffer 234 -> Word64 -- ^ extra prefix length 235 -> Registers -- ^ v 236 -> Word8 -- ^ sep 237 -> BS.ByteString -- ^ dat 238 -> IO () 239 _hash_vsb rp bp el v sep dat@(BI.PS _ _ l) 240 | l >= 31 = do 241 -- first block is complete: v || sep || dat[0:31] 242 let !b0 = parse_vsb v sep dat 243 update rp bp b0 244 -- hash remaining complete blocks from dat[31:] 245 let !rest = BU.unsafeDrop 31 dat 246 !restLen = l - 31 247 hash_blocks rp bp rest 248 -- handle final padding 249 let !finLen = restLen `rem` 64 250 !fin = BU.unsafeDrop (restLen - finLen) rest 251 !total = el + 33 + fi l 252 if finLen < 56 253 then update rp bp (parse_pad1 fin total) 254 else do 255 let !(# pen, ult #) = parse_pad2 fin total 256 update rp bp pen 257 update rp bp ult 258 | otherwise = do 259 -- message < 64 bytes total, straight to padding 260 let !total = el + 33 + fi l 261 if 33 + l < 56 262 then update rp bp (parse_pad1_vsb v sep dat total) 263 else do 264 let !(# pen, ult #) = parse_pad2_vsb v sep dat total 265 update rp bp pen 266 update rp bp ult 267 {-# INLINABLE _hash_vsb #-}