Fix chunking, tests now pass

fpco · Jan 13, 2024 · fac5e72 · fac5e72
1 parent 9772ca0
commit fac5e72
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 35 deletions.
diff --git a/Network/HTTP/ReverseProxy.hs b/Network/HTTP/ReverseProxy.hs
@@ -38,39 +38,38 @@ module Network.HTTP.ReverseProxy
     -}
     ) where
 
-import           Blaze.ByteString.Builder       (Builder, fromByteString
-                                                , toByteString, toLazyByteString)
-import           Control.Applicative            ((<$>), (<|>))
-import           Control.Monad                  (forM_, unless)
-import           Data.ByteString                (ByteString)
-import qualified Data.ByteString                as S
-import qualified Data.ByteString.Char8          as S8
-import qualified Data.ByteString.Lazy           as L
-import qualified Data.CaseInsensitive           as CI
+import           Blaze.ByteString.Builder             (Builder, fromByteString, toByteString, toLazyByteString)
+import           Control.Applicative                  ((<$>), (<|>))
+import           Control.Monad                        (unless, when)
+import           Data.ByteString                      (ByteString)
+import qualified Data.ByteString                      as S
+import           Data.ByteString.Builder.HTTP.Chunked (chunkedTransferEncoding, chunkedTransferTerminator)
+import qualified Data.ByteString.Char8                as S8
+import qualified Data.ByteString.Lazy                 as L
+import qualified Data.CaseInsensitive                 as CI
 import           Data.Conduit
-import qualified Data.Conduit.List              as CL
-import qualified Data.Conduit.Network           as DCN
-import           Data.Functor.Identity          (Identity (..))
+import qualified Data.Conduit.List                    as CL
+import qualified Data.Conduit.Network                 as DCN
+import           Data.Functor.Identity                (Identity (..))
 import           Data.IORef
-import           Data.List.NonEmpty             (NonEmpty (..))
-import qualified Data.List.NonEmpty             as NE
-import           Data.Maybe                     (fromMaybe, listToMaybe)
-import           Data.Monoid                    (mappend, mconcat, (<>))
-import           Data.Set                       (Set)
-import qualified Data.Set                       as Set
-import           Data.Streaming.Network         (AppData, readLens)
-import qualified Data.Text                      as T
-import qualified Data.Text.Encoding             as TE
-import qualified Data.Text.Lazy                 as TL
-import qualified Data.Text.Lazy.Encoding        as TLE
-import           Data.Word8                     (isSpace, _colon, _cr)
-import           GHC.Generics                   (Generic)
-import           Network.HTTP.Client            (BodyReader, brRead)
-import qualified Network.HTTP.Client            as HC
-import qualified Network.HTTP.Types             as HT
-import qualified Network.Wai                    as WAI
-import           Network.Wai.Logger             (showSockAddr)
-import           UnliftIO                       (MonadIO, liftIO, MonadUnliftIO, timeout, SomeException, try, bracket, concurrently_)
+import           Data.List.NonEmpty                   (NonEmpty (..))
+import qualified Data.List.NonEmpty                   as NE
+import           Data.Maybe                           (fromMaybe, isJust, listToMaybe)
+import           Data.Monoid                          (mappend, mconcat, (<>))
+import           Data.Set                             (Set)
+import qualified Data.Set                             as Set
+import           Data.Streaming.Network               (AppData, readLens)
+import qualified Data.Text                            as T
+import qualified Data.Text.Encoding                   as TE
+import           Data.Word8                           (isSpace, _colon, _cr)
+import           GHC.Generics                         (Generic)
+import           Network.HTTP.Client                  (BodyReader, brRead)
+import qualified Network.HTTP.Client                  as HC
+import qualified Network.HTTP.Types                   as HT
+import qualified Network.HTTP.Types.Header            as H
+import qualified Network.Wai                          as WAI
+import           Network.Wai.Logger                   (showSockAddr)
+import           UnliftIO                             (MonadIO, liftIO, MonadUnliftIO, timeout, SomeException, try, bracket, concurrently_)
 
 -- | Host\/port combination to which we want to proxy.
 data ProxyDest = ProxyDest
@@ -447,20 +446,32 @@ waiProxyToSettings getDest wps' manager req0 sendResponse = do
                           sendToClient $ L.toStrict $ toLazyByteString $
                             fromByteString (S8.pack $ show (HC.responseVersion res)) <> " " <> fromByteString (S8.pack $ show code) <> " " <> fromByteString message <> "\r\n"
 
-                          let headers = (filter (\(key, v) -> not (key `Set.member` strippedHeaders) ||
-                                                                   key == "content-length" && (noChunked || v == "0"))
-                                        (HC.responseHeaders res))
+                          -- Handle HTTP chunking, just as Warp does for WAI.responseStream
+                          let requestIsChunked = not noChunked
+                          let responseHasLength = isJust (lookup "content-length" (HC.responseHeaders res))
+                          let needsChunked = requestIsChunked && not responseHasLength
+
+                          let headers' = (filter (\(key, v) -> not (key `Set.member` strippedHeaders) ||
+                                                                    key == "content-length" && (noChunked || v == "0"))
+                                         (HC.responseHeaders res))
+                          let headers
+                                | needsChunked = (H.hTransferEncoding, "chunked") : headers'
+                                | otherwise = headers
                           sendToClient $ L.toStrict $ toLazyByteString $
                               mconcat [renderHeader h | h <- headers]
                               <> "\r\n"
 
                           -- It may look strange that we don't handle 'Flush' here, but 'Flush' is not actually used anywhere
                           -- except at the end of the stream in the conduit above.
+                          let sendChunk
+                                | needsChunked = sendToClient . toByteString . chunkedTransferEncoding
+                                | otherwise = sendToClient . toByteString
                           runConduit $ src .| conduit .| CL.mapM_ (\mb ->
                               case mb of
                                   Flush -> return ()
-                                  Chunk b -> sendToClient (toByteString b)
+                                  Chunk b -> sendChunk b
                               )
+                          when needsChunked $ sendToClient (toByteString chunkedTransferTerminator)
 
   where
     renderHeader :: HT.Header -> Builder

diff --git a/http-reverse-proxy.cabal b/http-reverse-proxy.cabal
@@ -20,6 +20,7 @@ library
     buildable: False
   build-depends:       base                   >= 4.11   && < 5
                      , text                   >= 0.11
+                     , bsb-http-chunked       >= 0.0.0.4
                      , bytestring             >= 0.9
                      , case-insensitive       >= 0.4
                      , http-types             >= 0.6