[tordnsel/master] fix frame reading complexity - tor-commits

16 Apr 2016

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

    