Merge branch 'zlibstream'

5 files changed, 208 insertions, 34 deletions

diff --git a/CMakeLists.txt b/CMakeLists.txt
index c16860a8ed..39a3a5759c 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -15,12 +15,13 @@ add_library(nbt++ STATIC
     src/value_initializer.cpp
 
     src/io/izlibstream.cpp
+    src/io/ozlibstream.cpp
     src/io/stream_reader.cpp
     src/io/stream_writer.cpp
 
     src/text/json_formatter.cpp)
 
-find_package(zlib REQUIRED)
+find_package(ZLIB REQUIRED)
 find_package(CxxTest)
 
 option(NBT_BUILD_TESTS "Build the unit tests. Requires CxxTest." ${CXXTEST_FOUND})
diff --git a/include/io/izlibstream.h b/include/io/izlibstream.h
index 6d58d56b86..43d2d2c02e 100644
--- a/include/io/izlibstream.h
+++ b/include/io/izlibstream.h
@@ -38,7 +38,8 @@ public:
      * @param input the istream to wrap
      * @param bufsize the size of the internal buffers
      * @param window_bits the base two logarithm of the maximum window size that
-     * zlib will use. This parameter also determines which type of input to expect.
+     * zlib will use.
+     * This parameter also determines which type of input to expect.
      * The default argument will autodetect between zlib and gzip data.
      * Refer to the zlib documentation of inflateInit2 for more details.
      *
