Fix decoding of multi-part ASCII armor.

3 files changed, 38 insertions, 34 deletions

diff --git a/Codec/Encryption/OpenPGP/ASCIIArmor/Decode.hs b/Codec/Encryption/OpenPGP/ASCIIArmor/Decode.hs
index e69087c..0376abc 100644
--- a/Codec/Encryption/OpenPGP/ASCIIArmor/Decode.hs
+++ b/Codec/Encryption/OpenPGP/ASCIIArmor/Decode.hs
@@ -26,17 +26,17 @@ import Data.Serialize.Put (runPut, putWord32be)
 import Data.String (IsString, fromString)
 import Data.Word (Word32)
 
-decode :: (Integral a, Read a, Show a, IsString e) => ByteString -> Either e ([Armor a])
+decode :: IsString e => ByteString -> Either e [Armor]
 decode bs = go (parse parseArmors bs)
     where
         go (Fail t c e) = Left (fromString e)
         go (Partial cont) = go (cont B.empty)
         go (Done _ r) = Right r
 
-parseArmors :: (Integral a, Read a, Show a) => Parser ([Armor a])
+parseArmors :: Parser [Armor]
 parseArmors = many parseArmor
 
-parseArmor :: (Integral a, Read a, Show a) => Parser (Armor a)
+parseArmor :: Parser Armor
 parseArmor = do
     atype <- prefixed beginLine <?> "begin line"
     headers <- armorHeaders <?> "headers"
@@ -45,23 +45,27 @@ parseArmor = do
     endLine atype <?> "end line"
     return $ Armor atype headers payload
 
-beginLine :: (Integral a, Read a, Show a) => Parser (ArmorType a)
+beginLine :: Parser ArmorType
 beginLine = do
-    string "-----BEGIN PGP "
-    atype <- message <|> pubkey <|> privkey<|> parts <|> signature
-    string "-----"
-    many (satisfy (inClass " \t"))
-    lineEnding
+    string "-----BEGIN PGP " <?> "leading minus-hyphens"
+    atype <- pubkey <|> privkey <|> parts <|> message <|> signature
+    string "-----" <?> "trailing minus-hyphens"
+    many (satisfy (inClass " \t")) <?> "whitespace"
+    lineEnding <?> "line ending"
     return atype
     where
-        message = string "MESSAGE" >> return ArmorMessage
-        pubkey = string "PUBLIC KEY BLOCK" >> return ArmorPublicKeyBlock
-        privkey = string "PRIVATE KEY BLOCK" >> return ArmorPrivateKeyBlock
-        signature = string "SIGNATURE" >> return ArmorSignature
-        parts = do
-            string "MESSAGE, PART "
-            firstnum <- read . BC8.unpack . B.pack <$> many1 (satisfy isDigit_w8)
-            return $ ArmorSplitMessageIndefinite firstnum
+        message = string "MESSAGE" *> return ArmorMessage
+        pubkey = string "PUBLIC KEY BLOCK" *> return ArmorPublicKeyBlock
+        privkey = string "PRIVATE KEY BLOCK" *> return ArmorPrivateKeyBlock
+        signature = string "SIGNATURE" *> return ArmorSignature
+        parts = string "MESSAGE, PART " *> (partsdef <|> partsindef)
+        partsdef = do
+            firstnum <- num
+            word8 (fromIntegral . fromEnum $ '/')
+            secondnum <- num
+            return $ ArmorSplitMessage (B.pack firstnum) (B.pack secondnum)
+        partsindef = ArmorSplitMessageIndefinite . B.pack <$> num
+        num = many1 (satisfy isDigit_w8) <?> "number"
 
 lineEnding :: Parser ByteString
 lineEnding = string "\n" <|> string "\r\n"
@@ -82,17 +86,17 @@ armorHeader = do
 blankishLine ::  Parser ByteString
 blankishLine = many (satisfy (inClass " \t")) >> lineEnding
 
-endLine :: (Integral a, Read a, Show a) => ArmorType a -> Parser ByteString
+endLine :: ArmorType -> Parser ByteString
 endLine atype = do
     string $ "-----END PGP " `B.append` aType atype `B.append` "-----"
     lineEnding
 
