Implement SHAKE output not divisible by 8 bits

Crypto/Hash/SHAKE.hs

@@ -20,19 +20,21 @@ module Crypto.Hash.SHAKE
     (  SHAKE128 (..), SHAKE256 (..)
     ) where
 
+import           Control.Monad (when)
 import           Crypto.Hash.Types
-import           Foreign.Ptr (Ptr)
+import           Foreign.Ptr (Ptr, castPtr)
+import           Foreign.Storable (Storable(..))
+import           Data.Bits
 import           Data.Data
 import           Data.Typeable
 import           Data.Word (Word8, Word32)
 
 import           Data.Proxy (Proxy(..))
-import           GHC.TypeLits (Nat, KnownNat, natVal)
+import           GHC.TypeLits (Nat, KnownNat, type (+))
 import           Crypto.Internal.Nat
 
 -- | SHAKE128 (128 bits) extendable output function.  Supports an arbitrary
--- digest size (multiple of 8 bits), to be specified as a type parameter
--- of kind 'Nat'.
+-- digest size, to be specified as a type parameter of kind 'Nat'.
 --
 -- Note: outputs from @'SHAKE128' n@ and @'SHAKE128' m@ for the same input are
 -- correlated (one being a prefix of the other).  Results are unrelated to
@@ -40,9 +42,9 @@ import           Crypto.Internal.Nat
 data SHAKE128 (bitlen :: Nat) = SHAKE128
     deriving (Show, Data, Typeable)
 
-instance (IsDivisibleBy8 bitlen, KnownNat bitlen) => HashAlgorithm (SHAKE128 bitlen) where
+instance KnownNat bitlen => HashAlgorithm (SHAKE128 bitlen) where
     type HashBlockSize           (SHAKE128 bitlen)  = 168
-    type HashDigestSize          (SHAKE128 bitlen) = Div8 bitlen
+    type HashDigestSize          (SHAKE128 bitlen) = Div8 (bitlen + 7)
     type HashInternalContextSize (SHAKE128 bitlen) = 376
     hashBlockSize  _          = 168
     hashDigestSize _          = byteLen (Proxy :: Proxy bitlen)
@@ -52,8 +54,7 @@ instance (IsDivisibleBy8 bitlen, KnownNat bitlen) => HashAlgorithm (SHAKE128 bit
     hashInternalFinalize      = shakeFinalizeOutput (Proxy :: Proxy bitlen)
 
 -- | SHAKE256 (256 bits) extendable output function.  Supports an arbitrary
--- digest size (multiple of 8 bits), to be specified as a type parameter
--- of kind 'Nat'.
+-- digest size, to be specified as a type parameter of kind 'Nat'.
 --
 -- Note: outputs from @'SHAKE256' n@ and @'SHAKE256' m@ for the same input are
 -- correlated (one being a prefix of the other).  Results are unrelated to
@@ -61,9 +62,9 @@ instance (IsDivisibleBy8 bitlen, KnownNat bitlen) => HashAlgorithm (SHAKE128 bit
 data SHAKE256 (bitlen :: Nat) = SHAKE256
     deriving (Show, Data, Typeable)
 
-instance (IsDivisibleBy8 bitlen, KnownNat bitlen) => HashAlgorithm (SHAKE256 bitlen) where
+instance KnownNat bitlen => HashAlgorithm (SHAKE256 bitlen) where
     type HashBlockSize           (SHAKE256 bitlen) = 136
-    type HashDigestSize          (SHAKE256 bitlen) = Div8 bitlen
+    type HashDigestSize          (SHAKE256 bitlen) = Div8 (bitlen + 7)
     type HashInternalContextSize (SHAKE256 bitlen) = 344
     hashBlockSize  _          = 136
     hashDigestSize _          = byteLen (Proxy :: Proxy bitlen)
@@ -72,7 +73,7 @@ instance (IsDivisibleBy8 bitlen, KnownNat bitlen) => HashAlgorithm (SHAKE256 bit
     hashInternalUpdate        = c_sha3_update
     hashInternalFinalize      = shakeFinalizeOutput (Proxy :: Proxy bitlen)
 
-shakeFinalizeOutput :: (IsDivisibleBy8 bitlen, KnownNat bitlen)
+shakeFinalizeOutput :: KnownNat bitlen
                     => proxy bitlen
                     -> Ptr (Context a)
                     -> Ptr (Digest a)
@@ -80,6 +81,16 @@ shakeFinalizeOutput :: (IsDivisibleBy8 bitlen, KnownNat bitlen)
 shakeFinalizeOutput d ctx dig = do
     c_sha3_finalize_shake ctx
     c_sha3_output ctx dig (byteLen d)
+    shakeTruncate d (castPtr dig)
+
+shakeTruncate :: KnownNat bitlen => proxy bitlen -> Ptr Word8 -> IO ()
+shakeTruncate d ptr =
+    when (bits > 0) $ do
+        byte <- peekElemOff ptr index
+        pokeElemOff ptr index (byte .&. mask)
+  where
+    mask = (1 `shiftL` bits) - 1
+    (index, bits) = integralNatVal d `divMod` 8
 
 foreign import ccall unsafe "cryptonite_sha3_init"
     c_sha3_init :: Ptr (Context a) -> Word32 -> IO ()

Crypto/Internal/Nat.hs

@@ -15,8 +15,8 @@ module Crypto.Internal.Nat
 
 import           GHC.TypeLits
 
-byteLen :: (KnownNat bitlen, IsDivisibleBy8 bitlen, Num a) => proxy bitlen -> a
-byteLen d = fromInteger (natVal d `div` 8)
+byteLen :: (KnownNat bitlen, Num a) => proxy bitlen -> a
+byteLen d = fromInteger ((natVal d + 7) `div` 8)
 
 integralNatVal :: (KnownNat bitlen, Num a) => proxy bitlen -> a
 integralNatVal = fromInteger . natVal

tests/Hash.hs

@@ -8,7 +8,9 @@ module Hash
 import Crypto.Hash
 
 import qualified Data.ByteString as B
+import           Data.ByteArray (convert)
 import qualified Data.ByteArray.Encoding as B (convertToBase, Base(..))
+import           GHC.TypeLits
 import Imports
 
 v0,v1,v2 :: ByteString
@@ -234,7 +236,25 @@ makeTestChunk (hashName, hashAlg, _) =
         runhash hashAlg inp `propertyEq` runhashinc hashAlg (chunkS ckLen inp)
     ]
 
+-- SHAKE128 truncation example with expected byte at final position
+-- <https://csrc.nist.gov/CSRC/media/Projects/Cryptographic-Standards-and-Guidelines/documents/examples/ShakeTruncation.pdf>
+shake128TruncationBytes = [0x01, 0x03, 0x07, 0x0f, 0x0f, 0x2f, 0x6f, 0x6f]
+
+makeTestSHAKE128Truncation i byte =
+    testCase (show i) $ xof 4088 `B.snoc` byte @=? xof (4088 + i)
+  where
+    hashEmpty :: KnownNat n => proxy n -> Digest (SHAKE128 n)
+    hashEmpty _ = hash B.empty
+
+    xof n = case someNatVal n of
+                Nothing          -> error ("invalid Nat: " ++ show n)
+                Just (SomeNat p) -> convert (hashEmpty p)
+
 tests = testGroup "hash"
     [ testGroup "KATs" (map makeTestAlg expected)
     , testGroup "Chunking" (concatMap makeTestChunk expected)
+    , testGroup "Truncating"
+        [ testGroup "SHAKE128"
+            (zipWith makeTestSHAKE128Truncation [1..] shake128TruncationBytes)
+        ]
     ]