diff --git a/include/io/ozlibstream.h b/include/io/ozlibstream.h
index 0e19aaaa2b..d4ea8c2840 100644
--- a/include/io/ozlibstream.h
+++ b/include/io/ozlibstream.h
@@ -29,12 +29,21 @@ namespace zlib
 
 /**
  * @brief Stream buffer used by zlib::ozlibstream
- * @see ozlibstream
+ * @sa ozlibstream
  */
 class deflate_streambuf : public std::streambuf
 {
 public:
-    explicit deflate_streambuf(std::ostream& output, int level = -1, int window_bits = 15, int mem_level = 8, int strategy = Z_DEFAULT_STRATEGY);
+    /**
+     * @param output the ostream to wrap
+     * @param bufsize the size of the internal buffers
+     * @param level the compression level, ranges from 0 to 9
+     *
+     * Refer to the zlib documentation of deflateInit2 for details about the arguments.
+     *
+     * @throw zlib_error if zlib encounters a problem during initialization
+     */
+    explicit deflate_streambuf(std::ostream& output, size_t bufsize = 32768, int level = Z_DEFAULT_COMPRESSION, int window_bits = 15, int mem_level = 8, int strategy = Z_DEFAULT_STRATEGY);
     ~deflate_streambuf() noexcept;
 
     std::ostream& get_ostr() const { return os; }
@@ -56,28 +65,22 @@ private:
  *
  * This ostream wraps another ostream. Data written to an ozlibstream will be
  * deflated (compressed) with zlib and written to the wrapped ostream.
+ *
+ * @sa deflate_streambuf
  */
 class ozlibstream : public std::ostream
 {
 public:
     /**
      * @param output the ostream to wrap
-     * @param level the compression level. Ranges from 0 to 9, or -1 for the default
-     * @param gzip_header whether to write a gzip header rather than a zlib header
+     * @param level the compression level, ranges from 0 to 9
+     * @param gzip if true, the output will be in gzip format rather than zlib
+     * @param bufsize the size of the internal buffers
      */
-    explicit ozlibstream(std::ostream& output, int level, bool gzip):
-        ozlibstream(output, level, 15 + (gzip ? 16 : 0))
+    explicit ozlibstream(std::ostream& output, int level = Z_DEFAULT_COMPRESSION, bool gzip = false, size_t bufsize = 32768):
+        std::ostream(&buf), buf(output, bufsize, level, 15 + (gzip ? 16 : 0))
     {}
 
-    /**
-     * @param output the ostream to wrap
-     * @param level the compression level. Ranges from 0 to 9, or -1 for the default
-     *
-     * Refer to the zlib documentation of deflateInit2 for a detailed explanation of the arguments.
-     */
-    explicit ozlibstream(std::ostream& output, int level = -1, int window_bits = 15, int mem_level = 8, int strategy = Z_DEFAULT_STRATEGY):
-        std::ostream(&buf), buf(output, level, window_bits, mem_level, strategy)
-    {}
     ///@return the wrapped ostream
     std::ostream& get_ostr() const { return buf.get_ostr(); }
 
diff --git a/src/io/ozlibstream.cpp b/src/io/ozlibstream.cpp
new file mode 100644
index 0000000000..97702a37e1
--- /dev/null
+++ b/src/io/ozlibstream.cpp
@@ -0,0 +1,87 @@
+/*
+ * libnbt++ - A library for the Minecraft Named Binary Tag format.
+ * Copyright (C) 2013, 2015  ljfa-ag
+ *
+ * This file is part of libnbt++.
+ *
+ * libnbt++ is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Lesser General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * libnbt++ is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with libnbt++.  If not, see <http://www.gnu.org/licenses/>.
+ */
+#include "io/ozlibstream.h"
+#include "io/zlib_error.h"
+
+namespace zlib
+{
+
+deflate_streambuf::deflate_streambuf(std::ostream& output, size_t bufsize, int level, int window_bits, int mem_level, int strategy):
+    os(output), in(bufsize), out(bufsize)
+{
+    zstr.zalloc = Z_NULL;
+    zstr.zfree = Z_NULL;
+    zstr.opaque = Z_NULL;
+    int ret = deflateInit2(&zstr, level, Z_DEFLATED, window_bits, mem_level, strategy);
+    if(ret != Z_OK)
+        throw zlib_error(zstr.msg, ret);
+
+    setp(in.data(), in.data() + in.size());
+}
+
+deflate_streambuf::~deflate_streambuf() noexcept
+{
+    try
+    {
+        deflate_chunk(Z_FINISH);
+    }
+    catch(...)
+    {
+        //ignore as we can't do anything about it
+    }
+    deflateEnd(&zstr);
+}
+
+void deflate_streambuf::deflate_chunk(int flush)
+{
+    zstr.next_in = reinterpret_cast<Bytef*>(pbase());
+    zstr.avail_in = pptr() - pbase();
+    do
+    {
+        zstr.next_out = reinterpret_cast<Bytef*>(out.data());
+        zstr.avail_out = out.size();
+        int ret = deflate(&zstr, flush);
+        if(ret != Z_OK && ret != Z_STREAM_END)
+            throw zlib_error(zstr.msg, ret);
+        int have = out.size() - zstr.avail_out;
+        if(!os.write(out.data(), have))
+            throw std::ios_base::failure("Could not write to the output stream");
+    } while(zstr.avail_out == 0);
+    setp(in.data(), in.data() + in.size());
+}
+
+deflate_streambuf::int_type deflate_streambuf::overflow(int_type ch)
+{
+    deflate_chunk();
+    if(ch != traits_type::eof())
+    {
+        *pptr() = ch;
+        pbump(1);
+    }
+    return ch;
+}
+
+int deflate_streambuf::sync()
+{
+    deflate_chunk();
+    return 0;
+}
+
+}
diff --git a/test/zlibstream_test.h b/test/zlibstream_test.h
index 8e9ca6b0f2..754f4e5d58 100644
--- a/test/zlibstream_test.h
+++ b/test/zlibstream_test.h
@@ -29,14 +29,16 @@ using namespace zlib;
 class zlibstream_test : public CxxTest::TestSuite
 {
 private:
-    std::stringbuf bigtest;
+    std::string bigtest;
 
 public:
     zlibstream_test()
     {
         std::ifstream bigtest_f("bigtest_uncompr", std::ios::binary);
-        bigtest_f >> &bigtest;
-        if(!bigtest_f || bigtest.str().size() == 0)
+        std::stringbuf bigtest_b;
+        bigtest_f >> &bigtest_b;
+        bigtest = bigtest_b.str();
+        if(!bigtest_f || bigtest.size() == 0)
             throw std::runtime_error("Could not read bigtest_uncompr file");
     }
 
@@ -49,13 +51,13 @@ public:
         //Small buffer so not all fits at once (the compressed file is 561 bytes)
         {
             izlibstream igzs(gzip_in, 256);
-            igzs.exceptions(std::ios::failbit);
+            igzs.exceptions(std::ios::failbit | std::ios::badbit);
             TS_ASSERT(igzs.good());
 
             TS_ASSERT_THROWS_NOTHING(igzs >> &data);
             TS_ASSERT(igzs);
             TS_ASSERT(igzs.eof());
-            TS_ASSERT_EQUALS(data.str(), bigtest.str());
+            TS_ASSERT_EQUALS(data.str(), bigtest);
         }
 
         //Clear and reuse buffers
@@ -65,13 +67,13 @@ public:
         //Now try the same with larger buffer (but not large enough for all output, uncompressed size 1561 bytes)
         {
             izlibstream igzs(gzip_in, 1000);
-            igzs.exceptions(std::ios::failbit);
+            igzs.exceptions(std::ios::failbit | std::ios::badbit);
             TS_ASSERT(igzs.good());
 
             TS_ASSERT_THROWS_NOTHING(igzs >> &data);
             TS_ASSERT(igzs);
             TS_ASSERT(igzs.eof());
-            TS_ASSERT_EQUALS(data.str(), bigtest.str());
+            TS_ASSERT_EQUALS(data.str(), bigtest);
         }
 
         data.str("");
@@ -80,13 +82,13 @@ public:
         //Now with large buffer
         {
             izlibstream igzs(gzip_in, 4000);
-            igzs.exceptions(std::ios::failbit);
+            igzs.exceptions(std::ios::failbit | std::ios::badbit);
             TS_ASSERT(igzs.good());
 
             TS_ASSERT_THROWS_NOTHING(igzs >> &data);
             TS_ASSERT(igzs);
             TS_ASSERT(igzs.eof());
-            TS_ASSERT_EQUALS(data.str(), bigtest.str());
+            TS_ASSERT_EQUALS(data.str(), bigtest);
         }
     }
 
@@ -97,13 +99,13 @@ public:
 
         std::stringbuf data;
         izlibstream izls(zlib_in, 256);
-        izls.exceptions(std::ios::failbit);
+        izls.exceptions(std::ios::failbit | std::ios::badbit);
         TS_ASSERT(izls.good());
 
         TS_ASSERT_THROWS_NOTHING(izls >> &data);
         TS_ASSERT(izls);
         TS_ASSERT(izls.eof());
-        TS_ASSERT_EQUALS(data.str(), bigtest.str());
+        TS_ASSERT_EQUALS(data.str(), bigtest);
     }
 
     void test_inflate_corrupt()
@@ -111,22 +113,102 @@ public:
         std::ifstream gzip_in("bigtest_corrupt.nbt", std::ios::binary);
         TS_ASSERT(gzip_in);
 
-        std::stringbuf data;
+        std::vector<char> buf(bigtest.size());
         {
             izlibstream igzs(gzip_in);
-            igzs.exceptions(std::ios::failbit);
-            TS_ASSERT_THROWS(igzs >> &data, zlib_error);
+            igzs.exceptions(std::ios::failbit | std::ios::badbit);
+            TS_ASSERT_THROWS(igzs.read(buf.data(), buf.size()), zlib_error);
+            TS_ASSERT(igzs.bad());
         }
 
         gzip_in.close();
+        gzip_in.clear();
         gzip_in.open("bigtest_eof.nbt", std::ios::binary);
         TS_ASSERT(gzip_in);
 
-        data.str("");
         {
             izlibstream igzs(gzip_in);
-            igzs.exceptions(std::ios::failbit);
-            TS_ASSERT_THROWS(igzs >> &data, zlib_error);
+            igzs.exceptions(std::ios::failbit | std::ios::badbit);
+            TS_ASSERT_THROWS(igzs.read(buf.data(), buf.size()), zlib_error);
+            TS_ASSERT(igzs.bad());
+        }
+    }
+
+    void test_deflate_zlib()
+    {
+        //Here we assume that inflating works and has already been tested
+        std::stringstream str;
+        std::stringbuf output;
+        //Small buffer
+        {
+            ozlibstream ozls(str, -1, false, 256);
+            ozls.exceptions(std::ios::failbit | std::ios::badbit);
+            ozls << bigtest;
+            TS_ASSERT(ozls.good());
+        }
+        TS_ASSERT(str.good());
+        {
+            izlibstream izls(str);
+            TS_ASSERT_THROWS_NOTHING(izls >> &output);
+            TS_ASSERT(izls);
+        }
+        TS_ASSERT_EQUALS(output.str(), bigtest);
+
+        str.clear(); str.str("");
+        output.str("");
+        //Medium sized buffer
+        //Write first half, then flush and write second half
+        {
+            ozlibstream ozls(str, 9, false, 512);
+            ozls.exceptions(std::ios::failbit | std::ios::badbit);
+
+            std::string half1 = bigtest.substr(0, bigtest.size()/2);
+            std::string half2 = bigtest.substr(bigtest.size()/2);
+            TS_ASSERT_THROWS_NOTHING(ozls << half1 << std::flush << half2);
+            TS_ASSERT(ozls.good());
+        }
+        TS_ASSERT(str.good());
+        {
+            izlibstream izls(str);
+            izls >> &output;
+            TS_ASSERT(izls);
+        }
+        TS_ASSERT_EQUALS(output.str(), bigtest);
+
+        str.clear(); str.str("");
+        output.str("");
+        //Large buffer
+        {
+            ozlibstream ozls(str, 1, false, 4000);
+            ozls.exceptions(std::ios::failbit | std::ios::badbit);
+            TS_ASSERT_THROWS_NOTHING(ozls << bigtest);
+            TS_ASSERT(ozls.good());
+        }
+        TS_ASSERT(str.good());
+        {
+            izlibstream izls(str);
+            izls >> &output;
+            TS_ASSERT(izls);
+        }
+        TS_ASSERT_EQUALS(output.str(), bigtest);
+    }
+
+    void test_deflate_gzip()
+    {
+        std::stringstream str;
+        std::stringbuf output;
+        {
+            ozlibstream ozls(str, -1, true);
+            ozls.exceptions(std::ios::failbit | std::ios::badbit);
+            TS_ASSERT_THROWS_NOTHING(ozls << bigtest);
+            TS_ASSERT(ozls.good());
+        }
+        TS_ASSERT(str.good());
+        {
+            izlibstream izls(str);
+            izls >> &output;
+            TS_ASSERT(izls);
         }
+        TS_ASSERT_EQUALS(output.str(), bigtest);
     }
 };