diff --git a/Network/HTTP/Proxy.hs b/Network/HTTP/Proxy.hs index 7a426bf..0da0682 100644 --- a/Network/HTTP/Proxy.hs +++ b/Network/HTTP/Proxy.hs @@ -38,12 +38,14 @@ module Network.HTTP.Proxy , runProxySettings , runProxySettingsSocket , defaultProxySettings + , retryWithRequest ) where import Blaze.ByteString.Builder (fromByteString) import Control.Concurrent.Async (race_) import Control.Exception -- (SomeException, catch, toException) +import Control.Monad ((>=>)) import Data.ByteString.Char8 (ByteString) import Data.Conduit (ConduitT, Flush (..), (.|), mapOutput, runConduit, yield) import Data.Conduit.Network @@ -119,6 +121,9 @@ data Settings = Settings -- ^ A function that allows the request to be modified before being run. Default: 'return . Right'. -- This only works for unencrypted HTTP requests (eg to upgrade the request to HTTPS) because -- HTTPS requests are encrypted. + , proxyHttpResponseModifier :: Request -> Response -> IO Response + -- ^ A function that allows the response to be modified before being sent + -- back to the client. Default: 'const return'. , proxyLogger :: ByteString -> IO () -- ^ A function for logging proxy internal state. Default: 'return ()'. , proxyUpstream :: Maybe UpstreamProxy @@ -155,6 +160,7 @@ defaultProxySettings = Settings , proxyOnException = defaultExceptionResponse , proxyTimeout = 30 , proxyHttpRequestModifier = return . Right + , proxyHttpResponseModifier = const return , proxyLogger = const $ return () , proxyUpstream = Nothing } @@ -165,13 +171,28 @@ defaultExceptionResponse e = [ (HT.hContentType, "text/plain; charset=utf-8") ] $ LBS.fromChunks [BS.pack $ show e] +data RetryException = RetryException Request + deriving Show + +instance Exception RetryException + +-- | For use in 'proxyHttpResponseModifier'. Ignore the current response and perform a new request. +retryWithRequest :: Request -> IO a +retryWithRequest = throwIO . RetryException + -- ----------------------------------------------------------------------------- -- Application == Wai.Request -> (Wai.Response -> IO Wai.ResponseReceived) -> IO Wai.ResponseReceived httpProxyApp :: Settings -> HC.Manager -> Application -httpProxyApp settings mgr wreq respond = do - mwreq <- proxyHttpRequestModifier settings $ proxyRequest wreq - either respond (doUpstreamRequest settings mgr respond . waiRequest wreq) mwreq +httpProxyApp settings mgr wreq respond = httpProxyApp' $ proxyRequest wreq + where + httpProxyApp' req = do + mwreq <- proxyHttpRequestModifier settings req + handle (\(RetryException req') -> httpProxyApp' req') $ case mwreq of + Left resp -> respond' req resp + Right req' -> doUpstreamRequest settings mgr (respond' req') $ waiRequest wreq req' + + respond' req' = proxyHttpResponseModifier settings req' >=> respond doUpstreamRequest :: Settings -> HC.Manager -> (Wai.Response -> IO Wai.ResponseReceived) -> Wai.Request -> IO Wai.ResponseReceived diff --git a/test/test-io.hs b/test/test-io.hs index 959f306..fe0db8c 100644 --- a/test/test-io.hs +++ b/test/test-io.hs @@ -137,6 +137,16 @@ requestTest = describe "Request:" $ do req <- addTestProxy testProxyPort <$> mkGetRequest Http "/whatever" result <- httpRun req "X-Test-Header: Blah" `BS.isInfixOf` resultBS result `shouldBe` True + it "Can add a response header." $ + withTestProxy proxySettingsAddResponseHeader $ \ testProxyPort -> do + req <- addTestProxy testProxyPort <$> mkGetRequest Http "/whatever" + result <- httpRun req + ("X-Test-Header", "Blah") `elem` resultHeaders result `shouldBe` True + it "Can retry a request." $ + withTestProxy proxySettingsRetry $ \ testProxyPort -> do + req <- addTestProxy testProxyPort <$> mkGetRequest Http "/whatever" + result <- httpRun req + resultBS result `shouldBe` "This is another page" it "Can rewrite HTTP to HTTPS." $ withTestProxy proxySettingsHttpsUpgrade $ \ testProxyPort -> do req <- addTestProxy testProxyPort <$> mkGetRequest Http "/secure" @@ -178,6 +188,21 @@ proxySettingsAddHeader = defaultProxySettings } } +proxySettingsAddResponseHeader :: Settings +proxySettingsAddResponseHeader = defaultProxySettings + { proxyHttpResponseModifier = \ _ resp -> return $ Wai.mapResponseHeaders ((CI.mk "X-Test-Header", "Blah") :) resp + } + +proxySettingsRetry :: Settings +proxySettingsRetry = defaultProxySettings + { proxyHttpResponseModifier = \ req _ -> + if BS.isSuffixOf "/another" $ requestPath req + then return $ simpleResponse HT.status200 "This is another page" + else retryWithRequest $ req { requestPath = redirect $ requestPath req } + } + where + redirect path = fst (BS.spanEnd (/= '/') path) `BS.append` "another" + proxySettingsHttpsUpgrade :: Settings proxySettingsHttpsUpgrade = defaultProxySettings { proxyHttpRequestModifier = \ req -> return . Right $ req { requestPath = httpsUpgrade $ requestPath req }