From 0de0de35367d130dd72fe11b489cf50d22b68933 Mon Sep 17 00:00:00 2001
From: Karol Kosek <krkk@serenityos.org>
Date: Fri, 17 Jun 2022 09:26:35 +0200
Subject: [PATCH] LibCompress: Implement a Zlib compressor

---
 Userland/Libraries/LibCompress/Zlib.cpp | 82 +++++++++++++++++++++++++
 Userland/Libraries/LibCompress/Zlib.h   | 23 +++++++
 2 files changed, 105 insertions(+)

diff --git a/Userland/Libraries/LibCompress/Zlib.cpp b/Userland/Libraries/LibCompress/Zlib.cpp
index f837e39995..eb0f43e173 100644
--- a/Userland/Libraries/LibCompress/Zlib.cpp
+++ b/Userland/Libraries/LibCompress/Zlib.cpp
@@ -4,7 +4,9 @@
  * SPDX-License-Identifier: BSD-2-Clause
  */
 
+#include <AK/MemoryStream.h>
 #include <AK/Span.h>
+#include <AK/TypeCasts.h>
 #include <AK/Types.h>
 #include <LibCompress/Deflate.h>
 #include <LibCompress/Zlib.h>
@@ -63,4 +65,84 @@ u32 Zlib::checksum()
     return m_checksum;
 }
 
+ZlibCompressor::ZlibCompressor(OutputStream& stream, ZlibCompressionLevel compression_level)
+    : m_output_stream(stream)
+{
+    // Zlib only defines Deflate as a compression method.
+    auto compression_method = ZlibCompressionMethod::Deflate;
+
+    write_header(compression_method, compression_level);
+
+    // FIXME: Find a way to compress with Deflate's "Best" compression level.
+    m_compressor = make<DeflateCompressor>(stream, static_cast<DeflateCompressor::CompressionLevel>(compression_level));
+}
+
+ZlibCompressor::~ZlibCompressor()
+{
+    VERIFY(m_finished);
+}
+
+void ZlibCompressor::write_header(ZlibCompressionMethod compression_method, ZlibCompressionLevel compression_level)
+{
+    ZlibHeader header {
+        .compression_method = compression_method,
+        .compression_info = 0,
+        .check_bits = 0,
+        .present_dictionary = false,
+        .compression_level = compression_level,
+    };
+    header.check_bits = 0b11111 - header.as_u16 % 31;
+
+    // FIXME: Support pre-defined dictionaries.
+
+    m_output_stream << header.as_u16;
+}
+
+size_t ZlibCompressor::write(ReadonlyBytes bytes)
+{
+    VERIFY(!m_finished);
+
+    size_t n_written = m_compressor->write(bytes);
+    m_adler32_checksum.update(bytes.trim(n_written));
+    return n_written;
+}
+
+bool ZlibCompressor::write_or_error(ReadonlyBytes bytes)
+{
+    if (write(bytes) < bytes.size()) {
+        set_fatal_error();
+        return false;
+    }
+
+    return true;
+}
+
+void ZlibCompressor::finish()
+{
+    VERIFY(!m_finished);
+
+    if (is<DeflateCompressor>(m_compressor.ptr()))
+        static_cast<DeflateCompressor*>(m_compressor.ptr())->final_flush();
+
+    NetworkOrdered<u32> adler_sum = m_adler32_checksum.digest();
+    m_output_stream << adler_sum;
+
+    m_finished = true;
+}
+
+Optional<ByteBuffer> ZlibCompressor::compress_all(ReadonlyBytes bytes, ZlibCompressionLevel compression_level)
+{
+    DuplexMemoryStream output_stream;
+    ZlibCompressor zlib_stream { output_stream, compression_level };
+
+    zlib_stream.write_or_error(bytes);
+
+    zlib_stream.finish();
+
+    if (zlib_stream.handle_any_error())
+        return {};
+
+    return output_stream.copy_into_contiguous_buffer();
+}
+
 }
diff --git a/Userland/Libraries/LibCompress/Zlib.h b/Userland/Libraries/LibCompress/Zlib.h
index e9d1dd5631..c526d5c123 100644
--- a/Userland/Libraries/LibCompress/Zlib.h
+++ b/Userland/Libraries/LibCompress/Zlib.h
@@ -6,10 +6,13 @@
 
 #pragma once
 
+#include <AK/BitStream.h>
 #include <AK/ByteBuffer.h>
 #include <AK/Optional.h>
+#include <AK/OwnPtr.h>
 #include <AK/Span.h>
 #include <AK/Types.h>
+#include <LibCrypto/Checksum/Adler32.h>
 
 namespace Compress {
 
@@ -57,4 +60,24 @@ private:
     ReadonlyBytes m_data_bytes;
 };
 
+class ZlibCompressor : public OutputStream {
+public:
+    ZlibCompressor(OutputStream&, ZlibCompressionLevel = ZlibCompressionLevel::Default);
+    ~ZlibCompressor();
+
+    size_t write(ReadonlyBytes) override;
+    bool write_or_error(ReadonlyBytes) override;
+    void finish();
+
+    static Optional<ByteBuffer> compress_all(ReadonlyBytes bytes, ZlibCompressionLevel = ZlibCompressionLevel::Default);
+
+private:
+    void write_header(ZlibCompressionMethod, ZlibCompressionLevel);
+
+    bool m_finished { false };
+    OutputBitStream m_output_stream;
+    OwnPtr<OutputStream> m_compressor;
+    Crypto::Checksum::Adler32 m_adler32_checksum;
+};
+
 }