Merge pull request #91 from scrive/fix-dechunker

Fix consumeChunks: actually strip chunk size-lines and trailing CRLFs

Author: stepcut

src/Happstack/Server/Internal/Handler.hs

@@ -4,6 +4,7 @@ module Happstack.Server.Internal.Handler
     ( request
     , parseResponse
     , putRequest
+    , consumeChunks
     ) where
 
 import qualified Paths_happstack_server as Paths
@@ -154,24 +155,32 @@ parseResponse inputStr =
                  mbCL
        return $ Response {rsCode=code,rsHeaders=headers,rsBody=body,rsFlags=RsFlags ContentLength,rsValidator=Nothing}
 
--- http://www.w3.org/Protocols/rfc2616/rfc2616-sec3.html
--- note this does NOT handle extenions
-consumeChunks::L.ByteString->(L.ByteString,L.ByteString)
-consumeChunks str = let (parts,tr,rest) = consumeChunksImpl str in (L.concat . (++ [tr]) .map snd $ parts,rest)
+-- | Decode an HTTP/1.1 Transfer-Encoding: chunked body.
+-- See http://www.w3.org/Protocols/rfc2616/rfc2616-sec3.html (chunk extensions are not handled).
+-- Returns @(decoded body, bytes following the chunked stream)@. The trailer (optional headers
+-- after the terminating 0-chunk, RFC 2616 §3.6.1) is dropped from the body — trailers are
+-- headers, not body content.
+consumeChunks :: L.ByteString -> (L.ByteString, L.ByteString)
+consumeChunks str = let (parts,_tr,rest) = consumeChunksImpl str in (L.concat . map snd $ parts, rest)
 
+-- | Worker for 'consumeChunks'. Returns @(chunks, trailer, remainder)@, where each chunk is
+-- @(declared-length, body-bytes)@ and the terminating 0-chunk has empty body.
 consumeChunksImpl :: L.ByteString -> ([(Int64, L.ByteString)], L.ByteString, L.ByteString)
 consumeChunksImpl str
     | L.null str = ([],L.empty,str)
-    | chunkLen == 0 = let (last,rest') = L.splitAt lenLine1 str
+    | chunkLen == 0 = let (_sizeLine,rest') = L.splitAt lenLine1 str
                           (tr',rest'') = getTrailer rest'
-                      in ([(0,last)],tr',rest'')
-    | otherwise = ((chunkLen,part):crest,tr,rest2)
+                      in ([(0,L.empty)],tr',rest'')
+    | otherwise = ((chunkLen,body):crest,tr,rest2)
     where
       line1 = head $ lazylines str
       lenLine1 = (L.length line1) + 1 -- endchar
       chunkLen = (fst $ head $ readHex $ L.unpack line1)
       len = chunkLen + lenLine1 + 2
-      (part,rest) = L.splitAt len str
+      -- Each frame on the wire is "<hex-size>\r\n<body-bytes>\r\n". Skip the size-line and
+      -- take exactly chunkLen bytes so the trailing CRLF stays out of the decoded body.
+      (frame,rest) = L.splitAt len str
+      body = L.take chunkLen . L.drop lenLine1 $ frame
       (crest,tr,rest2) = consumeChunksImpl rest
       getTrailer s = L.splitAt index s
           where index | crlfLC `L.isPrefixOf` s = 2

tests/Happstack/Server/Tests.hs

@@ -19,6 +19,7 @@ import Happstack.Server.FileServe.BuildingBlocks (sendFileResponse)
 import Happstack.Server.Cookie
 import Happstack.Server.Internal.Compression
 import Happstack.Server.Internal.Cookie
+import Happstack.Server.Internal.Handler (consumeChunks)
 import Happstack.Server.Internal.Multipart
 import Happstack.Server.Internal.MessageWrap
 import Happstack.Server.Internal.RFC822Headers (ContentDisposition(..), parseContentDisposition)
@@ -37,6 +38,7 @@ allTests =
                                 , cookieHeaderOrderTest
                                 , pContentDispositionFilename
                                 , applicativeTest
+                                , consumeChunksTest
                                 ]
 
 cookieParserTest :: Test
@@ -251,6 +253,21 @@ pContentDispositionFilename =
        c <- parseContentDisposition doesNotWorkWithOldParserButWithNew
        assertEqual "parseContentDisposition" c (ContentDisposition "form-data" [("filename","file.pdf"),("name","file")])
 
+consumeChunksTest :: Test
+consumeChunksTest =
+  "consumeChunks decodes Transfer-Encoding: chunked" ~:
+    [ -- Single chunk
+      consumeChunks (pack "5\r\nhello\r\n0\r\n\r\n") @?= (pack "hello", pack "")
+      -- Multiple chunks concatenated
+    , consumeChunks (pack "5\r\nhello\r\n5\r\nworld\r\n0\r\n\r\n")
+        @?= (pack "helloworld", pack "")
+      -- Empty body (only terminator)
+    , consumeChunks (pack "0\r\n\r\n") @?= (pack "", pack "")
+      -- Trailing bytes after final chunk are returned as remainder
+    , consumeChunks (pack "5\r\nhello\r\n0\r\n\r\nNEXT")
+        @?= (pack "hello", pack "NEXT")
+    ]
+
 applicativeTest :: Test
 applicativeTest =
   "applicativeTest" ~: