version 1.8.51.8.5

5 files changed, 110 insertions, 34 deletions

diff --git a/Network/HTTP/Conduit.hs b/Network/HTTP/Conduit.hs
index 8683095..f472736 100644
--- a/Network/HTTP/Conduit.hs
+++ b/Network/HTTP/Conduit.hs
@@ -140,7 +140,6 @@ module Network.HTTP.Conduit
     ) where
 
 import qualified Data.ByteString as S
-import qualified Data.ByteString.Char8 as S8
 import qualified Data.ByteString.Lazy as L
 
 import qualified Network.HTTP.Types as W
@@ -149,15 +148,15 @@ import Data.Default (def)
 import Control.Exception.Lifted (throwIO)
 import Control.Monad ((<=<))
 import Control.Monad.IO.Class (MonadIO (liftIO))
-import Control.Monad.Trans.Control (MonadBaseControl)
+import Control.Monad.Trans.Resource
+import Control.Monad.Trans.State (get, put, evalStateT)
+import Control.Monad.Trans (lift)
 
 import Control.Exception (fromException, toException)
 import qualified Data.Conduit as C
 import qualified Data.Conduit.Binary as CB
 import qualified Data.Conduit.Internal as CI
-import Data.Conduit.List (sinkNull)
 import Data.Conduit.Blaze (builderToByteString)
-import Data.Conduit (MonadResource)
 import Control.Exception.Lifted (try, SomeException)
 
 import Data.Time.Clock
@@ -167,6 +166,7 @@ import Network.HTTP.Conduit.Response
 import Network.HTTP.Conduit.Manager
 import Network.HTTP.Conduit.ConnInfo
 import Network.HTTP.Conduit.Cookies
+import Network.HTTP.Conduit.Internal (httpRedirect)
 
 -- | The most low-level function for initiating an HTTP request.
 --
@@ -196,7 +196,7 @@ http req0 manager = do
     res@(Response status _version hs body) <-
         if redirectCount req0 == 0
             then httpRaw req0 manager
-            else go (redirectCount req0) req0 def []
+            else go (redirectCount req0) req0 def
     case checkStatus req0 status hs of
         Nothing -> return res
         Just exc -> do
@@ -212,32 +212,20 @@ http req0 manager = do
                         return exc
             liftIO $ throwIO exc'
   where
-    go (-1) _ _ ress = liftIO . throwIO . TooManyRedirects =<< mapM lbsResponse ress
-    go count req'' cookie_jar'' ress = do
+    go count req''' cookie_jar''' = (`evalStateT` cookie_jar''') $
+      httpRedirect
+      count
+      (\req'' -> do
+        cookie_jar'' <- get
         now <- liftIO getCurrentTime
         let (req', cookie_jar') = insertCookiesIntoRequest req'' (evictExpiredCookies cookie_jar'' now) now
-        res <- httpRaw req' manager
+        res <- lift $ httpRaw req' manager
         let (cookie_jar, _) = updateCookieJar res req' now cookie_jar'
-        case getRedirectedRequest req' (responseHeaders res) (W.statusCode (responseStatus res)) of
-            Just req -> do
-                -- Allow the original connection to return to the
-                -- connection pool immediately by flushing the body.
-                -- If the response body is too large, don't flush, but
-                -- instead just close the connection.
-                let maxFlush = 1024
-                    readMay bs =
-                        case S8.readInt bs of
-                            Just (i, bs') | S.null bs' -> Just i
-                            _ -> Nothing
-                    sink =
-                        case lookup "content-length" (responseHeaders res) >>= readMay of
-                            Just i | i > maxFlush -> return ()
-                            _ -> CB.isolate maxFlush C.=$ sinkNull
-                responseBody res C.$$+- sink
-
-                -- And now perform the actual redirect
-                go (count - 1) req cookie_jar (res:ress)
-            Nothing -> return res
+        put cookie_jar
+        let mreq = getRedirectedRequest req' (responseHeaders res) (W.statusCode (responseStatus res))
+        return (res, mreq))
+      lift
+      req'''
 
 -- | Get a 'Response' without any redirect following.
 httpRaw
diff --git a/Network/HTTP/Conduit/Internal.hs b/Network/HTTP/Conduit/Internal.hs
index e17fea2..a73e874 100644
--- a/Network/HTTP/Conduit/Internal.hs
+++ b/Network/HTTP/Conduit/Internal.hs
@@ -1,5 +1,59 @@
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE RankNTypes #-}
 module Network.HTTP.Conduit.Internal
     ( module Network.HTTP.Conduit.Parser
+    , getUri
+    , setUri
+    , setUriRelative
+    , httpRedirect
     ) where
 
 import Network.HTTP.Conduit.Parser
+
+import qualified Data.ByteString as S
+import qualified Data.ByteString.Char8 as S8
+
+import Control.Exception.Lifted (throwIO)
+import Control.Monad.Trans.Resource
+
+import qualified Data.Conduit as C
+import qualified Data.Conduit.Binary as CB
+import Data.Conduit.List (sinkNull)
+
+import Network.HTTP.Conduit.Request
+import Network.HTTP.Conduit.Response
+
+-- | Redirect loop
+httpRedirect
+     :: (MonadBaseControl IO m, MonadResource m, Monad m1)
+     => Int -- ^ 'redirectCount'
+     -> (Request m1 -> m (Response (C.ResumableSource m1 S.ByteString), Maybe (Request m1))) -- ^ function which performs a request and returns a response, and possibly another request if there's a redirect.
+     -> (forall a. m1 a -> m a) -- ^ 'liftResourceT'
+     -> Request m1
+     -> m (Response (C.ResumableSource m1 S.ByteString))
+httpRedirect count0 http' lift' req0 = go count0 req0 []
+  where
+    go (-1) _ ress = throwIO . TooManyRedirects =<< lift' (mapM lbsResponse ress)
+    go count req' ress = do
+        (res, mreq) <- http' req'
+        case mreq of
+            Just req -> do
+                -- Allow the original connection to return to the
+                -- connection pool immediately by flushing the body.
+                -- If the response body is too large, don't flush, but
+                -- instead just close the connection.
+                let maxFlush = 1024
+                    readMay bs =
+                        case S8.readInt bs of
+                            Just (i, bs') | S.null bs' -> Just i
+                            _ -> Nothing
+                    sink =
+                        case lookup "content-length" (responseHeaders res) >>= readMay of
+                            Just i | i > maxFlush -> return ()
+                            _ -> CB.isolate maxFlush C.=$ sinkNull
+                lift' $ responseBody res C.$$+- sink
+
+                -- And now perform the actual redirect
+                go (count - 1) req (res:ress)
+            Nothing -> return res
diff --git a/Network/HTTP/Conduit/Request.hs b/Network/HTTP/Conduit/Request.hs
index 8a9a3bb..be530d5 100644
--- a/Network/HTTP/Conduit/Request.hs
+++ b/Network/HTTP/Conduit/Request.hs
@@ -11,6 +11,8 @@ module Network.HTTP.Conduit.Request
     , Proxy (..)
     , parseUrl
     , setUriRelative
+    , getUri
+    , setUri
     , browserDecompress
     , HttpException (..)
     , alwaysDecompress
diff --git a/http-conduit.cabal b/http-conduit.cabal
index eb9280a..89c6c0a 100644
--- a/http-conduit.cabal
+++ b/http-conduit.cabal
@@ -1,5 +1,5 @@
 name:            http-conduit
-version:         1.8.4.5
+version:         1.8.5
 license:         BSD3
 license-file:    LICENSE
 author:          Michael Snoyman <michael@snoyman.com>
diff --git a/test/main.hs b/test/main.hs
index 607b1df..9ea57af 100644
--- a/test/main.hs
+++ b/test/main.hs
@@ -3,28 +3,31 @@
 import Test.Hspec
 import qualified Data.ByteString as S
 import qualified Data.ByteString.Char8 as S8
+import qualified Data.ByteString.Lazy.Char8 as L8
 import Test.HUnit
 import Network.Wai hiding (requestBody)
-import qualified Network.Wai
+import qualified Network.Wai as Wai
 import Network.Wai.Handler.Warp (run)
 import Network.HTTP.Conduit
 import Data.ByteString.Base64 (encode)
 import Control.Concurrent (forkIO, killThread, threadDelay)
 import Network.HTTP.Types
 import Control.Exception.Lifted (try, SomeException)
+import qualified Control.Exception as E (catch)
 import Network.HTTP.Conduit.ConnInfo
 import Network (withSocketsDo)
 import CookieTest (cookieTest)
 import Data.Conduit.Network (runTCPServer, serverSettings, HostPreference (HostAny), appSink, appSource)
 import Data.Conduit (($$), yield, Flush (Chunk))
+import Control.Monad (void, forever)
 import Control.Monad.Trans.Resource (register)
 import Control.Monad.IO.Class (liftIO)
-import Control.Monad (forever)
 import Data.ByteString.UTF8 (fromString)
 import Data.Conduit.List (sourceList)
 import Data.CaseInsensitive (mk)
-import Data.List (partition)
+import Data.List (partition, lookup)
 import qualified Data.Conduit.List as CL
+import qualified Data.Text as T
 import qualified Data.Text.Encoding as TE
 import qualified Data.ByteString.Lazy as L
 import Blaze.ByteString.Builder (fromByteString)
@@ -34,6 +37,21 @@ app req =
     case pathInfo req of
         [] -> return $ responseLBS status200 [] "homepage"
         ["cookies"] -> return $ responseLBS status200 [tastyCookie] "cookies"
+        ["cookie_redir1"] -> return $ responseLBS status303 [tastyCookie, (hLocation, "/checkcookie")] ""
+        ["checkcookie"] -> return $ case lookup hCookie $ Wai.requestHeaders req of
+                                Just "flavor=chocolate-chip" -> responseLBS status200 [] "nom-nom-nom"
+                                _ -> responseLBS status412 [] "Baaaw where's my chocolate?"
+        ["infredir", i'] ->
+            let i = read $ T.unpack i'
+            in return $ responseLBS status303
+                    [(hLocation, S.append "/infredir/" $ S8.pack $ show $ i+1)]
+                    (L8.pack $ show i)
+        ["infredirrepeat", i'] ->
+            let i = read $ T.unpack i'
+            in return $ responseLBS status303
+                    [(hLocation, S.append "/infredirrepeat/" $ S8.pack $ show $ i+1)
+                    ,(hContentLength, "2048")]
+                    (L8.pack $ take 2048 $ unwords $ repeat $ show i)
         _ -> return $ responseLBS status404 [] "not found"
 
     where tastyCookie = (mk (fromString "Set-Cookie"), fromString "flavor=chocolate-chip;")
@@ -67,6 +85,13 @@ main = withSocketsDo $ hspec $ do
                     (setCookieHeaders, _) = partition ((== setCookie) . fst) headers
                 liftIO $ assertBool "response contains a 'set-cookie' header" $ length setCookieHeaders > 0
             killThread tid
+        it "redirects set cookies" $ do
+            tid <- forkIO $ run 13010 app
+            request <- parseUrl "http://127.0.0.1:13010/cookie_redir1"
+            withManager $ \manager -> do
+                _ <- register $ killThread tid
+                Response _ _ _ body <- httpLbs request manager
+                liftIO $ body @?= "nom-nom-nom"
     describe "manager" $ do
         it "closes all connections" $ do
             clearSocketsList
@@ -120,7 +145,14 @@ main = withSocketsDo $ hspec $ do
                     , ("hello%20world", "hello world")
                     , ("hello%20world%3f%23", "hello world?#")
                     ]
-
+        it "TooManyRedirects: redirect request body is preserved" $ do
+            tid <- forkIO $ run 13009 app
+            let Just req = parseUrl "http://127.0.0.1:13009/infredir/0"
+            let go (res, i) = liftIO $ responseBody res @?= (L8.pack $ show i)
+            E.catch (withManager $ \manager -> do
+                void $ register $ killThread tid
+                void $ http req{redirectCount=5} manager) $
+                \(TooManyRedirects redirs) -> mapM_ go (zip redirs [5,4..0])
     describe "chunked request body" $ do
         it "works" $ do
             tid <- forkIO echo
@@ -218,7 +250,7 @@ redir =
 
 echo :: IO ()
 echo = run 13007 $ \req -> do
-    bss <- Network.Wai.requestBody req $$ CL.consume
+    bss <- Wai.requestBody req $$ CL.consume
     return $ responseLBS status200 [] $ L.fromChunks bss
 
 noStatusMessage :: IO ()