Handle decoding multiple ASCII armor messages from a single bytestream.

6 files changed, 71 insertions, 18 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
diff --git a/openpgp-asciiarmor.cabal b/openpgp-asciiarmor.cabal
index 4b7e25d..ff2289a 100644
--- a/openpgp-asciiarmor.cabal
+++ b/openpgp-asciiarmor.cabal
@@ -12,6 +12,8 @@ Category:            Codec, Data
 Build-type:          Simple
 Extra-source-files: tests/suite.hs
   , tests/data/msg1.asc
+  , tests/data/msg1a.asc
+  , tests/data/msg1b.asc
   , tests/data/msg1.gpg
 
 Cabal-version:       >= 1.10
diff --git a/tests/data/msg1a.asc b/tests/data/msg1a.asc
new file mode 100644
index 0000000..64dc079
--- /dev/null
+++ b/tests/data/msg1a.asc
@@ -0,0 +1,11 @@
+This file contains an ASCII-armored PGP message enclosed between
+
+-----BEGIN PGP MESSAGE-----
+Version: OpenPrivacy 0.99
+
+yDgBO22WxBHv7O8X7O/jygAEzol56iUKiXmV+XmpCtmpqQUKiQrFqclFqUDBovzS
+vBSFjNSiVHsuAA==
+=njUN
+-----END PGP MESSAGE-----
+
+two lines of arbitrary text.
diff --git a/tests/data/msg1b.asc b/tests/data/msg1b.asc
new file mode 100644
index 0000000..1176420
--- /dev/null
+++ b/tests/data/msg1b.asc
@@ -0,0 +1,28 @@
+This file contains three ASCII-armored PGP messages interspersed
+
+-----BEGIN PGP MESSAGE-----
+Version: OpenPrivacy 0.99
+
+yDgBO22WxBHv7O8X7O/jygAEzol56iUKiXmV+XmpCtmpqQUKiQrFqclFqUDBovzS
+vBSFjNSiVHsuAA==
+=njUN
+-----END PGP MESSAGE-----
+-----BEGIN PGP MESSAGE-----
+Version: OpenPrivacy 0.99
+
+yDgBO22WxBHv7O8X7O/jygAEzol56iUKiXmV+XmpCtmpqQUKiQrFqclFqUDBovzS
+vBSFjNSiVHsuAA==
+=njUN
+-----END PGP MESSAGE-----
+
+with arbitrary text.
+
+-----BEGIN PGP MESSAGE-----
+Version: OpenPrivacy 0.99
+
+yDgBO22WxBHv7O8X7O/jygAEzol56iUKiXmV+XmpCtmpqQUKiQrFqclFqUDBovzS
+vBSFjNSiVHsuAA==
+=njUN
+-----END PGP MESSAGE-----
+
+All three messages are identical.
diff --git a/tests/suite.hs b/tests/suite.hs
index 272295a..45a2fb2 100644
--- a/tests/suite.hs
+++ b/tests/suite.hs
@@ -3,7 +3,7 @@ import Test.Framework.Providers.HUnit
 
 import Test.HUnit
 
-import Codec.Encryption.OpenPGP.ASCIIArmor (armor, decodeArmor)
+import Codec.Encryption.OpenPGP.ASCIIArmor (armor, decode)
 import Codec.Encryption.OpenPGP.ASCIIArmor.Types
 
 import Data.ByteString (ByteString)
@@ -15,13 +15,15 @@ import Data.Word (Word32)
 testCRC24 :: ByteString -> Word32 -> Assertion
 testCRC24 bs crc = assertEqual "crc24" crc (crc24 bs)
 
-testArmorDecode :: FilePath -> FilePath -> Assertion
-testArmorDecode fp target = do
+testArmorDecode :: FilePath -> [FilePath] -> Assertion
+testArmorDecode fp targets = do
     bs <- B.readFile $ "tests/data/" ++ fp
-    tbs <- B.readFile $ "tests/data/" ++ target
-    case decodeArmor bs of
+    tbss <- mapM (\target -> B.readFile $ "tests/data/" ++ target) targets
+    case decode bs of
         Left e -> assertFailure $ "Decode failed (" ++ e ++ ") on " ++ fp
-        Right (Armor at hdrs pl) -> do assertEqual ("for " ++ fp) tbs pl
+        Right as -> do assertEqual ("for " ++ fp) tbss (map getPayload as)
+    where
+        getPayload (Armor _ _ pl) = pl
 
 testArmorEncode :: FilePath -> FilePath -> Assertion
 testArmorEncode fp target = do
@@ -36,7 +38,9 @@ tests = [
     , testCase "CRC24: hOpenPGP and friends" (testCRC24 (BC8.pack "hOpenPGP and friends") 11940960)
     ]
  , testGroup "ASCII armor" [
-      testCase "Decode sample armor" (testArmorDecode "msg1.asc" "msg1.gpg")
+      testCase "Decode sample armor" (testArmorDecode "msg1.asc" ["msg1.gpg"])
+    , testCase "Decode sample armor with cruft" (testArmorDecode "msg1a.asc" ["msg1.gpg"])
+    , testCase "Decode multiple sample armors" (testArmorDecode "msg1b.asc" ["msg1.gpg","msg1.gpg","msg1.gpg"])
     , testCase "Encode sample armor" (testArmorEncode "msg1.gpg" "msg1.asc")
     ]
  ]