commit 054dba9d9e6d1a6a4f3970e5fea9642114464af5 Author: David Kaloper david@numm.org Date: Thu Oct 17 01:14:17 2013 +0200
fix frame reading complexity --- src/TorDNSEL/ExitTest/Request.hs | 25 +++++++++++----------- src/TorDNSEL/TorControl/Internals.hs | 36 ++++++++++++-------------------- src/TorDNSEL/Util.hsc | 40 +++++++++++++++++++++++++++--------- 3 files changed, 56 insertions(+), 45 deletions(-)
diff --git a/src/TorDNSEL/ExitTest/Request.hs b/src/TorDNSEL/ExitTest/Request.hs index 4634e8d..84f502a 100644 --- a/src/TorDNSEL/ExitTest/Request.hs +++ b/src/TorDNSEL/ExitTest/Request.hs @@ -30,6 +30,7 @@ import Control.Arrow ((***)) import Control.Applicative import Control.Monad import Data.Monoid +import Data.Maybe import qualified Data.ByteString.Char8 as B import Data.Char (isSpace, toLower) import qualified Data.Map as M @@ -67,27 +68,27 @@ createRequest host port cookie = getRequest :: Handle -> IO (Maybe Cookie) getRequest client = CB.sourceHandle client $= CB.isolate maxReqLen $$ do - mh <- getHeaders - case checkHeaders mh of + reqline <- line + hs <- accHeaders [] + case checkHeaders reqline hs of Nothing -> return Nothing Just _ -> Just . Cookie <$> takeC cookieLen
where maxReqLen = 2048 + cookieLen - line = frameC "\r\n" + line = fromMaybe "" <$> frame "\r\n"
- getHeaders = - (,) <$> line - <*> (decodeHeaders <$> muntil B.null line) - where - decodeHeaders = M.fromList . - map ((B.map toLower *** B.dropWhile isSpace . B.tail) - . B.break (== ':')) + accHeaders hs = line >>= \ln -> + if ln == "" then return $ M.fromList hs + else accHeaders (parseHeader ln : hs)
- checkHeaders (reqLine, headers) = do + parseHeader = (B.map toLower *** B.dropWhile isSpace . B.tail) . + B.break (== ':') + + checkHeaders reqline headers = do contentType <- "content-type" `M.lookup` headers contentLen <- readInt =<< "content-length" `M.lookup` headers - guard $ reqLine `elem` ["POST / HTTP/1.0", "POST / HTTP/1.1"] + guard $ reqline `elem` ["POST / HTTP/1.0", "POST / HTTP/1.1"] guard $ contentType == "application/octet-stream" guard $ contentLen == cookieLen
diff --git a/src/TorDNSEL/TorControl/Internals.hs b/src/TorDNSEL/TorControl/Internals.hs index 58c64ef..0d299d6 100644 --- a/src/TorDNSEL/TorControl/Internals.hs +++ b/src/TorDNSEL/TorControl/Internals.hs @@ -137,11 +137,10 @@ import Control.Concurrent.Chan (newChan, readChan, writeChan) import Control.Concurrent.MVar (MVar, newMVar, newEmptyMVar, takeMVar, tryPutMVar, withMVar, modifyMVar_) import qualified Control.Exception as E -import Control.Monad (when, unless, liftM, mzero, mplus, forever) +import Control.Monad (when, unless, liftM, mzero, mplus) import Control.Monad.Error (MonadError(..)) import Control.Monad.Fix (fix) import Control.Monad.State (StateT(StateT), get, put, lift, evalStateT) -import Control.Applicative import qualified Data.ByteString.Char8 as B import Data.ByteString (ByteString) import Data.Char (isSpace, isAlphaNum, isDigit, isAlpha, toLower) @@ -830,13 +829,13 @@ startIOManager handle = do startSocketReader :: Handle -> ([Reply] -> IO ()) -> IO ThreadId startSocketReader handle sendRepliesToIOManager = forkLinkIO $ CB.sourceHandle handle $= - repliesC $$ + c_replies $$ CL.mapM_ sendRepliesToIOManager
--- | Conduit taking lines to 'Reply' blocks. -replyC :: Conduit B.ByteString IO [Reply] -replyC = - line0 [] +-- | Stream decoded 'Reply' groups. +c_replies :: Conduit B.ByteString IO [Reply] +c_replies = + frames (B.pack "\r\n") =$= line0 [] where
line0 acc = await >>= return () `maybe` \line -> do @@ -844,13 +843,13 @@ replyC = code' <- either (monadThrow . ProtocolError) return $ parseReplyCode code case () of - _ | typ == B.pack "-" -> line0 (Reply code' text [] : acc) - | typ == B.pack "+" -> line0 . (: acc) . Reply code' text =<< rest [] - | typ == B.pack " " -> do - yield $ reverse (Reply code' text [] : acc) - line0 [] - | otherwise -> monadThrow $ ProtocolError $ - cat "Malformed reply line type " (esc 1 typ) '.' + _ | typ == B.pack "-" -> line0 (acc' []) + | typ == B.pack "+" -> rest [] >>= line0 . acc' + | typ == B.pack " " -> yield (reverse $ acc' []) >> line0 [] + | otherwise -> monadThrow $ + ProtocolError $ cat "Malformed reply line type " (esc 1 typ) '.' + where + acc' xs = Reply code' text xs : acc
rest acc = await >>= \mline -> case mline of @@ -859,15 +858,6 @@ replyC = | line == B.pack "." -> return $ reverse (line:acc) | otherwise -> rest (line:acc)
--- | Conduit taking raw 'ByteString' to 'Reply' blocks. -repliesC :: Conduit B.ByteString IO [Reply] -repliesC = - CB.lines =$= CL.map strip =$= replyC - where - strip bs = case unsnoc bs of - Just (bs', '\r') -> bs' - _ -> bs - -------------------------------------------------------------------------------- -- Data types
diff --git a/src/TorDNSEL/Util.hsc b/src/TorDNSEL/Util.hsc index f866eb9..2c57e0e 100644 --- a/src/TorDNSEL/Util.hsc +++ b/src/TorDNSEL/Util.hsc @@ -65,7 +65,8 @@ module TorDNSEL.Util (
-- * Conduit utilities , takeC - , frameC + , frames + , frame
-- * Network functions , bindUDPSocket @@ -416,15 +417,34 @@ foreign import ccall unsafe "ntohl" ntohl :: Word32 -> Word32 takeC :: Monad m => Int -> C.ConduitM ByteString o m ByteString takeC = fmap (mconcat . BL.toChunks) . CB.take
--- | Take a prefix up to delimiter. --- FIXME This is worst-case quadratic. -frameC :: Monad m => ByteString -> C.ConduitM ByteString o m ByteString -frameC delim = loop $ B.pack "" where - loop acc = C.await >>= - return acc `maybe` \bs -> - case B.breakSubstring delim $ acc <> bs of - (h, t) | B.null t -> loop h - | otherwise -> h <$ C.leftover (B.drop (B.length delim) t) +-- | Take a "frame" - delimited sequence - from the input. +-- Returns 'Nothing' if the delimiter does not appear before the stream ends. +frame :: MonadIO m => ByteString -> C.ConduitM ByteString a m (Maybe ByteString) +frame delim = input $ B.pack "" + where + input front = C.await >>= + (Nothing <$ C.leftover front) `maybe` \bs -> + + let (front', bs') = (<> bs) `second` + B.splitAt (B.length front - d_len + 1) front + + in case B.breakSubstring delim bs' of + (part, rest) | B.null rest -> input (front' <> bs') + | otherwise -> do + leftover $ B.drop d_len rest + return $ Just $ front' <> part + + d_len = B.length delim + +-- | Stream delimited chunks. +frames :: MonadIO m => ByteString -> C.Conduit ByteString m ByteString +frames delim = frame delim >>= + return () `maybe` ((>> frames delim) . C.yield) + +leftover :: Monad m => ByteString -> C.Conduit ByteString m o +leftover bs | B.null bs = return () + | otherwise = C.leftover bs +
-- | Convert a 'UTCTime' to a string in ISO 8601 format. showUTCTime :: UTCTime -> String