Split large UID sets into bounded FETCH commands

src/Network/HaskellNet/IMAP.hs

@@ -266,7 +266,7 @@
     do (c, num) <- sendCommand' conn $ "AUTHENTICATE " ++ show at
        let challenge =
                if BS.take 2 c == BS.pack "+ "
                then A.b64Decode $ BS.unpack $ head $
                     dropWhile (isSpace . BS.last) $ BS.inits $ BS.drop 2 c
                else ""
        bsPutCrLf (stream conn) $ BS.pack $
@@ -475,16 +475,48 @@
 -- | Fetch arbitrary data items for an exact set of UIDs.
 --
 -- Unlike 'fetchByByteStringR', this does not fetch messages whose UIDs happen
--- to lie between sparse search results.
+-- to lie between sparse search results. Large UID sets are split across
+-- commands whose generated command lines are at most approximately 1000
+-- octets, as recommended by RFC 2683 section 3.2.1.5.
 fetchByByteStringSet :: IMAPConnection -> [UID] -> String
                      -> IO [(UID, [(String, ByteString)])]
 fetchByByteStringSet _ [] _ = return []
 fetchByByteStringSet conn uids command =
-    fetchCommandBS conn
-        ("UID FETCH "++intercalate "," (map show uids)++" "++command) proc
-    where proc (n, ps) =
+    case chunkUIDsByLength availableUIDLength uids of
+      Nothing -> fail "UID FETCH command is too long to contain a UID"
+      Just uidChunks -> concat <$> mapM fetchUIDChunk uidChunks
+    where fetchPrefix = "UID FETCH "
+          fetchSuffix = " " ++ command
+          availableUIDLength = maxGeneratedCommandLength
+                             - length fetchPrefix - length fetchSuffix
+          fetchUIDChunk uidChunk =
+              fetchCommandBS conn
+                  (fetchPrefix ++ intercalate "," (map show uidChunk)
+                               ++ fetchSuffix) proc
+          proc (n, ps) =
               (maybe (toEnum (fromIntegral n)) (read . BS.unpack) (lookup' "UID" ps), ps)
 
+-- RFC 2683 recommends that clients limit generated command lines to
+-- approximately 1000 octets. Reserve space for the six-octet command tag,
+-- its separating space, and the terminating CRLF added by
+-- 'sendCommandNoResponse'.
+maxGeneratedCommandLength :: Int
+maxGeneratedCommandLength = 1000 - 6 - 1 - 2
+
+chunkUIDsByLength :: Int -> [UID] -> Maybe [[UID]]
+chunkUIDsByLength maxLength = go [] 0
+  where
+    go [] _ [] = Just []
+    go current _ [] = Just [reverse current]
+    go current currentLength (uid:rest)
+      | uidLength > maxLength = Nothing
+      | nextLength <= maxLength = go (uid:current) nextLength rest
+      | otherwise = (reverse current :) <$> go [uid] uidLength rest
+      where
+        uidLength = length (show uid)
+        separatorLength = if null current then 0 else 1
+        nextLength = currentLength + separatorLength + uidLength
+
 fetchCommand :: IMAPConnection -> String
              -> ((Integer, [(String, String)]) -> b) -> IO [b]
 fetchCommand conn command proc =

test/IMAPParsersTest.hs

@@ -4,6 +4,7 @@ import Data.ByteString (ByteString)
 import qualified Data.ByteString as B
 import qualified Data.ByteString.Char8 as BS
 import Data.IORef
+import Data.List (intercalate)
 import Network.HaskellNet.BSStream
 import qualified Network.HaskellNet.IMAP as IMAP
 import Network.HaskellNet.IMAP.Connection
@@ -480,6 +481,37 @@ imapFetchTest =
           [] @=? fetched
           actual <- written
           B.empty @=? actual
+    , "fetchByByteStringSet splits long command lines" ~: TestCase $ do
+          let uids = take 123 [1000000,1000002..]
+              firstChunk = take 121 uids
+              secondChunk = drop 121 uids
+              firstUID = 1000240
+              secondUID = 1000242
+              firstStructure = BS.pack "(\"TEXT\" \"PLAIN\")"
+              secondStructure = BS.pack "(\"APPLICATION\" \"PDF\")"
+          (conn, written) <- scriptedConnection
+              [ line ("* 1 FETCH (BODYSTRUCTURE (\"TEXT\" \"PLAIN\") UID "
+                      ++ show firstUID ++ ")")
+              , line "000000 OK FETCH completed"
+              , line ("* 2 FETCH (BODYSTRUCTURE (\"APPLICATION\" \"PDF\") UID "
+                      ++ show secondUID ++ ")")
+              , line "000001 OK FETCH completed"
+              ]
+          fetched <- IMAP.fetchByByteStringSet conn uids "BODYSTRUCTURE"
+          [ (firstUID, [("BODYSTRUCTURE", firstStructure),
+                        ("UID", BS.pack $ show firstUID)])
+            , (secondUID, [("BODYSTRUCTURE", secondStructure),
+                           ("UID", BS.pack $ show secondUID)])
+            ] @=? fetched
+          actual <- written
+          let expected = B.append
+                  (commandBytes $ "000000 UID FETCH "
+                                  ++ intercalate "," (map show firstChunk)
+                                  ++ " BODYSTRUCTURE")
+                  (commandBytes $ "000001 UID FETCH "
+                                  ++ intercalate "," (map show secondChunk)
+                                  ++ " BODYSTRUCTURE")
+          expected @=? actual
     , "fetchByString keeps scalar and literal values compatible" ~: TestCase $ do
           let headers = BS.pack "hello"
           (conn, _) <- scriptedConnection