From 19e631c9f4621bfbe517e686f9bedc00ec1191fd Mon Sep 17 00:00:00 2001
From: Victor Denisov <denisovenator@gmail.com>
Date: Wed, 27 Apr 2016 00:10:12 -0700
Subject: [PATCH] Implement readExactly using hGet

---
 Database/MongoDB/Internal/Connection.hs | 37 ++-----------------------
 Database/MongoDB/Internal/Protocol.hs   |  4 +--
 2 files changed, 4 insertions(+), 37 deletions(-)

diff --git a/Database/MongoDB/Internal/Connection.hs b/Database/MongoDB/Internal/Connection.hs
index d921941..c0abdd7 100644
--- a/Database/MongoDB/Internal/Connection.hs
+++ b/Database/MongoDB/Internal/Connection.hs
@@ -23,49 +23,16 @@ import System.IO.Error (mkIOError, eofErrorType)
 --
 -- `read` should return `ByteString.null` on EOF
 data Connection = Connection {
-    read :: IO ByteString,
-    unread :: ByteString -> IO (),
+    readExactly :: Int -> IO ByteString,
     write :: ByteString -> IO (),
     flush :: IO (),
     close :: IO ()}
 
-readExactly :: Connection -> Int -> IO Lazy.ByteString
--- ^ Read specified number of bytes
---
--- If EOF is reached before N bytes then raise EOF exception.
-readExactly conn count = go mempty count
-  where
-  go acc n = do
-    -- read until get enough bytes
-    chunk <- read conn
-    when (ByteString.null chunk) $
-      ioError eof
-    let len = ByteString.length chunk
-    if len >= n
-      then do
-        let (res, rest) = ByteString.splitAt n chunk
-        unless (ByteString.null rest) $
-          unread conn rest
-        return (acc <> Lazy.ByteString.fromStrict res)
-      else go (acc <> Lazy.ByteString.fromStrict chunk) (n - len)
-  eof = mkIOError eofErrorType "Database.MongoDB.Internal.Connection"
-                  Nothing Nothing
-
 fromHandle :: Handle -> IO Connection
 -- ^ Make connection form handle
 fromHandle handle = do
-  restRef <- newIORef mempty
   return Connection
-    { read = do
-        rest <- readIORef restRef
-        writeIORef restRef mempty
-        if ByteString.null rest
-          -- 32k corresponds to the default chunk size
-          -- used in bytestring package
-          then ByteString.hGetSome handle (32 * 1024)
-          else return rest
-    , unread = \rest ->
-        modifyIORef restRef (rest <>)
+    { readExactly = ByteString.hGet handle
     , write = ByteString.hPut handle
     , flush = hFlush handle
     , close = hClose handle
diff --git a/Database/MongoDB/Internal/Protocol.hs b/Database/MongoDB/Internal/Protocol.hs
index 53db3be..1f3b534 100644
--- a/Database/MongoDB/Internal/Protocol.hs
+++ b/Database/MongoDB/Internal/Protocol.hs
@@ -117,8 +117,8 @@ readMessage :: Connection -> IO Response
 -- ^ read response from a connection
 readMessage conn = readResp  where
     readResp = do
-        len <- fromEnum . decodeSize <$> Connection.readExactly conn 4
-        runGet getReply <$> Connection.readExactly conn len
+        len <- fromEnum . decodeSize . L.fromStrict <$> Connection.readExactly conn 4
+        runGet getReply . L.fromStrict <$> Connection.readExactly conn len
     decodeSize = subtract 4 . runGet getInt32
 
 type FullCollection = Text