Streamline: Added hability to limit the input

src/System/IO/Uniform/Streamline.hs

@@ -37,6 +37,7 @@ module System.IO.Uniform.Streamline (
   startTls,
   isSecure,
   transformTarget,
+  limitInput,
   echoTo,
   setEcho
   ) where
@@ -52,7 +53,7 @@ import Data.Default.Class
 import Control.Monad.Trans.Class
 import Control.Monad.Trans.Interruptible
 import Control.Monad.Trans.Control
-import Control.Monad (ap, liftM)
+import Control.Monad
 import Control.Monad.Base
 import Control.Monad.IO.Class
 import System.IO.Error
@@ -65,10 +66,10 @@ import Data.IP (IP)
 import qualified Data.Attoparsec.ByteString as A
 
 -- | Internal state for a Streamline monad
-data StreamlineState = StreamlineState {str :: SomeIO, buff :: ByteString, isEOF :: Bool, echo :: Maybe Handle}
+data StreamlineState = StreamlineState {str :: SomeIO, buff :: ByteString, isEOF :: Bool, echo :: Maybe Handle, inLimit :: Int, sentEmpty :: Bool}
 instance Default StreamlineState where
   -- | Will open StdIO
-  def = StreamlineState (SomeIO Std.StdIO) BS.empty False Nothing
+  def = StreamlineState (SomeIO Std.StdIO) BS.empty False Nothing (-1) False
 
 -- | Monad that emulates character stream IO over block IO.
 newtype Streamline m a = Streamline {withTarget' :: StreamlineState -> m (a, StreamlineState)}
@@ -76,22 +77,57 @@ newtype Streamline m a = Streamline {withTarget' :: StreamlineState -> m (a, Str
 blockSize :: Int
 blockSize = 4096
 
-readF :: MonadIO m => StreamlineState -> m ByteString
-readF cl = case echo cl of
-  Just h -> do
-    l <- liftIO $ S.uRead (str cl) blockSize
-    liftIO $ BS.hPutStr h "< "
-    liftIO $ BS.hPutStr h l
-    return l
-  Nothing -> liftIO $ S.uRead (str cl) blockSize
-
-writeF :: MonadIO m => StreamlineState -> ByteString -> m ()
-writeF cl l = case echo cl of
+readF :: MonadIO m => Streamline m ()
+readF = -- Must try just not to read more than the limit, actual limiting is done by takeBuff
+  Streamline $ \cl -> if not . BS.null . buff $ cl then return ((), cl)
+  else do
+  let lim = inLimit cl
+      sz = if lim < 0 then blockSize
+           else if lim <= blockSize then lim
+                else blockSize
+  l <- liftIO $ S.uRead (str cl) sz
+  let cl' = cl{buff= l}
+  case echo cl of
+    Just h -> do
+      liftIO $ BS.hPutStr h "< "
+      liftIO $ BS.hPutStr h l
+    Nothing -> return ()
+  return ((), cl')
+
+-- | Takes the buffer for processing
+takeBuff :: MonadIO m => Streamline m ByteString
+takeBuff = do
+  readF
+  Streamline $ \cl -> 
+    let lim = inLimit cl
+        eof = isEOF cl
+        b = buff cl
+    in if eof then eofError "System.IO.Uniform.Streamline"
+       else if lim < 0 then return (b, cl{buff="", isEOF=BS.null b})
+            else let (r, b') = BS.splitAt lim b
+                 in return (r, cl{
+                                                        isEOF = BS.null b || sentEmpty cl,
+                                                        sentEmpty = BS.null r,
+                                                        buff = b',
+                                                        inLimit = lim - BS.length r
+                                                        })
+
+-- | Pushes remaining data back into the buffer
+pushBuff :: Monad m => ByteString -> Streamline m ()
+pushBuff dt = Streamline $ \cl -> let
+  lim = inLimit cl
+  b = buff cl
+  in if lim <= 0 then return ((), cl{buff = BS.append dt b})
+     else return ((), cl{buff = BS.append dt b, inLimit = lim - BS.length dt})
+
+writeF :: MonadIO m => ByteString -> Streamline m ()
+writeF l = Streamline $ \cl -> case echo cl of
   Just h -> do
     liftIO $ BS.hPutStr h "> "
     liftIO $ BS.hPutStr h l
     liftIO $ S.uPut (str cl) l
-  Nothing -> liftIO $ S.uPut (str cl) l
+    return ((), cl)
+  Nothing -> liftIO $ S.uPut (str cl) l >> return ((), cl)
 
 -- | > withServer f serverIP port
 --
@@ -153,16 +189,13 @@ instance MonadIO m => MonadIO (Streamline m) where
 
 -- | Sends data over the IO target.
 send :: MonadIO m => ByteString -> Streamline m ()
-send r = Streamline $ \cl -> do
-  writeF cl r
-  return ((), cl)
+send r = writeF r
 
 -- | Sends data from a lazy byte string
 send' :: MonadIO m => LBS.ByteString -> Streamline m ()
-send' r = Streamline $ \cl -> do
+send' r = do
   let dd = LBS.toChunks r
-  mapM_ (writeF cl) dd
-  return ((), cl)
+  mapM_ writeF dd
 
 {- |
 Very much like Attoparsec's runScanner:
@@ -181,29 +214,24 @@ runScanner state scanner = do
   (rt, st) <- runScanner' state scanner
   return (LBS.toStrict rt, st)
 
--- | Equivalent to runScanner, but returns a strict, completely
---   evaluated ByteString.
+-- | Equivalent to runScanner, but returns a lazy ByteString
 runScanner' :: MonadIO m => s -> IOScanner s -> Streamline m (LBS.ByteString, s)
-runScanner' state scanner = Streamline $ \d ->
-  do
-    (tx, st, d') <- in_scan d state
-    return ((LBS.fromChunks tx, st), d')
+runScanner' state scanner = do
+  (tx, st) <- in_scan state
+  return (LBS.fromChunks tx, st)
   where
-    --in_scan :: StreamlineState -> s -> m ([ByteString], s, StreamlineState)
-    in_scan d st
-      | isEOF d = eofError "System.IO.Uniform.Streamline.scan'"
-      | BS.null (buff d) = do
-        dt <- readF d
-        if BS.null dt
-          then return ([], st, d{isEOF=True})
-          else in_scan d{buff=dt} st
-      | otherwise = case sscan scanner st 0 (BS.unpack . buff $ d) of
-        AllInput st' -> do
-          (tx', st'', d') <- in_scan d{buff=""} st'
-          return (buff d:tx', st'', d')
-        SplitAt n st' -> let
-          (r, i) = BS.splitAt n (buff d)
-          in return ([r], st', d{buff=i})
+    --in_scan :: MonadIO m => s -> Streamline m ([ByteString], s)
+    in_scan st = do
+      d <- takeBuff
+      if BS.null d then return ([], st)
+        else case sscan scanner st 0 $ BS.unpack d of
+               AllInput st' -> do
+                 (tx', st'') <- in_scan st'
+                 return (d:tx', st'')
+               SplitAt n st' -> do
+                 let (r, i) = BS.splitAt n d
+                 pushBuff i
+                 return ([r], st')
     -- I'll avoid rebuilding a list on high level code. The ByteString functions are way better.
     sscan :: (s -> Word8 -> IOScannerState s) -> s -> Int -> [Word8] -> ScanResult s
     sscan _ s0 _ [] = AllInput s0
@@ -231,29 +259,27 @@ recieveLine = recieveTill "\n"
 recieveLine' :: MonadIO m => Streamline m LBS.ByteString
 recieveLine' = recieveTill' "\n"
 
--- | Recieves the given number of bytes.
+{- |
+Recieves the given number of bytes, or less in case of end of file.
+-}
 recieveN :: MonadIO m => Int -> Streamline m ByteString
 recieveN n = LBS.toStrict <$> recieveN' n
 
 -- | Lazy version of recieveN
 recieveN' :: MonadIO m => Int -> Streamline m LBS.ByteString
-recieveN' n | n <= 0 = return ""
-            | otherwise = Streamline $ \cl ->
-            do
-              (tt, cl') <- recieve cl n
-              return (LBS.fromChunks tt, cl')
+recieveN' n = LBS.fromChunks <$> recieve n
   where
-    recieve d b
-      | isEOF d = eofError "System.IO.Uniform.Streamline.recieveN"
-      | BS.null . buff $ d = do
-        dt <- readF d
-        recieve d{buff=dt}{isEOF=BS.null dt} b
-      | b <= (BS.length . buff $ d) = let
-        (r, dt) = BS.splitAt b $ buff d
-        in return ([r], d{buff=dt})
+    recieve sz
+      | sz <= 0 = return []
       | otherwise = do
-        (r, d') <- recieve d{buff=""} $ b - (BS.length . buff $ d)
-        return (buff d : r, d')
+          d <- takeBuff
+          if BS.null d then return []
+            else do
+            let (h, t) = BS.splitAt sz d
+                sz' = sz - BS.length h
+            unless (BS.null t) $ pushBuff t
+            r <- recieve sz'
+            return $ h : r
 
 -- | Recieves data until it matches the argument.
 --   Returns all of it, including the matching data.
@@ -277,43 +303,26 @@ startTls st = Streamline $ \cl -> do
 --  streamlined IO target. Returns both the parser
 --  result and the string consumed by it.
 runAttoparsecAndReturn :: MonadIO m => A.Parser a -> Streamline m (ByteString, Either String a)
-runAttoparsecAndReturn p = Streamline $ \cl ->
-  if isEOF cl
-  then eofError "System.IO.Uniform.Streamline.runAttoparsecAndReturn"
-  else do
-    let c = A.parse p $ buff cl
-    (cl', i, a) <- liftIO $ continueResult cl c
-    return ((i, a), cl')
+runAttoparsecAndReturn p = do
+  d <- takeBuff
+  let c = A.parse p d
+  continueResult c d
   where
-    continueResult :: StreamlineState -> A.Result a -> IO (StreamlineState, ByteString, Either String a)
-    -- tx eof ds 
-    continueResult cl c = case c of
-      A.Fail i _ msg -> return (cl{buff=i}, BS.take (BS.length (buff cl) - BS.length i) (buff cl), Left msg)
-      A.Done i r -> return (cl{buff=i}, BS.take (BS.length (buff cl) - BS.length i) (buff cl), Right r)
+    continueResult c d = case c of
+      A.Fail i _ msg -> do
+        pushBuff i
+        return (BS.take (BS.length d - BS.length i) d, Left msg)
+      A.Done i r -> do
+        pushBuff i
+        return (BS.take (BS.length d - BS.length i) d, Right r)
       A.Partial c' -> do
-        d <- readF cl
-        let cl' = cl{buff=BS.append (buff cl) d}{isEOF=BS.null d}
-        continueResult cl' (c' d)
+        dt <- takeBuff
+        continueResult (c' dt) dt
 
 -- | Runs an Attoparsec parser over the data read from the
 --  streamlined IO target. Returning the parser result.
 runAttoparsec :: MonadIO m => A.Parser a -> Streamline m (Either String a)
-runAttoparsec p = Streamline $ \cl -> 
-  if isEOF cl
-  then eofError "System.IO.Uniform.Streamline.runAttoparsec"
-  else do
-    let c = A.parse p $ buff cl
-    (cl', a) <- liftIO $ continueResult cl c
-    return (a, cl')
-  where
-    continueResult :: StreamlineState -> A.Result a -> IO (StreamlineState, Either String a)
-    continueResult cl c = case c of
-        A.Fail i _ msg -> return (cl{buff=i}, Left msg)
-        A.Done i r -> return (cl{buff=i}, Right r)
-        A.Partial c' -> do
-          d <- readF cl
-          let eof' = BS.null d
-          continueResult cl{buff=d}{isEOF=eof'} (c' d)
+runAttoparsec p = snd <$> runAttoparsecAndReturn p
 
 -- | Indicates whether transport layer security is being used.
 isSecure :: Monad m => Streamline m Bool
@@ -335,6 +344,14 @@ Discards all buffered data in the process.
 transformTarget :: (UniformIO a, Monad m) => (SomeIO -> a) -> Streamline m ()
 transformTarget w = Streamline $ \cl -> return ((), cl{str = SomeIO . w . str $ cl})
 
+{- |
+Limits the input to the given number of bytes, emulating an end of file after them.
+
+If the limit is negative, the input will not be limited.
+-}
+limitInput :: Monad m => Int -> Streamline m ()
+limitInput n = Streamline $ \cl -> return ((), cl{inLimit = n})
+
 {- |
 Sets echo of the streamlined IO target.
 

test/LimitedInput.hs

@@ -0,0 +1,53 @@
+{-# LANGUAGE OverloadedStrings #-}
+
+module LimitedInput where
+
+import Distribution.TestSuite
+import Base (simpleTest)
+import System.IO.Error
+import System.IO.Uniform.ByteString
+import qualified System.IO.Uniform.Streamline as S
+import System.Timeout (timeout)
+import Data.ByteString (ByteString)
+import qualified Data.ByteString as BS
+import qualified Data.ByteString.Char8 as C8
+import qualified Data.Attoparsec.ByteString as A
+import Control.Monad
+
+tests :: IO [Test]
+tests = return [
+  simpleTest "recieveLine" $ recvTest False "A test\n" 3 S.recieveLine,
+  simpleTest "recieveN" $ recvTest False "A test\n" 3 $ S.recieveN 30,
+  simpleTest "runAttoparsec" $ recvTest False "abcedefghijklmnopqrstuvwxyz" 3 $ runParser parseAlphabet,
+  simpleTest "recieveLine" $ recvTest True "A test\n" 3 S.recieveLine,
+  simpleTest "recieveN" $ recvTest True "A test\n" 3 $ S.recieveN 30,
+  simpleTest "runAttoparsec" $ recvTest True "abcedefghijklmnopqrstuvwxyz" 3 $ runParser parseAlphabet
+  ]
+
+recvTest :: Bool -> ByteString -> Int -> S.Streamline IO ByteString -> IO Progress
+recvTest readFirst input limit f = do
+  r <- tryIOError $ test readFirst
+  case r of
+    Left e -> return . Finished . Fail $ "IO error: " ++ show e
+    Right v -> return v
+  where
+    test recv = do
+      to <- timeout 1000000 $ do
+        (r, _) <- withByteStringIO input $ \t -> do
+          S.withTarget t $ do
+            when recv $ S.recieveN 1 >> return ()
+            S.limitInput limit
+            f
+        let extra = if readFirst then 1 else 0
+        if BS.length r <= limit + extra
+          then return . Finished $ Pass
+          else return . Finished . Fail $ "Got too much text: " ++ show r
+      case to of
+        Nothing -> return . Finished . Fail $ "Execution blocked"
+        Just v -> return v
+
+parseAlphabet :: A.Parser ByteString
+parseAlphabet = A.string "abcedefghijklmnopqrstuvwxyz"
+
+runParser :: A.Parser ByteString -> S.Streamline IO ByteString
+runParser p = fst <$> S.runAttoparsecAndReturn p

uniform-io.cabal

@@ -120,3 +120,19 @@ Test-suite blocking
     Base
   ghc-options: -Wall -fno-warn-unused-do-bind -fwarn-incomplete-patterns -threaded
   default-language: Haskell2010
+
+Test-suite limited_input
+  type: detailed-0.9
+  test-module: LimitedInput
+  hs-source-dirs:
+    test
+  build-depends:
+    base >=4.7,
+    Cabal >= 1.9.2,
+    bytestring >=0.10 && <1.0,
+    attoparsec >=0.10 && <1.0,
+    uniform-io
+  other-modules:
+    Base
+  ghc-options: -Wall -fno-warn-unused-do-bind -fwarn-incomplete-patterns -threaded
+  default-language: Haskell2010