diff --git a/Network/HTTP/ReverseProxy.hs b/Network/HTTP/ReverseProxy.hs --- a/Network/HTTP/ReverseProxy.hs +++ b/Network/HTTP/ReverseProxy.hs @@ -17,6 +17,7 @@ , wpsTimeout , wpsSetIpHeader , wpsProcessBody + , wpsUpgradeToRaw , SetIpHeader (..) {- FIXME -- * WAI to Raw @@ -25,12 +26,12 @@ ) where import BasicPrelude +import Debug.Trace import Data.Conduit import Data.Default.Class (def) import qualified Network.Wai as WAI import qualified Network.HTTP.Client as HC import Network.HTTP.Client (BodyReader, brRead) -import qualified Network.HTTP.Client.Internal as HC import Control.Exception (bracketOnError) import Blaze.ByteString.Builder (fromByteString) import Data.Word8 (isSpace, _colon, _cr) @@ -45,9 +46,15 @@ import Control.Concurrent.Lifted (fork, killThread) import Data.Default.Class (Default (..)) import Network.Wai.Logger (showSockAddr) -import Blaze.ByteString.Builder (Builder) import qualified Data.Set as Set import Data.IORef +#if MIN_VERSION_wai(2, 1, 0) +import qualified Data.ByteString.Lazy as L +import Control.Concurrent.Async (concurrently) +import Blaze.ByteString.Builder (Builder, toLazyByteString) +#else +import Blaze.ByteString.Builder (Builder) +#endif -- | Host\/port combination to which we want to proxy. data ProxyDest = ProxyDest @@ -164,6 +171,14 @@ -- ^ Post-process the response body returned from the host. -- -- Since 0.2.1 + , wpsUpgradeToRaw :: WAI.Request -> Bool + -- ^ Determine if the request should be upgraded to a raw proxy connection, + -- as is needed for WebSockets. Requires WAI 2.1 or higher and a WAI + -- handler with raw response support (e.g., Warp) to work. + -- + -- Default: check if the upgrade header is websocket. + -- + -- Since 0.3.1 } -- | How to set the X-Real-IP request header. @@ -179,8 +194,76 @@ , wpsTimeout = Nothing , wpsSetIpHeader = SIHFromSocket , wpsProcessBody = const Nothing + , wpsUpgradeToRaw = \req -> + traceShow ((CI.mk <$> lookup "upgrade" (WAI.requestHeaders req)) == Just "websocket") $ + (CI.mk <$> lookup "upgrade" (WAI.requestHeaders req)) == Just "websocket" } +tryWebSockets :: WaiProxySettings -> ByteString -> Int -> WAI.Request -> IO WAI.Response -> IO WAI.Response +#if MIN_VERSION_wai(2, 1, 0) +tryWebSockets wps host port req fallback + | wpsUpgradeToRaw wps req = do + putStrLn "here0" + return $ flip WAI.responseRaw backup $ \fromClientBody toClient -> + DCN.runTCPClient settings $ \server -> + let toServer = DCN.appSink server + fromServer = DCN.appSource server + fromClient = do + mapM_ yield $ L.toChunks $ toLazyByteString headers + fromClientBody + headers = renderHeaders req $ fixReqHeaders wps req + in void $ concurrently + (do + putStrLn "here1" + fromClient $$ toServer) + (do + putStrLn "here2" + fromServer $$ toClient) + | otherwise = fallback + where + backup = WAI.responseLBS HT.status500 [("Content-Type", "text/plain")] + "http-reverse-proxy detected WebSockets request, but server does not support responseRaw" + settings = DCN.clientSettings port host + +renderHeaders :: WAI.Request -> HT.RequestHeaders -> Builder +renderHeaders req headers + = fromByteString (WAI.requestMethod req) + <> fromByteString " " + <> fromByteString (WAI.rawPathInfo req) + <> fromByteString (WAI.rawQueryString req) + <> (if WAI.httpVersion req == HT.http11 + then fromByteString " HTTP/1.1" + else fromByteString " HTTP/1.0") + <> mconcat (map goHeader headers) + <> fromByteString "\r\n\r\n" + where + goHeader (x, y) + = fromByteString "\r\n" + <> fromByteString (CI.original x) + <> fromByteString ": " + <> fromByteString y +#else +tryWebSockets _ _ _ _ = id +#endif + +strippedHeaders :: Set HT.HeaderName +strippedHeaders = Set.fromList + ["content-length", "transfer-encoding", "accept-encoding", "content-encoding"] + +fixReqHeaders :: WaiProxySettings -> WAI.Request -> HT.RequestHeaders +fixReqHeaders wps req = + addXRealIP $ filter (\(key, _) -> not $ key `Set.member` strippedHeaders) + $ WAI.requestHeaders req + where + addXRealIP = + case wpsSetIpHeader wps of + SIHFromSocket -> (("X-Real-IP", S8.pack $ showSockAddr $ WAI.remoteHost req):) + SIHFromHeader -> + case lookup "x-real-ip" (WAI.requestHeaders req) <|> lookup "X-Forwarded-For" (WAI.requestHeaders req) of + Nothing -> id + Just ip -> (("X-Real-IP", ip):) + SIHNone -> id + waiProxyToSettings :: (WAI.Request -> IO WaiProxyResponse) -> WaiProxySettings -> HC.Manager @@ -194,22 +277,14 @@ WPRModifiedRequest req pd -> Right (pd, req) case edest of Left response -> return response - Right (ProxyDest host port, req) -> do + Right (ProxyDest host port, req) -> tryWebSockets wps host port req $ do let req' = def { HC.method = WAI.requestMethod req , HC.host = host , HC.port = port , HC.path = WAI.rawPathInfo req , HC.queryString = WAI.rawQueryString req - , HC.requestHeaders = filter (\(key, _) -> not $ key `Set.member` strippedHeaders) $ - (case wpsSetIpHeader wps of - SIHFromSocket -> (("X-Real-IP", S8.pack $ showSockAddr $ WAI.remoteHost req):) - SIHFromHeader -> - case lookup "x-real-ip" (WAI.requestHeaders req) <|> lookup "X-Forwarded-For" (WAI.requestHeaders req) of - Nothing -> id - Just ip -> (("X-Real-IP", ip):) - SIHNone -> id) - $ WAI.requestHeaders req + , HC.requestHeaders = fixReqHeaders wps req , HC.requestBody = body , HC.redirectCount = 0 , HC.checkStatus = \_ _ _ -> Nothing @@ -243,8 +318,6 @@ , filter (\(key, _) -> not $ key `Set.member` strippedHeaders) $ HC.responseHeaders res , src $= conduit ) - where - strippedHeaders = Set.fromList ["content-length", "transfer-encoding", "accept-encoding", "content-encoding"] -- | Get the HTTP headers for the first request on the stream, returning on -- consumed bytes as leftovers. Has built-in limits on how many bytes it will diff --git a/http-reverse-proxy.cabal b/http-reverse-proxy.cabal --- a/http-reverse-proxy.cabal +++ b/http-reverse-proxy.cabal @@ -1,5 +1,5 @@ name: http-reverse-proxy -version: 0.3.0.1 +version: 0.3.1 synopsis: Reverse proxy HTTP requests, either over raw sockets or with WAI description: Provides a simple means of reverse-proxying HTTP requests. The raw approach uses the same technique as leveraged by keter, whereas the WAI approach performs full request/response parsing via WAI and http-conduit. homepage: https://github.com/fpco/http-reverse-proxy @@ -33,6 +33,7 @@ , wai-logger , resourcet , containers + , async test-suite test type: exitcode-stdio-1.0 diff --git a/test/main.hs b/test/main.hs --- a/test/main.hs +++ b/test/main.hs @@ -12,11 +12,14 @@ import qualified Data.ByteString.Char8 as S8 import qualified Data.ByteString.Lazy.Char8 as L8 import Data.Conduit (Flush (..), await, runResourceT, - yield, ($$+-)) + yield, ($$+-), (=$), ($$)) import Data.Conduit.Network (HostPreference (HostIPv4, HostAny), ServerSettings, bindPort, runTCPServer, serverAfterBind, - serverSettings) + serverSettings, runTCPClient, clientSettings, appSource, appSink) +import qualified Data.Conduit.Binary as CB +import Data.Char (toUpper) +import qualified Data.Conduit.List as CL import qualified Data.Conduit.Network import qualified Data.IORef as I import qualified Network.HTTP.Conduit as HC @@ -24,14 +27,19 @@ WaiProxyResponse (..), defaultOnExc, rawProxyTo, waiProxyTo{- FIXME, waiToRaw-}) -import Network.HTTP.Types (status200) +import Network.HTTP.Types (status200, status500) import Network.Socket (sClose) import Network.Wai (responseSource, rawPathInfo, responseLBS) import qualified Network.Wai import Network.Wai.Handler.Warp (defaultSettings, runSettings, +#if MIN_VERSION_warp(2, 1, 0) + setBeforeMainLoop, + setPort) +#else settingsBeforeMainLoop, settingsPort) +#endif import System.IO.Unsafe (unsafePerformIO) import System.Timeout.Lifted (timeout) import Test.Hspec (describe, hspec, it, shouldBe) @@ -54,12 +62,22 @@ port <- getPort baton <- newEmptyMVar bracket - (forkIO $ runSettings defaultSettings - { settingsPort = port - , settingsBeforeMainLoop = putMVar baton () - } app `onException` putMVar baton ()) + (forkIO $ runSettings (settings port baton) + app `onException` putMVar baton ()) killThread (const $ takeMVar baton >> f port) + where +#if MIN_VERSION_warp(2, 1, 0) + settings port baton + = setPort port + $ setBeforeMainLoop (putMVar baton ()) + defaultSettings +#else + settings port baton = defaultSettings + { settingsPort = port + , settingsBeforeMainLoop = putMVar baton () + } +#endif withCApp :: Data.Conduit.Network.Application IO -> (Int -> IO ()) -> IO () withCApp app f = do @@ -135,6 +153,19 @@ $ Network.Wai.KnownLength $ fromIntegral $ S.length body) +#endif +#if MIN_VERSION_warp(2, 1, 0) + it "upgrade to raw" $ + let app _ = return $ flip Network.Wai.responseRaw fallback $ \src sink -> + src $$ CL.iterM print =$ CL.map (S8.map toUpper) =$ sink + fallback = responseLBS status500 [] "fallback used" + in withMan $ \manager -> + withWApp app $ \port1 -> + withWApp (waiProxyTo (const $ return $ WPRProxyDest $ ProxyDest "127.0.0.1" port1) defaultOnExc manager) $ \port2 -> + runTCPClient (clientSettings port2 "127.0.0.1") $ \ad -> do + yield "GET / HTTP/1.1\r\nUpgrade: websockET\r\n\r\n" $$ appSink ad + yield "hello" $$ appSink ad + (appSource ad $$ CB.take 5) >>= (`shouldBe` "HELLO") #endif {- FIXME describe "waiToRaw" $ do