Rewrite compression with OOP

native/jni/magiskboot/compress.cpp

@@ -4,392 +4,25 @@
 #include <string.h>
 #include <sys/mman.h>
 
-#include <zlib.h>
-#include <lzma.h>
-#include <lz4.h>
-#include <lz4frame.h>
-#include <lz4hc.h>
-#include <bzlib.h>
 #include <logging.h>
 #include <utils.h>
 
 #include "magiskboot.h"
+#include "compress.h"
 
-#define CHUNK 0x40000
+int64_t decompress(format_t type, int fd, const void *from, size_t size) {
-
+	auto cmp = get_decoder(type);
-// Mode: 0 = decode; 1 = encode
+	int64_t ret = cmp->one_step(fd, from, size);
-size_t gzip(int mode, int fd, const void *buf, size_t size) {
+	delete cmp;
-	size_t ret = 0, have, total = 0;
+	return ret;
-	z_stream strm;
-	unsigned char out[CHUNK];
-
-	strm.zalloc = Z_NULL;
-	strm.zfree = Z_NULL;
-	strm.opaque = Z_NULL;
-
-	switch(mode) {
-		case 0:
-			ret = inflateInit2(&strm, 15 | 16);
-			break;
-		case 1:
-			ret = deflateInit2(&strm, 9, Z_DEFLATED, 15 | 16, 8, Z_DEFAULT_STRATEGY);
-			break;
 }
 
-	if (ret != Z_OK)
+int64_t compress(format_t type, int fd, const void *from, size_t size) {
-		LOGE("Unable to init zlib stream\n");
+	auto cmp = get_encoder(type);
-
+	int64_t ret = cmp->one_step(fd, from, size);
-	strm.next_in = (Bytef *) buf;
+	delete cmp;
-	strm.avail_in = size;
+	return ret;
-
-	do {
-		strm.avail_out = CHUNK;
-		strm.next_out = out;
-		switch(mode) {
-			case 0:
-				ret = inflate(&strm, Z_FINISH);
-				break;
-			case 1:
-				ret = deflate(&strm, Z_FINISH);
-				break;
 }
-		if (ret == Z_STREAM_ERROR)
-			LOGE("Error when running gzip\n");
-		have = CHUNK - strm.avail_out;
-		total += xwrite(fd, out, have);
-	} while (strm.avail_out == 0);
-
-	switch(mode) {
-		case 0:
-			inflateEnd(&strm);
-			break;
-		case 1:
-			deflateEnd(&strm);
-			break;
-	}
-	return total;
-}
-
-// Mode: 0 = decode xz/lzma; 1 = encode xz; 2 = encode lzma
-size_t lzma(int mode, int fd, const void *buf, size_t size) {
-	size_t have, total = 0;
-	lzma_ret ret = LZMA_OK;
-	lzma_stream strm = LZMA_STREAM_INIT;
-	lzma_options_lzma opt;
-	unsigned char out[CHUNK];
-
-	// Initialize preset
-	lzma_lzma_preset(&opt, 9);
-	lzma_filter filters[] = {
-		{ .id = LZMA_FILTER_LZMA2, .options = &opt },
-		{ .id = LZMA_VLI_UNKNOWN, .options = nullptr },
-	};
-
-	switch(mode) {
-		case 0:
-			ret = lzma_auto_decoder(&strm, UINT64_MAX, 0);
-			break;
-		case 1:
-			ret = lzma_stream_encoder(&strm, filters, LZMA_CHECK_CRC32);
-			break;
-		case 2:
-			ret = lzma_alone_encoder(&strm, &opt);
-			break;
-	}
-
-
-	if (ret != LZMA_OK)
-		LOGE("Unable to init lzma stream\n");
-
-	strm.next_in = static_cast<const uint8_t *>(buf);
-	strm.avail_in = size;
-
-	do {
-		strm.avail_out = CHUNK;
-		strm.next_out = out;
-		ret = lzma_code(&strm, LZMA_FINISH);
-		if (ret != LZMA_OK && ret != LZMA_STREAM_END)
-			LOGE("LZMA error %d!\n", ret);
-		have = CHUNK - strm.avail_out;
-		total += xwrite(fd, out, have);
-	} while (strm.avail_out == 0);
-
-	lzma_end(&strm);
-	return total;
-}
-
-// Mode: 0 = decode; 1 = encode
-size_t lz4(int mode, int fd, const uint8_t *buf, size_t size) {
-	LZ4F_decompressionContext_t dctx;
-	LZ4F_compressionContext_t cctx;
-	LZ4F_frameInfo_t info;
-
-	size_t blockSize, outCapacity, avail_in, ret = 0, pos = 0, total = 0;
-	size_t have, read;
-	uint8_t *out = nullptr;
-
-	// Initialize context
-	switch(mode) {
-		case 0:
-			ret = LZ4F_createDecompressionContext(&dctx, LZ4F_VERSION);
-			break;
-		case 1:
-			ret = LZ4F_createCompressionContext(&cctx, LZ4F_VERSION);
-			break;
-	}
-
-	if (LZ4F_isError(ret))
-		LOGE("Context creation error: %s\n", LZ4F_getErrorName(ret));
-
-	// Allocate out buffer
-	blockSize = 1 << 22;
-	switch(mode) {
-		case 0:
-			// Read header
-			read = blockSize;
-			ret = LZ4F_getFrameInfo(dctx, &info, buf, &read);
-			if (LZ4F_isError(ret))
-				LOGE("LZ4F_getFrameInfo error: %s\n", LZ4F_getErrorName(ret));
-			switch (info.blockSizeID) {
-				case LZ4F_default:
-				case LZ4F_max64KB:  outCapacity = 1 << 16; break;
-				case LZ4F_max256KB: outCapacity = 1 << 18; break;
-				case LZ4F_max1MB:   outCapacity = 1 << 20; break;
-				case LZ4F_max4MB:   outCapacity = 1 << 22; break;
-				default:
-					LOGE("Impossible unless more block sizes are allowed\n");
-			}
-			pos += read;
-			break;
-		case 1:
-			outCapacity = LZ4F_compressFrameBound(blockSize, nullptr);
-			break;
-	}
-
-	out = new uint8_t[outCapacity];
-
-	// Write header
-	if (mode == 1) {
-		LZ4F_preferences_t prefs = LZ4F_preferences_t();
-		prefs.autoFlush = 1;
-		prefs.compressionLevel = 9;
-		prefs.frameInfo.blockMode = LZ4F_blockIndependent;
-		prefs.frameInfo.blockSizeID = LZ4F_max4MB;
-		prefs.frameInfo.blockChecksumFlag = LZ4F_noBlockChecksum;
-		prefs.frameInfo.contentChecksumFlag = LZ4F_contentChecksumEnabled;
-		have = ret = LZ4F_compressBegin(cctx, out, size, &prefs);
-		if (LZ4F_isError(ret))
-			LOGE("Failed to start compression: error %s\n", LZ4F_getErrorName(ret));
-		total += xwrite(fd, out, have);
-	}
-
-	do {
-		if (pos + blockSize >= size) {
-			avail_in = size - pos;
-		} else {
-			avail_in = blockSize;
-		}
-
-		do {
-			switch(mode) {
-				case 0:
-					have = outCapacity;
-					read = avail_in;
-					ret = LZ4F_decompress(dctx, out, &have, buf + pos, &read, nullptr);
-					break;
-				case 1:
-					read = avail_in;
-					have = ret = LZ4F_compressUpdate(cctx, out, outCapacity, buf + pos, avail_in, nullptr);
-					break;
-			}
-			if (LZ4F_isError(ret))
-				LOGE("LZ4 coding error: %s\n", LZ4F_getErrorName(ret));
-
-			total += xwrite(fd, out, have);
-			// Update status
-			pos += read;
-			avail_in -= read;
-		} while(avail_in != 0 && ret != 0);
-
-	} while(pos < size && ret != 0);
-
-	switch(mode) {
-		case 0:
-			LZ4F_freeDecompressionContext(dctx);
-			break;
-		case 1:
-			have = ret = LZ4F_compressEnd(cctx, out, outCapacity, nullptr);
-			if (LZ4F_isError(ret))
-				LOGE("Failed to end compression: error %s\n", LZ4F_getErrorName(ret));
-
-			total += xwrite(fd, out, have);
-
-			LZ4F_freeCompressionContext(cctx);
-			break;
-	}
-
-	delete[] out;
-	return total;
-}
-
-// Mode: 0 = decode; 1 = encode
-size_t bzip2(int mode, int fd, const void* buf, size_t size) {
-	size_t ret = 0, have, total = 0;
-	bz_stream strm;
-	char out[CHUNK];
-
-	strm.bzalloc = nullptr;
-	strm.bzfree = nullptr;
-	strm.opaque = nullptr;
-
-	switch(mode) {
-		case 0:
-			ret = BZ2_bzDecompressInit(&strm, 0, 0);
-			break;
-		case 1:
-			ret = BZ2_bzCompressInit(&strm, 9, 0, 0);
-			break;
-	}
-
-	if (ret != BZ_OK)
-		LOGE("Unable to init bzlib stream\n");
-
-	strm.next_in = (char *) buf;
-	strm.avail_in = size;
-
-	do {
-		strm.avail_out = CHUNK;
-		strm.next_out = out;
-		switch(mode) {
-			case 0:
-				ret = BZ2_bzDecompress(&strm);
-				break;
-			case 1:
-				ret = BZ2_bzCompress(&strm, BZ_FINISH);
-				break;
-		}
-		have = CHUNK - strm.avail_out;
-		total += xwrite(fd, out, have);
-	} while (strm.avail_out == 0);
-
-	switch(mode) {
-		case 0:
-			BZ2_bzDecompressEnd(&strm);
-			break;
-		case 1:
-			BZ2_bzCompressEnd(&strm);
-			break;
-	}
-	return total;
-}
-
-#define LZ4_LEGACY_BLOCKSIZE  0x800000
-
-// Mode: 0 = decode; 1 = encode
-size_t lz4_legacy(int mode, int fd, const uint8_t *buf, size_t size) {
-	size_t pos = 0;
-	int have;
-	char *out;
-	unsigned block_size, insize, total = 0;
-
-	switch(mode) {
-		case 0:
-			out = new char[LZ4_LEGACY_BLOCKSIZE];
-			// Skip magic
-			pos += 4;
-			break;
-		case 1:
-			out = new char[LZ4_COMPRESSBOUND(LZ4_LEGACY_BLOCKSIZE)];
-			// Write magic
-			total += xwrite(fd, "\x02\x21\x4c\x18", 4);
-			break;
-	}
-
-	do {
-		switch(mode) {
-			case 0:
-				// Read block size
-				block_size = *(unsigned *)(buf + pos);
-				pos += 4;
-				if (block_size > LZ4_COMPRESSBOUND(LZ4_LEGACY_BLOCKSIZE))
-					goto done;
-				have = LZ4_decompress_safe((const char *) buf + pos, out, block_size, LZ4_LEGACY_BLOCKSIZE);
-				if (have < 0)
-					LOGE("Cannot decode lz4_legacy block\n");
-				pos += block_size;
-				break;
-			case 1:
-				if (pos + LZ4_LEGACY_BLOCKSIZE >= size)
-					insize = size - pos;
-				else
-					insize = LZ4_LEGACY_BLOCKSIZE;
-				have = LZ4_compress_HC((const char *) buf + pos, out, insize, LZ4_COMPRESSBOUND(LZ4_LEGACY_BLOCKSIZE), 9);
-				if (have == 0)
-					LOGE("lz4_legacy compression error\n");
-				pos += insize;
-				// Write block size
-				total += xwrite(fd, &have, sizeof(have));
-				break;
-		}
-		// Write main data
-		total += xwrite(fd, out, have);
-	} while(pos < size);
-
-done:
-	if (mode == 1) {
-		// Append original size to output
-		unsigned uncomp = size;
-		xwrite(fd, &uncomp, sizeof(uncomp));
-	}
-	delete[] out;
-	return total;
-}
-
-long long decompress(format_t type, int fd, const void *from, size_t size) {
-	const uint8_t *buf = (uint8_t *) from;
-	switch (type) {
-		case GZIP:
-			return gzip(0, fd, buf, size);
-		case XZ:
-			return lzma(0, fd, buf, size);
-		case LZMA:
-			return lzma(0, fd, buf, size);
-		case BZIP2:
-			return bzip2(0, fd, buf, size);
-		case LZ4:
-			return lz4(0, fd, buf, size);
-		case LZ4_LEGACY:
-			return lz4_legacy(0, fd, buf, size);
-		default:
-			// Unsupported
-			return -1;
-	}
-}
-
-long long compress(format_t type, int fd, const void *from, size_t size) {
-	const uint8_t *buf = (uint8_t *) from;
-	switch (type) {
-		case GZIP:
-			return gzip(1, fd, buf, size);
-		case XZ:
-			return lzma(1, fd, buf, size);
-		case LZMA:
-			return lzma(2, fd, buf, size);
-		case BZIP2:
-			return bzip2(1, fd, buf, size);
-		case LZ4:
-			return lz4(1, fd, buf, size);
-		case LZ4_LEGACY:
-			return lz4_legacy(1, fd, buf, size);
-		default:
-			// Unsupported
-			return -1;
-	}
-}
-
-/*
- * Below are utility functions for commandline
- */
 
 void decompress(char *from, const char *to) {
 	int strip = 1;
@@ -515,3 +148,436 @@ void compress(const char *method, const char *from, const char *to) {
 		unlink(from);
 }
 
+
+/* Compression Streams */
+
+Compression *get_encoder(format_t type) {
+	switch (type) {
+		case XZ:
+			return new XZEncoder();
+		case LZMA:
+			return new LZMAEncoder();
+		case BZIP2:
+			return new BZEncoder();
+		case LZ4:
+			return new LZ4FEncoder();
+		case LZ4_LEGACY:
+			return new LZ4Encoder();
+		case GZIP:
+		default:
+			return new GZEncoder();
+	}
+}
+
+Compression *get_decoder(format_t type) {
+	switch (type) {
+		case XZ:
+		case LZMA:
+			return new LZMADecoder();
+		case BZIP2:
+			return new BZDecoder();
+		case LZ4:
+			return new LZ4FDecoder();
+		case LZ4_LEGACY:
+			return new LZ4Decoder();
+		case GZIP:
+		default:
+			return new GZDecoder();
+	}
+}
+
+Compression::Compression() : fn([](auto, auto) -> void {}) {}
+
+void Compression::set_outfn(std::function<void(const void *, size_t)> &&fn) {
+	this->fn = std::move(fn);
+}
+
+void Compression::set_outfd(int fd) {
+	fn = [=](const void *out, size_t len) -> void {
+		xwrite(fd, out, len);
+	};
+}
+
+int64_t Compression::one_step(int outfd, const void *in, size_t size) {
+	set_outfd(outfd);
+	if (!update(in, size))
+		return -1;
+	return finalize();
+}
+
+GZStream::GZStream(int mode) : mode(mode), strm({}) {
+	switch(mode) {
+		case 0:
+			inflateInit2(&strm, 15 | 16);
+			break;
+		case 1:
+			deflateInit2(&strm, 9, Z_DEFLATED, 15 | 16, 8, Z_DEFAULT_STRATEGY);
+			break;
+	}
+}
+
+bool GZStream::update(const void *in, size_t size) {
+	return update(in, size, Z_NO_FLUSH);
+}
+
+uint64_t GZStream::finalize() {
+	update(nullptr, 0, Z_FINISH);
+	uint64_t total = strm.total_out;
+	switch(mode) {
+		case 0:
+			inflateEnd(&strm);
+			break;
+		case 1:
+			deflateEnd(&strm);
+			break;
+	}
+	return total;
+}
+
+bool GZStream::update(const void *in, size_t size, int flush) {
+	int ret;
+	strm.next_in = (Bytef *) in;
+	strm.avail_in = size;
+	do {
+		strm.next_out = outbuf;
+		strm.avail_out = sizeof(outbuf);
+		switch(mode) {
+			case 0:
+				ret = inflate(&strm, flush);
+				break;
+			case 1:
+				ret = deflate(&strm, flush);
+				break;
+		}
+		if (ret == Z_STREAM_ERROR) {
+			LOGW("Gzip %s failed (%d)\n", mode ? "encode" : "decode", ret);
+			return false;
+		}
+		fn(outbuf, sizeof(outbuf) - strm.avail_out);
+	} while (strm.avail_out == 0);
+	return true;
+}
+
+BZStream::BZStream(int mode) : mode(mode), strm({}) {
+	switch(mode) {
+		case 0:
+			BZ2_bzDecompressInit(&strm, 0, 0);
+			break;
+		case 1:
+			BZ2_bzCompressInit(&strm, 9, 0, 0);
+			break;
+	}
+}
+
+bool BZStream::update(const void *in, size_t size) {
+	return update(in, size, BZ_RUN);
+}
+
+uint64_t BZStream::finalize() {
+	update(nullptr, 0, BZ_FINISH);
+	uint64_t total = ((uint64_t) strm.total_out_hi32 << 32) + strm.total_out_lo32;
+	switch(mode) {
+		case 0:
+			BZ2_bzDecompressEnd(&strm);
+			break;
+		case 1:
+			BZ2_bzCompressEnd(&strm);
+			break;
+	}
+	return total;
+}
+
+bool BZStream::update(const void *in, size_t size, int flush) {
+	int ret;
+	strm.next_in = (char *) in;
+	strm.avail_in = size;
+	do {
+		strm.avail_out = sizeof(outbuf);
+		strm.next_out = outbuf;
+		switch(mode) {
+			case 0:
+				ret = BZ2_bzDecompress(&strm);
+				break;
+			case 1:
+				ret = BZ2_bzCompress(&strm, flush);
+				break;
+		}
+		if (ret < 0) {
+			LOGW("Bzip2 %s failed (%d)\n", mode ? "encode" : "decode", ret);
+			return false;
+		}
+		fn(outbuf, sizeof(outbuf) - strm.avail_out);
+	} while (strm.avail_out == 0);
+	return true;
+}
+
+LZMAStream::LZMAStream(int mode) : mode(mode), strm(LZMA_STREAM_INIT) {
+	lzma_options_lzma opt;
+	int ret;
+
+	// Initialize preset
+	lzma_lzma_preset(&opt, 9);
+	lzma_filter filters[] = {
+			{ .id = LZMA_FILTER_LZMA2, .options = &opt },
+			{ .id = LZMA_VLI_UNKNOWN, .options = nullptr },
+	};
+
+	switch(mode) {
+		case 0:
+			ret = lzma_auto_decoder(&strm, UINT64_MAX, 0);
+			break;
+		case 1:
+			ret = lzma_stream_encoder(&strm, filters, LZMA_CHECK_CRC32);
+			break;
+		case 2:
+			ret = lzma_alone_encoder(&strm, &opt);
+			break;
+	}
+}
+
+bool LZMAStream::update(const void *in, size_t size) {
+	return update(in, size, LZMA_RUN);
+}
+
+uint64_t LZMAStream::finalize() {
+	update(nullptr, 0, LZMA_FINISH);
+	uint64_t total = strm.total_out;
+	lzma_end(&strm);
+	return total;
+}
+
+bool LZMAStream::update(const void *in, size_t size, lzma_action flush) {
+	int ret;
+	strm.next_in = (uint8_t *) in;
+	strm.avail_in = size;
+	do {
+		strm.avail_out = sizeof(outbuf);
+		strm.next_out = outbuf;
+		ret = lzma_code(&strm, flush);
+		if (ret != LZMA_OK && ret != LZMA_STREAM_END) {
+			LOGW("LZMA %s failed (%d)\n", mode ? "encode" : "decode", ret);
+			return false;
+		}
+		fn(outbuf, sizeof(outbuf) - strm.avail_out);
+	} while (strm.avail_out == 0);
+	return true;
+}
+
+LZ4FDecoder::LZ4FDecoder() : outbuf(nullptr), total(0) {
+	LZ4F_createDecompressionContext(&ctx, LZ4F_VERSION);
+}
+
+LZ4FDecoder::~LZ4FDecoder() {
+	LZ4F_freeDecompressionContext(ctx);
+	delete[] outbuf;
+}
+
+bool LZ4FDecoder::update(const void *in, size_t size) {
+	auto inbuf = (const uint8_t *) in;
+	if (!outbuf)
+		read_header(inbuf, size);
+	size_t read, write;
+	LZ4F_errorCode_t ret;
+	do {
+		read = size;
+		write = outCapacity;
+		ret = LZ4F_decompress(ctx, outbuf, &write, inbuf, &read, nullptr);
+		if (LZ4F_isError(ret)) {
+			LOGW("LZ4 decode error: %s\n", LZ4F_getErrorName(ret));
+			return false;
+		}
+		size -= read;
+		inbuf += read;
+		total += write;
+		fn(outbuf, write);
+	} while (size != 0 || write != 0);
+	return true;
+}
+
+uint64_t LZ4FDecoder::finalize() {
+	return total;
+}
+
+void LZ4FDecoder::read_header(const uint8_t *&in, size_t &size) {
+	size_t read = size;
+	LZ4F_frameInfo_t info;
+	LZ4F_getFrameInfo(ctx, &info, in, &read);
+	switch (info.blockSizeID) {
+		case LZ4F_default:
+		case LZ4F_max64KB:  outCapacity = 1 << 16; break;
+		case LZ4F_max256KB: outCapacity = 1 << 18; break;
+		case LZ4F_max1MB:   outCapacity = 1 << 20; break;
+		case LZ4F_max4MB:   outCapacity = 1 << 22; break;
+	}
+	outbuf = new uint8_t[outCapacity];
+	in += read;
+	size -= read;
+}
+
+LZ4FEncoder::LZ4FEncoder() : outbuf(nullptr), outCapacity(0) {
+	LZ4F_createCompressionContext(&ctx, LZ4F_VERSION);
+}
+
+LZ4FEncoder::~LZ4FEncoder() {
+	LZ4F_freeCompressionContext(ctx);
+	delete[] outbuf;
+}
+
+bool LZ4FEncoder::update(const void *in, size_t size) {
+	if (!outbuf)
+		write_header();
+	auto inbuf = (const uint8_t *) in;
+	size_t read, write;
+	do {
+		read = size > CHUNK ? CHUNK : size;
+		write = LZ4F_compressUpdate(ctx, outbuf, outCapacity, inbuf, read, nullptr);
+		if (LZ4F_isError(write)) {
+			LOGW("LZ4 encode error: %s\n", LZ4F_getErrorName(write));
+			return false;
+		}
+		size -= read;
+		inbuf += read;
+		total += write;
+		fn(outbuf, write);
+	} while (size != 0);
+	return true;
+}
+
+uint64_t LZ4FEncoder::finalize() {
+	size_t write = LZ4F_compressEnd(ctx, outbuf, outCapacity, nullptr);
+	total += write;
+	fn(outbuf, write);
+	return total;
+}
+
+void LZ4FEncoder::write_header() {
+	LZ4F_preferences_t prefs {
+		.autoFlush = 1,
+		.compressionLevel = 9,
+		.frameInfo = {
+			.blockMode = LZ4F_blockIndependent,
+			.blockSizeID = LZ4F_max4MB,
+			.blockChecksumFlag = LZ4F_noBlockChecksum,
+			.contentChecksumFlag = LZ4F_contentChecksumEnabled
+		}
+	};
+	outCapacity = LZ4F_compressBound(CHUNK, &prefs);
+	outbuf = new uint8_t[outCapacity];
+	size_t write = LZ4F_compressBegin(ctx, outbuf, outCapacity, &prefs);
+	total += write;
+	fn(outbuf, write);
+}
+
+LZ4Decoder::LZ4Decoder() : init(false), buf_off(0), total(0), block_sz(0) {
+	outbuf = new char[LZ4_UNCOMPRESSED];
+	buf = new char[LZ4_COMPRESSED];
+}
+
+LZ4Decoder::~LZ4Decoder() {
+	delete[] outbuf;
+	delete[] buf;
+}
+
+bool LZ4Decoder::update(const void *in, size_t size) {
+	const char *inbuf = (const char *) in;
+	if (!init) {
+		// Skip magic
+		inbuf += 4;
+		size -= 4;
+		init = true;
+	}
+	int write;
+	size_t consumed;
+	do {
+		if (block_sz == 0) {
+			block_sz = *((unsigned *) inbuf);
+			inbuf += sizeof(unsigned);
+			size -= sizeof(unsigned);
+		} else if (buf_off + size >= block_sz) {
+			consumed = block_sz - buf_off;
+			memcpy(buf + buf_off, inbuf, consumed);
+			inbuf += consumed;
+			size -= consumed;
+
+			write = LZ4_decompress_safe(buf, outbuf, block_sz, LZ4_UNCOMPRESSED);
+			if (write < 0) {
+				LOGW("LZ4HC decompression failure (%d)\n", write);
+				return false;
+			}
+			fn(outbuf, write);
+			total += write;
+
+			// Reset
+			buf_off = 0;
+			block_sz = 0;
+		} else {
+			// Copy to internal buffer
+			memcpy(buf + buf_off, inbuf, size);
+			buf_off += size;
+			size = 0;
+		}
+	} while (size != 0);
+	return true;
+}
+
+uint64_t LZ4Decoder::finalize() {
+	return total;
+}
+
+LZ4Encoder::LZ4Encoder() : init(false), buf_off(0), out_total(0), in_total(0) {
+	outbuf = new char[LZ4_COMPRESSED];
+	buf = new char[LZ4_UNCOMPRESSED];
+}
+
+LZ4Encoder::~LZ4Encoder() {
+	delete[] outbuf;
+	delete[] buf;
+}
+
+bool LZ4Encoder::update(const void *in, size_t size) {
+	if (!init) {
+		fn("\x02\x21\x4c\x18", 4);
+		init = true;
+	}
+	in_total += size;
+	const char *inbuf = (const char *) in;
+	size_t consumed;
+	int write;
+	do {
+		if (buf_off + size >= LZ4_UNCOMPRESSED) {
+			consumed = LZ4_UNCOMPRESSED - buf_off;
+			memcpy(buf + buf_off, inbuf, consumed);
+			inbuf += consumed;
+			size -= consumed;
+
+			write = LZ4_compress_HC(buf, outbuf, LZ4_UNCOMPRESSED, LZ4_COMPRESSED, 9);
+			if (write == 0) {
+				LOGW("LZ4HC compression failure\n");
+				return false;
+			}
+			fn(&write, sizeof(write));
+			fn(outbuf, write);
+			out_total += write + sizeof(write);
+
+			// Reset buffer
+			buf_off = 0;
+		} else {
+			// Copy to internal buffer
+			memcpy(buf + buf_off, inbuf, size);
+			buf_off += size;
+			size = 0;
+		}
+	} while (size != 0);
+	return true;
+}
+
+uint64_t LZ4Encoder::finalize() {
+	if (buf_off) {
+		int write = LZ4_compress_HC(buf, outbuf, buf_off, LZ4_COMPRESSED, 9);
+		fn(&write, sizeof(write));
+		fn(outbuf, write);
+		out_total += write + sizeof(write);
+	}
+	fn(&in_total, sizeof(in_total));
+	return out_total + sizeof(in_total);
+}

native/jni/magiskboot/compress.h

@@ -0,0 +1,187 @@
+#pragma once
+
+#include <functional>
+
+#include <zlib.h>
+#include <bzlib.h>
+#include <lzma.h>
+#include <lz4.h>
+#include <lz4frame.h>
+#include <lz4hc.h>
+
+#include "format.h"
+
+#define CHUNK 0x40000
+
+class Compression {
+public:
+	virtual ~Compression() = default;
+	void set_outfn(std::function<void(const void *, size_t)> &&fn);
+	void set_outfd(int fd);
+	int64_t one_step(int outfd, const void *in, size_t size);
+	virtual bool update(const void *in, size_t size) = 0;
+	virtual uint64_t finalize() = 0;
+
+	template <class T>
+	static int64_t one_step(int outfd, const void *in, size_t size) {
+		T cmp;
+		return cmp.one_step(outfd, in, size);
+	}
+
+protected:
+	Compression();
+	std::function<void (const void*, size_t)> fn;
+};
+
+class GZStream : public Compression {
+public:
+	bool update(const void *in, size_t size) override;
+	uint64_t finalize() override;
+
+protected:
+	explicit GZStream(int mode);
+
+private:
+	int mode;
+	z_stream strm;
+	uint8_t outbuf[CHUNK];
+
+	bool update(const void *in, size_t size, int flush);
+};
+
+class GZDecoder : public GZStream {
+public:
+	GZDecoder() : GZStream(0) {};
+};
+
+class GZEncoder : public GZStream {
+public:
+	GZEncoder() : GZStream(1) {};
+};
+
+class BZStream : public Compression {
+public:
+	bool update(const void *in, size_t size) override;
+	uint64_t finalize() override;
+
+protected:
+	explicit BZStream(int mode);
+
+private:
+	int mode;
+	bz_stream strm;
+	char outbuf[CHUNK];
+
+	bool update(const void *in, size_t size, int flush);
+};
+
+class BZDecoder : public BZStream {
+public:
+	BZDecoder() : BZStream(0) {};
+};
+
+class BZEncoder : public BZStream {
+public:
+	BZEncoder() : BZStream(1) {};
+};
+
+class LZMAStream : public Compression {
+public:
+	bool update(const void *in, size_t size) override;
+	uint64_t finalize() override;
+
+protected:
+	explicit LZMAStream(int mode);
+
+private:
+	int mode;
+	lzma_stream strm;
+	uint8_t outbuf[CHUNK];
+
+	bool update(const void *in, size_t size, lzma_action flush);
+};
+
+class LZMADecoder : public LZMAStream {
+public:
+	LZMADecoder() : LZMAStream(0) {}
+};
+
+class XZEncoder : public LZMAStream {
+public:
+	XZEncoder() : LZMAStream(1) {}
+};
+
+class LZMAEncoder : public LZMAStream {
+public:
+	LZMAEncoder() : LZMAStream(2) {}
+};
+
+class LZ4FDecoder : public Compression {
+public:
+	LZ4FDecoder();
+	~LZ4FDecoder() override;
+	bool update(const void *in, size_t size) override;
+	uint64_t finalize() override;
+
+private:
+	LZ4F_decompressionContext_t ctx;
+	uint8_t *outbuf;
+	size_t outCapacity;
+	uint64_t total;
+
+	void read_header(const uint8_t *&in, size_t &size);
+};
+
+class LZ4FEncoder : public Compression {
+public:
+	LZ4FEncoder();
+	~LZ4FEncoder() override;
+	bool update(const void *in, size_t size) override;
+	uint64_t finalize() override;
+
+private:
+	LZ4F_compressionContext_t ctx;
+	uint8_t *outbuf;
+	size_t outCapacity;
+	uint64_t total;
+
+	void write_header();
+};
+
+#define LZ4_UNCOMPRESSED 0x800000
+#define LZ4_COMPRESSED   LZ4_COMPRESSBOUND(LZ4_UNCOMPRESSED)
+
+class LZ4Decoder : public Compression {
+public:
+	LZ4Decoder();
+	~LZ4Decoder() override;
+	bool update(const void *in, size_t size) override;
+	uint64_t finalize() override;
+
+private:
+	char *outbuf;
+	char *buf;
+	bool init;
+	unsigned block_sz;
+	int buf_off;
+	uint64_t total;
+};
+
+class LZ4Encoder : public Compression {
+public:
+	LZ4Encoder();
+	~LZ4Encoder() override;
+	bool update(const void *in, size_t size) override;
+	uint64_t finalize() override;
+
+private:
+	char *outbuf;
+	char *buf;
+	bool init;
+	int buf_off;
+	uint64_t out_total;
+	unsigned in_total;
+};
+
+Compression *get_encoder(format_t type);
+Compression *get_decoder(format_t type);

native/jni/magiskboot/magiskboot.h

@@ -23,13 +23,8 @@ void decompress(char *from, const char *to);
 int dtb_commands(const char *cmd, int argc, char *argv[]);
 
 // Compressions
-size_t gzip(int mode, int fd, const void *buf, size_t size);
+int64_t compress(format_t type, int fd, const void *from, size_t size);
-size_t lzma(int mode, int fd, const void *buf, size_t size);
+int64_t decompress(format_t type, int fd, const void *from, size_t size);
-size_t lz4(int mode, int fd, const uint8_t *buf, size_t size);
-size_t bzip2(int mode, int fd, const void *buf, size_t size);
-size_t lz4_legacy(int mode, int fd, const uint8_t *buf, size_t size);
-long long compress(format_t type, int fd, const void *from, size_t size);
-long long decompress(format_t type, int fd, const void *from, size_t size);
 
 // Pattern
 int patch_verity(void **buf, uint32_t *size, int patch);

native/jni/magiskboot/main.cpp

@@ -139,13 +139,13 @@ int main(int argc, char *argv[]) {
 	} else if (argc > 2 && strcmp(argv[1], "--repack") == 0) {
 		repack(argv[2], argc > 3 ? argv[3] : NEW_BOOT);
 	} else if (argc > 2 && strcmp(argv[1], "--decompress") == 0) {
-		decompress(argv[2], argc > 3 ? argv[3] : NULL);
+		decompress(argv[2], argc > 3 ? argv[3] : nullptr);
 	} else if (argc > 2 && strncmp(argv[1], "--compress", 10) == 0) {
 		const char *method;
 		method = strchr(argv[1], '=');
-		if (method == NULL) method = "gzip";
+		if (method == nullptr) method = "gzip";
 		else method++;
-		compress(method, argv[2], argc > 3 ? argv[3] : NULL);
+		compress(method, argv[2], argc > 3 ? argv[3] : nullptr);
 	} else if (argc > 4 && strcmp(argv[1], "--hexpatch") == 0) {
 		hexpatch(argv[2], argv[3], argv[4]);
 	} else if (argc > 2 && strcmp(argv[1], "--cpio") == 0) {