From 123e3f8e76e6c6f0f3292cbb7bbe3c1ea9f34ad0 Mon Sep 17 00:00:00 2001
From: Alexey Milovidov <milovidov@yandex-team.ru>
Date: Sun, 26 Jun 2011 21:30:59 +0000
Subject: [PATCH] dbms: IO: fixed error [#CONV-2546].

---
 dbms/include/DB/IO/CompressedReadBuffer.h   | 20 ++++++++++----------
 dbms/include/DB/IO/CompressedWriteBuffer.h  |  8 ++++----
 dbms/include/DB/IO/ReadBuffer.h             |  3 +++
 dbms/include/DB/IO/ReadBufferFromIStream.h  | 17 +++++++++--------
 dbms/include/DB/IO/WriteBufferFromOStream.h |  6 +++---
 5 files changed, 29 insertions(+), 25 deletions(-)

diff --git a/dbms/include/DB/IO/CompressedReadBuffer.h b/dbms/include/DB/IO/CompressedReadBuffer.h
index 0068456049a..306ff9102c0 100644
--- a/dbms/include/DB/IO/CompressedReadBuffer.h
+++ b/dbms/include/DB/IO/CompressedReadBuffer.h
@@ -26,13 +26,6 @@ private:
 	std::vector<char> compressed_buffer;
 	char scratch[QLZ_SCRATCH_DECOMPRESS];
 
-public:
-	CompressedReadBuffer(ReadBuffer & in_)
-		: in(in_),
-		compressed_buffer(QUICKLZ_HEADER_SIZE)
-	{
-	}
-
 	bool nextImpl()
 	{
 		if (in.eof())
@@ -40,13 +33,13 @@ public:
 
 		uint128 checksum;
 		in.readStrict(reinterpret_cast<char *>(&checksum), sizeof(checksum));
-		
+
 		in.readStrict(&compressed_buffer[0], QUICKLZ_HEADER_SIZE);
 
 		size_t size_compressed = qlz_size_compressed(&compressed_buffer[0]);
 		if (size_compressed > DBMS_COMPRESSED_READ_BUFFER_MAX_COMPRESSED_SIZE)
 			throw Exception("Too large size_compressed. Most likely corrupted data.", ErrorCodes::TOO_LARGE_SIZE_COMPRESSED);
-			
+
 		size_t size_decompressed = qlz_size_decompressed(&compressed_buffer[0]);
 
 		compressed_buffer.resize(size_compressed);
@@ -56,13 +49,20 @@ public:
 
 		if (checksum != CityHash128(&compressed_buffer[0], size_compressed))
 			throw Exception("Checksum doesnt match: corrupted data.", ErrorCodes::CHECKSUM_DOESNT_MATCH);
-				
+
 		qlz_decompress(&compressed_buffer[0], &internal_buffer[0], scratch);
 
 		working_buffer = Buffer(working_buffer.begin(), working_buffer.begin() + size_decompressed);
 
 		return true;
 	}
+
+public:
+	CompressedReadBuffer(ReadBuffer & in_)
+		: in(in_),
+		compressed_buffer(QUICKLZ_HEADER_SIZE)
+	{
+	}
 };
 
 }
diff --git a/dbms/include/DB/IO/CompressedWriteBuffer.h b/dbms/include/DB/IO/CompressedWriteBuffer.h
index 66ac9b9249b..d02ecd376f2 100644
--- a/dbms/include/DB/IO/CompressedWriteBuffer.h
+++ b/dbms/include/DB/IO/CompressedWriteBuffer.h
@@ -23,14 +23,11 @@ private:
 
 	size_t compressed_bytes;
 
-public:
-	CompressedWriteBuffer(WriteBuffer & out_) : out(out_), compressed_bytes(0) {}
-
 	void nextImpl()
 	{
 		size_t uncompressed_size = pos - working_buffer.begin();
 		compressed_buffer.resize(uncompressed_size + QUICKLZ_ADDITIONAL_SPACE);
-		
+
 		size_t compressed_size = qlz_compress(
 			working_buffer.begin(),
 			&compressed_buffer[0],
@@ -44,6 +41,9 @@ public:
 		compressed_bytes += compressed_size;
 	}
 
+public:
+	CompressedWriteBuffer(WriteBuffer & out_) : out(out_), compressed_bytes(0) {}
+
 	/// Объём сжатых данных
 	size_t getCompressedBytes()
 	{
diff --git a/dbms/include/DB/IO/ReadBuffer.h b/dbms/include/DB/IO/ReadBuffer.h
index 1526fb468cd..501cb394155 100644
--- a/dbms/include/DB/IO/ReadBuffer.h
+++ b/dbms/include/DB/IO/ReadBuffer.h
@@ -57,6 +57,9 @@ public:
 	{
 		bytes_read += pos - working_buffer.begin();
 		bool res = nextImpl();
+		if (!res)
+			working_buffer = Buffer(working_buffer.begin(), working_buffer.begin());
+		
 		pos = working_buffer.begin();
 		return res;
 	}
diff --git a/dbms/include/DB/IO/ReadBufferFromIStream.h b/dbms/include/DB/IO/ReadBufferFromIStream.h
index 50d579d191e..f41176a6570 100644
--- a/dbms/include/DB/IO/ReadBufferFromIStream.h
+++ b/dbms/include/DB/IO/ReadBufferFromIStream.h
@@ -17,25 +17,26 @@ class ReadBufferFromIStream : public ReadBuffer
 private:
 	std::istream & istr;
 
-public:
-	ReadBufferFromIStream(std::istream & istr_) : istr(istr_) {}
-
 	bool nextImpl()
 	{
 		istr.read(working_buffer.begin(), DEFAULT_READ_BUFFER_SIZE);
+		size_t gcount = istr.gcount();
 
-		working_buffer = Buffer(working_buffer.begin(), working_buffer.begin() + istr.gcount());
-
-		if (working_buffer.end() == working_buffer.begin())
+		if (!gcount)
 		{
 			if (istr.eof())
 				return false;
-			if (!istr.good())
+			else
 				throw Exception("Cannot read from istream", ErrorCodes::CANNOT_READ_FROM_ISTREAM);
 		}
-		
+		else
+			working_buffer = Buffer(working_buffer.begin(), working_buffer.begin() + gcount);
+
 		return true;
 	}
+
+public:
+	ReadBufferFromIStream(std::istream & istr_) : istr(istr_) {}
 };
 
 }
diff --git a/dbms/include/DB/IO/WriteBufferFromOStream.h b/dbms/include/DB/IO/WriteBufferFromOStream.h
index 8aa0e2d096c..e121602494f 100644
--- a/dbms/include/DB/IO/WriteBufferFromOStream.h
+++ b/dbms/include/DB/IO/WriteBufferFromOStream.h
@@ -17,9 +17,6 @@ class WriteBufferFromOStream : public WriteBuffer
 private:
 	std::ostream & ostr;
 
-public:
-	WriteBufferFromOStream(std::ostream & ostr_) : ostr(ostr_) {}
-
 	void nextImpl()
 	{
 		ostr.write(working_buffer.begin(), pos - working_buffer.begin());
@@ -29,6 +26,9 @@ public:
 			throw Exception("Cannot write to ostream", ErrorCodes::CANNOT_WRITE_TO_OSTREAM);
 	}
 
+public:
+	WriteBufferFromOStream(std::ostream & ostr_) : ostr(ostr_) {}
+
 	~WriteBufferFromOStream()
 	{
 		next();