-aType :: (Integral a, Read a, Show a) => ArmorType a -> ByteString
+aType :: ArmorType -> ByteString
 aType (ArmorMessage) = BC8.pack "MESSAGE"
 aType (ArmorPublicKeyBlock) = BC8.pack "PUBLIC KEY BLOCK"
 aType (ArmorPrivateKeyBlock) = BC8.pack "PRIVATE KEY BLOCK"
-aType (ArmorSplitMessage x y) = BC8.pack $ "MESSAGE, PART " ++ show x ++ "/" ++ show y
-aType (ArmorSplitMessageIndefinite x) = BC8.pack $ "MESSAGE, PART " ++ show x
+aType (ArmorSplitMessage x y) = BC8.pack "MESSAGE, PART " `B.append` x `B.append` BC8.singleton '/' `B.append` y
+aType (ArmorSplitMessageIndefinite x) = BC8.pack "MESSAGE, PART " `B.append` x
 aType (ArmorSignature) = BC8.pack "SIGNATURE"
 
 base64Data :: Parser ByteString
diff --git a/Codec/Encryption/OpenPGP/ASCIIArmor/Encode.hs b/Codec/Encryption/OpenPGP/ASCIIArmor/Encode.hs
index 99d4d54..28bb3e6 100644
--- a/Codec/Encryption/OpenPGP/ASCIIArmor/Encode.hs
+++ b/Codec/Encryption/OpenPGP/ASCIIArmor/Encode.hs
@@ -17,22 +17,22 @@ import Data.Serialize (put)
 import Data.Serialize.Put (runPut, putWord32be)
 import Data.String (IsString, fromString)
 
-encode :: (Integral a, Show a) => [Armor a] -> ByteString
+encode :: [Armor] -> ByteString
 encode = B.concat . map armor
 
-armor :: (Integral a, Show a) => Armor a -> ByteString
+armor :: Armor -> ByteString
 armor (Armor atype ahs bs) = beginLine atype `B.append` armorHeaders ahs `B.append` blankLine `B.append` armorData bs `B.append` armorChecksum bs `B.append` endLine atype
 
 blankLine :: ByteString
 blankLine = BC8.singleton '\n'
 
-beginLine :: (Integral a, Show a) => ArmorType a -> ByteString
+beginLine :: ArmorType -> ByteString
 beginLine atype = BC8.pack "-----BEGIN PGP " `B.append` aType atype `B.append` BC8.pack "-----\n"
 
-endLine :: (Integral a, Show a) => ArmorType a -> ByteString
+endLine :: ArmorType -> ByteString
 endLine atype = BC8.pack "-----END PGP " `B.append` aType atype `B.append` BC8.pack "-----\n"
 
-aType :: (Integral a, Show a) => ArmorType a -> ByteString
+aType :: ArmorType -> ByteString
 aType (ArmorMessage) = BC8.pack "MESSAGE"
 aType (ArmorPublicKeyBlock) = BC8.pack "PUBLIC KEY BLOCK"
 aType (ArmorPrivateKeyBlock) = BC8.pack "PRIVATE KEY BLOCK"
diff --git a/Codec/Encryption/OpenPGP/ASCIIArmor/Types.hs b/Codec/Encryption/OpenPGP/ASCIIArmor/Types.hs
index 8c7ef6f..46416c1 100644
--- a/Codec/Encryption/OpenPGP/ASCIIArmor/Types.hs
+++ b/Codec/Encryption/OpenPGP/ASCIIArmor/Types.hs
@@ -10,14 +10,14 @@ module Codec.Encryption.OpenPGP.ASCIIArmor.Types (
 
 import Data.ByteString (ByteString)
 
-data Armor a = Armor (ArmorType a) [(String, String)] ByteString
-   | ClearSigned [(String, String)] String (Armor a)
+data Armor = Armor ArmorType [(String, String)] ByteString
+   | ClearSigned [(String, String)] String Armor
     deriving (Show, Eq)
 
-data ArmorType a = ArmorMessage
-                 | ArmorPublicKeyBlock
-                 | ArmorPrivateKeyBlock
-                 | ArmorSplitMessage a a
-                 | ArmorSplitMessageIndefinite a
-                 | ArmorSignature
+data ArmorType = ArmorMessage
+               | ArmorPublicKeyBlock
+               | ArmorPrivateKeyBlock
+               | ArmorSplitMessage ByteString ByteString
+               | ArmorSplitMessageIndefinite ByteString
+               | ArmorSignature
     deriving (Show, Eq)