From 2e6eb504dcd4dd278ed5faa24b279d25113470d2 Mon Sep 17 00:00:00 2001
From: Evgeniy Gatov <egatov@yandex-team.ru>
Date: Mon, 14 Aug 2017 06:38:32 +0300
Subject: [PATCH] LimitReadBuffer fix [#CLICKHOUSE-2].

---
 dbms/src/IO/LimitReadBuffer.h | 60 +++++++++++++++++++++++++----------
 1 file changed, 43 insertions(+), 17 deletions(-)

diff --git a/dbms/src/IO/LimitReadBuffer.h b/dbms/src/IO/LimitReadBuffer.h
index 304c3302245..b99c8e95bf5 100644
--- a/dbms/src/IO/LimitReadBuffer.h
+++ b/dbms/src/IO/LimitReadBuffer.h
@@ -8,26 +8,52 @@ namespace DB
 
 /** Lets read from another ReadBuffer no more than the specified number of bytes.
   */
-class LimitReadBuffer : public ReadBuffer
-{
-private:
-    ReadBuffer & in;
-    size_t limit;
-
-    bool nextImpl() override
+    class LimitReadBuffer : public ReadBuffer
     {
-        if (count() >= limit || !in.next())
-            return false;
+    private:
+        ReadBuffer & in;
+        size_t limit;
 
-        working_buffer = in.buffer();
-        if (limit - count() < working_buffer.size())
-            working_buffer.resize(limit - count());
+        bool nextImpl() override
+        {
+            /// Let underlying buffer calculate read bytes in `next()` call.
+            in.position() = position();
 
-        return true;
-    }
+            if (bytes >= limit || !in.next())
+            {
+                return false;
+            }
 
-public:
-    LimitReadBuffer(ReadBuffer & in_, size_t limit_) : ReadBuffer(nullptr, 0), in(in_), limit(limit_) {}
-};
+            position() = in.position();
+
+            working_buffer = in.buffer();
+
+            if (limit - count() < working_buffer.size())
+            {
+                working_buffer.resize(limit - count());
+            }
+
+            return true;
+        }
+
+    public:
+        LimitReadBuffer(ReadBuffer & in_, size_t limit_) : ReadBuffer(in_.position(), 0), in(in_), limit(limit_)
+        {
+            working_buffer = in.buffer();
+
+            size_t bytes_in_buffer = working_buffer.end() - position();
+
+            working_buffer = Buffer(position(), working_buffer.end());
+
+            if (limit < bytes_in_buffer)
+                working_buffer.resize(limit);
+        }
+        virtual ~LimitReadBuffer() override
+        {
+            /// Update underlying buffer's position in case when limit wasn't reached.
+            if (working_buffer.size() != 0)
+                in.position() = position();
+        }
+    };
 
 }