Handle decoding multiple ASCII armor messages from a single bytestream.

2 files changed, 19 insertions, 11 deletions

diff --git a/Codec/Encryption/OpenPGP/ASCIIArmor.hs b/Codec/Encryption/OpenPGP/ASCIIArmor.hs
index 74c8c69..075a481 100644
--- a/Codec/Encryption/OpenPGP/ASCIIArmor.hs
+++ b/Codec/Encryption/OpenPGP/ASCIIArmor.hs
@@ -5,9 +5,9 @@
 
 module Codec.Encryption.OpenPGP.ASCIIArmor (
    armor
- , decodeArmor
+ , decode
  , parseArmor
 ) where
 
 import Codec.Encryption.OpenPGP.ASCIIArmor.Encode (armor)
-import Codec.Encryption.OpenPGP.ASCIIArmor.Decode (decodeArmor, parseArmor)
+import Codec.Encryption.OpenPGP.ASCIIArmor.Decode (decode, parseArmor)
diff --git a/Codec/Encryption/OpenPGP/ASCIIArmor/Decode.hs b/Codec/Encryption/OpenPGP/ASCIIArmor/Decode.hs
index 8c2a8a3..e69087c 100644
--- a/Codec/Encryption/OpenPGP/ASCIIArmor/Decode.hs
+++ b/Codec/Encryption/OpenPGP/ASCIIArmor/Decode.hs
@@ -6,13 +6,14 @@
 
 module Codec.Encryption.OpenPGP.ASCIIArmor.Decode (
    parseArmor
- , decodeArmor
+ , decode
 ) where
 
 import Codec.Encryption.OpenPGP.ASCIIArmor.Types
-import Control.Applicative (many, (<|>), (<$>))
+import Control.Applicative (many, (<|>), (<$>), Alternative, (*>))
 import Data.Attoparsec.ByteString (Parser, many1, string, inClass, notInClass, satisfy, word8, (<?>), parse, IResult(..))
-import Data.Attoparsec.ByteString.Char8 (isDigit_w8)
+import Data.Attoparsec.ByteString.Char8 (isDigit_w8, anyChar)
+import Data.Attoparsec.Combinator (manyTill)
 import Data.Bits (shiftL)
 import Data.ByteString (ByteString)
 import qualified Data.ByteString as B
@@ -25,15 +26,19 @@ import Data.Serialize.Put (runPut, putWord32be)
 import Data.String (IsString, fromString)
 import Data.Word (Word32)
 
-decodeArmor :: (Integral a, Read a, Show a, IsString e) => ByteString -> Either e (Armor a)
-decodeArmor bs = case parse parseArmor bs of
-                     Fail t c e -> Left (fromString e)
-                     Partial _ -> Left (fromString "what")
-                     Done _ r -> Right r
+decode :: (Integral a, Read a, Show a, IsString e) => ByteString -> Either e ([Armor a])
+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 = many parseArmor
 
 parseArmor :: (Integral a, Read a, Show a) => Parser (Armor a)
 parseArmor = do
-    atype <- beginLine <?> "begin line"
+    atype <- prefixed beginLine <?> "begin line"
     headers <- armorHeaders <?> "headers"
     blankishLine <?> "blank line"
     payload <- base64Data <?> "base64 data"
@@ -126,3 +131,6 @@ d24 = do
     b <- getWord8
     c <- getWord8
     return $ shiftL (fromIntegral a :: Word32) 16 + shiftL (fromIntegral b :: Word32) 8 + (fromIntegral c :: Word32)
+
+prefixed :: Parser a -> Parser a
+prefixed end = end <|> anyChar *> prefixed end