diff --git a/Network/HTTP/Proxy.hs b/Network/HTTP/Proxy.hs --- a/Network/HTTP/Proxy.hs +++ b/Network/HTTP/Proxy.hs @@ -28,12 +28,12 @@ -- module Network.HTTP.Proxy - ( runProxy - , runProxySettings + ( runProxy + , runProxySettings - , Settings (..) - , defaultSettings - ) + , Settings (..) + , defaultSettings + ) where import Prelude hiding (catch, lines) @@ -42,7 +42,7 @@ import qualified Data.ByteString as S import qualified Data.ByteString.Unsafe as SU import qualified Data.ByteString.Char8 as B -import qualified Data.ByteString.Lazy as L + import Network ( PortID(..) ) import Network.Socket ( accept, Family (..) @@ -68,7 +68,7 @@ import Data.Typeable (Typeable) -import Data.Enumerator (($$), (=$), (>>==)) +import Data.Enumerator hiding (filter, head, foldl', map) import qualified Data.Enumerator as E import qualified Data.Enumerator.List as EL import qualified Data.Enumerator.Binary as EB @@ -76,15 +76,18 @@ import Blaze.ByteString.Builder (copyByteString, Builder, toByteString, fromByteString) import Blaze.ByteString.Builder.Char8 (fromChar, fromShow) -import Data.Monoid (mappend, mconcat) +import Data.Monoid (mappend) -import Control.Monad.IO.Class (liftIO) +import Control.Monad.IO.Class (liftIO, MonadIO) +import Control.Monad.Trans.Class (lift) import qualified Network.HTTP.Proxy.Timeout as T import Data.List (delete, foldl') import Control.Monad (forever, when) import qualified Network.HTTP.Types as H import qualified Data.CaseInsensitive as CI import System.IO (hPutStrLn, stderr) +import Numeric (readDec) +import Data.Int (Int64) #if WINDOWS import Control.Concurrent (threadDelay) @@ -170,15 +173,15 @@ shutdownX :: Socket -> ShutdownCmd -> IO () shutdownX s ShutdownReceive = do - when verboseSockets $ putStrLn ("shutdown " ++ show s++" ShutdownReceive") + when verboseSockets $ putStrLn ("shutdown " ++ show s ++ " ShutdownReceive") shutdown s ShutdownReceive shutdownX s ShutdownSend = do - when verboseSockets $ putStrLn ("shutdown " ++ show s++" ShutdownSend") + when verboseSockets $ putStrLn ("shutdown " ++ show s ++ " ShutdownSend") shutdown s ShutdownSend shutdownX s ShutdownBoth = do - when verboseSockets $ putStrLn ("shutdown " ++ show s++" ShutdownBoth") + when verboseSockets $ putStrLn ("shutdown " ++ show s ++ " ShutdownBoth") shutdown s ShutdownBoth mkHeaders :: Monad m @@ -217,8 +220,8 @@ proxyPlain req _ | requestMethod req == "CONNECT" -> case B.split ':' (rawPathInfo req) of - [h, p] -> proxyConnect th tm onException conn h (read $ B.unpack p) req - _ -> failRequest th conn req "Bad request" ("Bad request '" `mappend` rawPathInfo req `mappend` "'.") + [h, p] -> proxyConnect th tm onException conn h (readDecimal $ B.unpack p) req + _ -> failRequest th conn req "Bad request" ("Bad request '" `mappend` rawPathInfo req `mappend` "'.") _ | otherwise -> failRequest th conn req "Unknown request" ("Unknown request '" `mappend` rawPathInfo req `mappend` "'.") @@ -226,31 +229,30 @@ proxyPlain req = do let urlStr = "http://" `mappend` serverName req `mappend` rawPathInfo req - `mappend` H.renderQuery True (queryString req) + `mappend` rawQueryString req close = let hasClose hdrs = (== "close") . CI.mk <$> lookup "connection" hdrs mClose = hasClose (requestHeaders req) -- RFC2616: Connection defaults to Close in HTTP/1.0 and Keep-Alive in HTTP/1.1 defaultClose = httpVersion req == H.HttpVersion 1 0 in fromMaybe defaultClose mClose - outHdrs = [(n,v) | (n,v) <- requestHeaders req, n `notElem` [ - "Host", "Content-Length" - ]] + outHdrs = [(n,v) | (n,v) <- requestHeaders req, n /= "Host"] liftIO $ putStrLn $ B.unpack (requestMethod req) ++ " " ++ B.unpack urlStr let contentLength = if requestMethod req == "GET" - then 0 - else read . B.unpack . fromMaybe "0" . lookup "content-length" . requestHeaders $ req - postBody <- mconcat . L.toChunks <$> EB.take contentLength - -- This loads it all into memory at once. - -- TO DO: Stream via iteratee - let enumPostBody = E.enumList 1 [fromByteString postBody] + then 0 + else readDecimal . B.unpack . fromMaybe "0" . lookup "content-length" . requestHeaders $ req + url <- - (\url -> url { HE.method = requestMethod req, - HE.requestHeaders = outHdrs, - HE.rawBody = True, - HE.requestBody = HE.RequestBodyEnum (fromIntegral contentLength) enumPostBody }) - <$> liftIO (HE.parseUrl (B.unpack urlStr)) - close' <- liftIO $ E.run_ $ HE.http url (handleHttpReply close) mgr + (\u -> u { HE.method = requestMethod req, + HE.requestHeaders = outHdrs, + HE.rawBody = True, + HE.requestBody = HE.RequestBodyEnum contentLength + $ joinE (enumIteratee contentLength lazyTakeMax) + $ EL.map fromByteString + }) + <$> lift (HE.parseUrl (B.unpack urlStr)) + + close' <- E.run_ $ HE.http url (handleHttpReply close) mgr if close' then return Nothing else serveConnection' @@ -265,6 +267,49 @@ mkHeaders (httpVersion req) status hdrs' $$ iterSocket th conn close return remoteClose +-- Create an Enumerator from an Iteratee. +-- The first parameter is the number of bytes to be pulled from the Iteratee. +-- The second is an Iteratee that can pull the data from the source in chunks. +-- The return value is an Enumerator that operates inside the Iteratee monad. +enumIteratee :: MonadIO m => Int64 + -> (Int -> Iteratee ByteString m ByteString) + -> Enumerator ByteString (Iteratee ByteString m) c +enumIteratee maxlen takeMax = inner 0 + where + blockLen = 32768 + inner count (Continue k) + | count >= maxlen = k (Chunks []) + | count + blockLen <= maxlen = do + bs <- lift $ takeMax $ fromIntegral blockLen + if B.null bs + then k EOF + else k (Chunks [bs]) >>== inner (count + blength bs) + | otherwise = do + bs <- lift $ takeMax $ fromIntegral (maxlen - count) + if B.null bs + then k EOF + else k (Chunks [bs]) >>== inner (count + blength bs) + inner _ step = returnI step + blength = fromIntegral . B.length + +-- Take up to maxlen bytes of data as lazily as possible. Avoid splitting and +-- joining ByteStrings as much as possible. +lazyTakeMax :: Monad m => Int -> Iteratee ByteString m ByteString +lazyTakeMax maxlen | maxlen <= 0 = return B.empty +lazyTakeMax maxlen = continue step + where + step (Chunks []) = continue step + + step (Chunks (x:xs)) + | B.length x < maxlen = yield x (Chunks xs) + | otherwise = + let (start, extra) = B.splitAt maxlen x + in yield start (Chunks (extra:xs)) + + step EOF = yield B.empty EOF + +-------------------------------------------------------------------------------- + failRequest :: T.Handle -> Socket -> Request -> ByteString -> ByteString -> E.Iteratee ByteString IO (Maybe (ThreadId, Socket)) failRequest th conn req headerMsg bodyMsg = do EB.isolate 0 =$ @@ -297,9 +342,9 @@ Left errorMsg -> failRequest th conn req errorMsg ("PROXY FAILURE\r\n" `mappend` errorMsg) -connectTo :: HostName -- Hostname - -> PortID -- Port Identifier - -> IO Socket -- Connected Socket +connectTo :: HostName -- Hostname + -> PortID -- Port Identifier + -> IO Socket -- Connected Socket connectTo hostname (Service serv) = connect' hostname serv connectTo hostname (PortNumber port) = connect' hostname (show port) connectTo _ (UnixSocket _) = error "Cannot connect to a UnixSocket" @@ -315,12 +360,12 @@ where tryToConnect addr = bracketOnError - (socket (addrFamily addr) (addrSocketType addr) (addrProtocol addr)) - sClose -- only done if there's an error - (\sock -> do - connect sock (addrAddress addr) - return sock - ) + (socket (addrFamily addr) (addrSocketType addr) (addrProtocol addr)) + sClose -- only done if there's an error + (\sock -> do + connect sock (addrAddress addr) + return sock + ) -- Returns the first action from a list which does not throw an exception. -- If all the actions throw exceptions (and the list of actions is not empty), @@ -461,17 +506,18 @@ liftIO $ shutdownX sock ShutdownReceive `catch` \(exc :: IOException) -> - putStrLn $ "couldn't shutdown read side of " ++ show sock++": " ++ show exc + putStrLn $ "couldn't shutdown read side of " ++ show sock ++ ": " ++ show exc E.continue k else k (E.Chunks [bs]) >>== inner inner step = E.returnI step ------ The functions below are not warp-specific and could be split out into a --separate package. -iterSocket :: T.Handle +iterSocket :: MonadIO m + => T.Handle -> Socket -> Bool -- ^ To close - -> E.Iteratee B.ByteString IO () + -> E.Iteratee B.ByteString m () iterSocket th sock toClose = E.continue step where @@ -518,7 +564,7 @@ Just x -> go x Nothing -> when (go' $ fromException e) - $ hPutStrLn stderr $ show e + $ hPutStrLn stderr $ show e , proxyTimeout = 30 } where @@ -605,3 +651,11 @@ ver = B.pack "Proxy/0.0" server = (key, ver) servers svr = (key, S.concat [svr, " ", ver]) + + +readDecimal :: Num a => String -> a +readDecimal s = + case readDec s of + [] -> 0 + (x, _):_ -> x + diff --git a/http-proxy.cabal b/http-proxy.cabal --- a/http-proxy.cabal +++ b/http-proxy.cabal @@ -1,5 +1,5 @@ Name: http-proxy -Version: 0.0 +Version: 0.0.1 License: BSD3 License-file: LICENSE Author: Michael Snoyman, Stephen Blackheath, Erik de Castro Lopo @@ -15,10 +15,10 @@ Description: http-proxy is a library for writing HTTP and HTTPS proxies. . - Use of the enumerator library provides file streaming from the upstream server - through the proxy to the client program. Memory usage of the proxy depends on - number of files being downlaoded simultaneously, not the size of the files - being downloaded. + Use of the enumerator library provides file streaming via the proxy in both + directions. Memory usage of the proxy scales linearly with the number of + simultaneous connections and is independent of the size of the files being + uploaded or downloaded. . Eventually, features like logging, request re-writing, disk caching etc will be made available via caller provided functions in the Settings data type.