Harden streaming fetch parser coverage

Author: mpscholten

src/Network/HaskellNet/IMAP.hs

@@ -431,11 +431,19 @@ fetchRPeek conn range =
     do ls <- fetchByByteStringR conn range "BODY.PEEK[]"
        return $ map (\(uid, vs) -> (uid, fromMaybe BS.empty $ lookup' "BODY[]" vs)) ls
 
+-- | Fetch arbitrary data items and return values as 'String's.
+--
+-- This is kept for compatibility. Prefer 'fetchByByteString' for message
+-- bodies or other data that may be large or non-textual.
 fetchByString :: IMAPConnection -> UID -> String
               -> IO [(String, String)]
 fetchByString conn uid command =
     map (\(key, value) -> (key, BS.unpack value)) <$> fetchByByteString conn uid command
 
+-- | Fetch arbitrary data items and return raw 'ByteString' values.
+--
+-- Literal values are read directly from the stream instead of first building a
+-- full response buffer, so this is the preferred API for large messages.
 fetchByByteString :: IMAPConnection -> UID -> String
                   -> IO [(String, ByteString)]
 fetchByByteString conn uid command =
@@ -444,13 +452,15 @@ fetchByByteString conn uid command =
          (_, pairs):_ -> return pairs
          [] -> return []
 
+-- | Range variant of 'fetchByString'.
 fetchByStringR :: IMAPConnection -> (UID, UID) -> String
                -> IO [(UID, [(String, String)])]
 fetchByStringR conn (s, e) command =
     map unpackFetch <$> fetchByByteStringR conn (s, e) command
     where unpackFetch (uid, pairs) =
               (uid, map (\(key, value) -> (key, BS.unpack value)) pairs)
 
+-- | Range variant of 'fetchByByteString'.
 fetchByByteStringR :: IMAPConnection -> (UID, UID) -> String
                    -> IO [(UID, [(String, ByteString)])]
 fetchByByteStringR conn (s, e) command =
@@ -476,6 +486,8 @@ fetchCommandBS conn command proc =
          BAD _ msg     -> fail ("BAD: " ++ msg)
          PREAUTH _ msg -> fail ("preauth: " ++ msg)
 
+-- Streaming FETCH response parser. Literal payloads are read with 'bsGet' so
+-- large message bodies do not require buffering the complete server response.
 getFetchResponseBS :: BSStream -> String
                    -> IO (ServerResponse, MboxUpdate, [(Integer, [(String, ByteString)])])
 getFetchResponseBS s tag = go Nothing Nothing []
@@ -652,12 +664,7 @@ parseQuotedValueBS input =
 parseAtomValueBS :: ByteString -> Maybe (ByteString, ByteString)
 parseAtomValueBS input =
     let (value, rest) = BS.span isAtomValueChar input
-    in if BS.null value
-       then Nothing
-       else let normalized = if BS.map toUpper value == BS.pack "NIL"
-                             then BS.empty
-                             else value
-            in Just (normalized, rest)
+    in if BS.null value then Nothing else Just (value, rest)
   where
     isAtomValueChar c = not $ c `elem` " (){%*\"\\]\r\n"
 

test/IMAPParsersTest.hs

@@ -255,6 +255,79 @@ imapFetchTest =
               ]
           fetched <- IMAP.fetch conn 42
           body @=? fetched
+    , "fetchPeek reads body response without setting Seen" ~: TestCase $ do
+          let body = BS.pack "peeked"
+          conn <- scriptedConnection
+              [ line ("* 12 FETCH (BODY[] {" ++ show (BS.length body) ++ "}")
+              , ReadBytes body
+              , line " UID 42)"
+              , okLine "FETCH completed"
+              ]
+          fetched <- IMAP.fetchPeek conn 42
+          body @=? fetched
+    , "fetchSize parses scalar size responses" ~: TestCase $ do
+          conn <- scriptedConnection
+              [ line "* 12 FETCH (RFC822.SIZE 12345 UID 42)"
+              , okLine "FETCH completed"
+              ]
+          fetched <- IMAP.fetchSize conn 42
+          12345 @=? fetched
+    , "fetchFlags parses parenthesized flag responses" ~: TestCase $ do
+          conn <- scriptedConnection
+              [ line "* 12 FETCH (FLAGS (\\Seen \\Deleted) UID 42)"
+              , okLine "FETCH completed"
+              ]
+          fetched <- IMAP.fetchFlags conn 42
+          [Seen, Deleted] @=? fetched
+    , "fetchHeaderFields matches normalized body section keys" ~: TestCase $ do
+          let headers = BS.pack "Subject: Hi\r\nFrom: a@example.com\r\n\r\n"
+          conn <- scriptedConnection
+              [ line ("* 12 FETCH (BODY[HEADER.FIELDS (SUBJECT FROM)] {"
+                       ++ show (BS.length headers) ++ "}")
+              , ReadBytes headers
+              , line " UID 42)"
+              , okLine "FETCH completed"
+              ]
+          fetched <- IMAP.fetchHeaderFields conn 42 ["Subject", "From"]
+          headers @=? fetched
+    , "fetchR maps sequence numbers to response UIDs" ~: TestCase $ do
+          let firstBody = BS.pack "one"
+              secondBody = BS.pack "two"
+          conn <- scriptedConnection
+              [ line ("* 1 FETCH (BODY[] {" ++ show (BS.length firstBody) ++ "}")
+              , ReadBytes firstBody
+              , line " UID 101)"
+              , line ("* 2 FETCH (BODY[] {" ++ show (BS.length secondBody) ++ "}")
+              , ReadBytes secondBody
+              , line " UID 102)"
+              , okLine "FETCH completed"
+              ]
+          fetched <- IMAP.fetchR conn (1, 2)
+          [(101, firstBody), (102, secondBody)] @=? fetched
+    , "fetchByString keeps scalar and literal values compatible" ~: TestCase $ do
+          let headers = BS.pack "hello"
+          conn <- scriptedConnection
+              [ line ("* 12 FETCH (RFC822.SIZE 123 FLAGS (\\Seen) BODY[HEADER] {"
+                       ++ show (BS.length headers) ++ "}")
+              , ReadBytes headers
+              , line " NILKEY NIL QUOTED \"world\" UID 42)"
+              , okLine "FETCH completed"
+              ]
+          fetched <- IMAP.fetchByString conn 42
+              "RFC822.SIZE FLAGS BODY[HEADER] NILKEY QUOTED"
+          [ ("RFC822.SIZE", "123")
+            , ("FLAGS", "(\\Seen)")
+            , ("BODY[HEADER]", "hello")
+            , ("NILKEY", "NIL")
+            , ("QUOTED", "\"world\"")
+            , ("UID", "42")
+            ] @=? fetched
+    , "store accepts FETCH data in STORE responses" ~: TestCase $ do
+          conn <- scriptedConnection
+              [ line "* 12 FETCH (UID 42 FLAGS (\\Seen))"
+              , okLine "STORE completed"
+              ]
+          IMAP.store conn 42 (IMAP.PlusFlags [Seen])
     ]
 
 testData = [ "base" ~: baseTest