Take n bytes of ByteString

getByteStringN.hs

import Data.ByteString (ByteString)
import Data.Iteratee
import qualified Data.ByteString as S
import qualified Data.ByteString.Builder as B
import qualified Data.ByteString.Lazy as L

getByteStringN :: Monad m => Int -> Iteratee ByteString m ByteString
getByteStringN = liftI . step mempty
    where
        step b n chunk
            | n <= 0 = idone (toBS b) chunk
        step b n (Chunk bytes)
            | len < n = liftI $ step (b <> fromBS bytes) (n - len)
            | otherwise = idone
                (toBS $ b <> fromBS (S.take n bytes))
                (Chunk (S.drop n bytes))
            where
                len = S.length bytes
        step _ _ (EOF Nothing) = throwErr $ toException EofException
        step _ _ (EOF (Just e)) = throwErr e
        fromBS = B.byteString
        toBS = L.toStrict . B.toLazyByteString