add checksum to encrypted data

Needed to verify decryption puzzles

1 files changed, 38 insertions, 27 deletions

diff --git a/Encryption.hs b/Encryption.hs
index 385f36a..bf2370c 100644
--- a/Encryption.hs
+++ b/Encryption.hs
@@ -65,28 +65,25 @@ genIV (Password password) = do
 	let ivlen = B.length (Raaz.toByteString exampleiv)
 	return $ fromMaybe (error "genIV fromByteString failed") $
 		Raaz.fromByteString $ B.take ivlen $
-		Raaz.toByteString $ Raaz.sha256 password
+			Raaz.toByteString $ Raaz.sha256 password
 
 cipher :: Raaz.AES 256 'Raaz.CBC
 cipher = Raaz.aes256cbc
 
-blocksize :: Int	
-blocksize = fromIntegral $ Raaz.blockSize cipher
-
 encrypt :: Tunables -> KeyEncryptionKey -> SecretKey -> EncryptedSecretKey
 encrypt tunables kek (SecretKey secret) = EncryptedSecretKey b (keyBruteForceCalc kek)
   where
 	-- Raaz does not seem to provide a high-level interface
-	-- for AES encryption, so use unsafeEncrypt, doing our own padding
-	-- of the secret key, so that it is a multiple of
-	-- the block size.
+	-- for AES encryption, so use unsafeEncrypt. The use of 
+	-- EncryptableBytes makes sure it's provided with a 
+	-- multiple of the AES block size.
 	b = Raaz.unsafeEncrypt cipher (keyEncryptionKey kek, keyEncryptionIV kek) $
-		getPaddedBytes $ toPaddedBytes tunables blocksize secret
+		getEncryptableBytes $ toEncryptableBytes tunables secret
 
 decrypt :: KeyEncryptionKey -> EncryptedSecretKey -> Maybe SecretKey
-decrypt kek (EncryptedSecretKey b _) = SecretKey <$> fromPaddedBytes pbs
+decrypt kek (EncryptedSecretKey b _) = SecretKey <$> fromEncryptableBytes pbs
   where
-	pbs = PaddedBytes $ Raaz.unsafeDecrypt cipher (keyEncryptionKey kek, keyEncryptionIV kek) b
+	pbs = EncryptableBytes $ Raaz.unsafeDecrypt cipher (keyEncryptionKey kek, keyEncryptionIV kek) b
 
 -- | A stream of all the key encryption keys that need to be tried to 
 -- decrypt. 
@@ -191,29 +188,43 @@ mixinRandomObstacle (RandomObstacle r) k = k'
 orBytes :: B.ByteString -> B.ByteString -> B.ByteString
 orBytes a b = B.pack $ map (uncurry (.|.)) $ zip (B.unpack a) (B.unpack b)
 
-newtype PaddedBytes = PaddedBytes { getPaddedBytes :: B.ByteString }
+-- | A bytestring that can be AES enctypted. It includes a checksum,
+-- and size, and is padded to the objectSize with NULs.
+--
+-- This is a multiple of the AES blocksize, as long as objectSize is,
+-- which should always be the case.
+newtype EncryptableBytes = EncryptableBytes { getEncryptableBytes :: B.ByteString }
 	deriving (Show)
 
--- Pad with NULs. Since the bytestring can itself include NULs, prefix
--- with the length. Length is itself padded with 0's.
-toPaddedBytes :: Tunables -> Int -> B.ByteString -> PaddedBytes
-toPaddedBytes tunables n b = PaddedBytes $
-	B8.pack paddedlen <> B.singleton 0 <> b <> padding
+toEncryptableBytes :: Tunables -> B.ByteString -> EncryptableBytes
+toEncryptableBytes tunables b = EncryptableBytes $ 
+	padBytes (objectSize tunables) $ 
+		checksum <> sep <> len <> sep <> b
+  where
+	checksum = Raaz.toByteString $ Raaz.sha256 b
+	len = B8.pack (show (B.length b))
+	sep = B.singleton 0
+
+padBytes :: Int -> B.ByteString -> B.ByteString
+padBytes n b = b <> padding
   where
 	len = B.length b
 	r = len `rem` n
 	padding
 		| r == 0 = B.empty
 		| otherwise = B.replicate (n - r) 0
-	paddedlen = 
-		let s = show len
-		in replicate (lensz - length s) '0' ++ s
-	lensz = length $ show $ objectSize tunables
-
-fromPaddedBytes :: PaddedBytes -> Maybe B.ByteString
-fromPaddedBytes (PaddedBytes b) = case B.break (== 0) b of
-	(header, rest)
-		| B.null header || B.null rest -> Nothing
+
+fromEncryptableBytes :: EncryptableBytes -> Maybe B.ByteString
+fromEncryptableBytes (EncryptableBytes b) = case B.break (== 0) b of
+	(checksum, rest)
+		| B.null checksum || B.null rest -> Nothing
 		| otherwise -> do
-			len <- readMaybe (B8.unpack header)
-			return $ B.take len $ B.drop 1 rest
+			case B.break (== 0) (B.drop 1 rest) of
+				(lenb, rest')
+					| B.null lenb || B.null rest' -> Nothing
+					| otherwise -> do
+						len <- readMaybe (B8.unpack lenb)
+						let d = B.take len $ B.drop 1 rest'
+						if checksum == Raaz.toByteString (Raaz.sha256 d)
+							then Just d
+							else Nothing