From 0d51997e468bd2a229839e956e07ae55d75a5e10 Mon Sep 17 00:00:00 2001
From: topjohnwu <topjohnwu@gmail.com>
Date: Tue, 28 Feb 2017 21:56:13 +0800
Subject: [PATCH] Add xz support

---
 jni/magiskboot/Android.mk   |  1 +
 jni/magiskboot/compress.c   | 60 ++++++++++++++++++++++++++++++++++++-
 jni/magiskboot/magiskboot.h |  2 +-
 jni/magiskboot/unpack.c     |  1 +
 4 files changed, 62 insertions(+), 2 deletions(-)

diff --git a/jni/magiskboot/Android.mk b/jni/magiskboot/Android.mk
index 3d9bc7657..a8f38a47f 100644
--- a/jni/magiskboot/Android.mk
+++ b/jni/magiskboot/Android.mk
@@ -3,6 +3,7 @@ LOCAL_PATH := $(call my-dir)
 include $(CLEAR_VARS)
 LOCAL_MODULE := magiskboot
 LOCAL_STATIC_LIBRARIES := liblzma
+LOCAL_C_INCLUDES := jni/ndk-compression/xz/src/liblzma/api/
 LOCAL_SRC_FILES := main.c unpack.c repack.c hexpatch.c parseimg.c compress.c
 LOCAL_LDLIBS += -lz
 include $(BUILD_EXECUTABLE)
diff --git a/jni/magiskboot/compress.c b/jni/magiskboot/compress.c
index b10f8a5c6..2b17d2524 100644
--- a/jni/magiskboot/compress.c
+++ b/jni/magiskboot/compress.c
@@ -1,3 +1,6 @@
+#include <zlib.h>
+#include <lzma.h>
+
 #include "magiskboot.h"
 
 void gzip(int dec, const char* filename, unsigned char* buf, size_t size) {
@@ -25,7 +28,7 @@ void gzip(int dec, const char* filename, unsigned char* buf, size_t size) {
 
 	do {
 		strm.next_in = buf + pos;
-		if (pos + CHUNK > size) {
+		if (pos + CHUNK >= size) {
 			strm.avail_in = size - pos;
 			pos = size;
 			flush = Z_FINISH;
@@ -57,3 +60,58 @@ void gzip(int dec, const char* filename, unsigned char* buf, size_t size) {
 	}
 	close(fd);
 }
+
+void lzma(int dec, const char* filename, unsigned char* buf, size_t size) {
+	int have, pos = 0;
+	lzma_ret ret;
+	lzma_stream strm = LZMA_STREAM_INIT;
+	lzma_options_lzma opt;
+	lzma_action action = LZMA_RUN;
+	unsigned char out[BUFSIZ];
+
+	int fd = open(filename, O_WRONLY | O_CREAT | O_TRUNC, 0644);
+
+	if (fd < 0)
+		error(1, "Unable to create %s", filename);
+
+	if (dec) {
+		ret = lzma_auto_decoder(&strm, UINT64_MAX, 0);
+	} else {
+		lzma_lzma_preset(&opt, LZMA_PRESET_DEFAULT);
+		lzma_filter filters[] = {
+			{ .id = LZMA_FILTER_LZMA2, .options = &opt },
+			{ .id = LZMA_VLI_UNKNOWN, .options = NULL },
+		};
+		ret = lzma_stream_encoder(&strm, filters, LZMA_CHECK_CRC64);
+	}
+
+	if (ret != LZMA_OK)
+		error(1, "Unable to init lzma");
+
+	do {
+		strm.next_in = buf + pos;
+		if (pos + BUFSIZ >= size) {
+			strm.avail_in = size - pos;
+			pos = size;
+			action = LZMA_FINISH;
+		} else {
+			strm.avail_in = BUFSIZ;
+			pos += BUFSIZ;
+		}
+
+		do {
+			strm.avail_out = BUFSIZ;
+			strm.next_out = out;
+			ret = lzma_code(&strm, action);
+			have = BUFSIZ - strm.avail_out;
+			if (write(fd, out, have) != have)
+				error(1, "Error in writing %s", filename);
+		} while (strm.avail_out == 0 && ret == LZMA_OK);
+
+		if (ret != LZMA_OK && ret != LZMA_STREAM_END)
+			error(1, "LZMA error %d!", ret);
+
+	} while (pos < size);
+
+	lzma_end(&strm);
+}
diff --git a/jni/magiskboot/magiskboot.h b/jni/magiskboot/magiskboot.h
index 99d68cd64..daabc7d8c 100644
--- a/jni/magiskboot/magiskboot.h
+++ b/jni/magiskboot/magiskboot.h
@@ -10,7 +10,6 @@
 #include <sys/mman.h>
 #include <fcntl.h>
 #include <string.h>
-#include <zlib.h>
 
 #include "bootimg.h"
 
@@ -58,5 +57,6 @@ void parse_img(unsigned char *orig, size_t size);
 
 // Compressions
 void gzip(int dec, const char* filename, unsigned char* buf, size_t size);
+void lzma(int dec, const char* filename, unsigned char* buf, size_t size);
 
 #endif
diff --git a/jni/magiskboot/unpack.c b/jni/magiskboot/unpack.c
index 8d7961055..fb94377ba 100644
--- a/jni/magiskboot/unpack.c
+++ b/jni/magiskboot/unpack.c
@@ -51,6 +51,7 @@ void unpack(const char* image) {
 			break;
 		case XZ:
 			sprintf(name, "%s.%s", RAMDISK_FILE, "xz");
+			lzma(1, RAMDISK_FILE, ramdisk, hdr.ramdisk_size);
 			break;
 		case LZMA:
 			sprintf(name, "%s.%s", RAMDISK_FILE, "lzma");