Refactor MagiskBoot compression methods

jni/magiskboot/bootimg.c

@@ -148,11 +148,13 @@ void unpack(const char* image) {
 		ramdisk += 512;
 		hdr.ramdisk_size -= 512;
 	}
-	if (decomp(ramdisk_type, RAMDISK_FILE, ramdisk, hdr.ramdisk_size)) {
+	int fd = open_new(RAMDISK_FILE);
+	if (decomp(ramdisk_type, fd, ramdisk, hdr.ramdisk_size) < 0) {
 		// Dump the compressed ramdisk
 		dump(ramdisk, hdr.ramdisk_size, RAMDISK_FILE ".raw");
 		LOGE(1, "Unknown ramdisk format! Dumped to %s\n", RAMDISK_FILE ".raw");
 	}
+	close(fd);
 
 	if (hdr.second_size) {
 		// Dump second
@@ -171,7 +173,6 @@ void unpack(const char* image) {
 void repack(const char* orig_image, const char* out_image) {
 	size_t size;
 	void *orig;
-	char name[PATH_MAX];
 
 	// There are possible two MTK headers
 	mtk_hdr mtk_kernel_hdr, mtk_ramdisk_hdr;
@@ -198,52 +199,51 @@ void repack(const char* orig_image, const char* out_image) {
 	// Skip a page for header
 	write_zero(fd, hdr.page_size);
 
-	// Restore kernel
 	if (mtk_kernel) {
 		mtk_kernel_off = lseek(fd, 0, SEEK_CUR);
-		restore_buf(fd, kernel, 512);
+		// Skip header
+		write_zero(fd, 512);
 		memcpy(&mtk_kernel_hdr, kernel, sizeof(mtk_kernel_hdr));
 	}
 	hdr.kernel_size = restore(KERNEL_FILE, fd);
 	file_align(fd, hdr.page_size, 1);
 
-	// Restore ramdisk
 	if (mtk_ramdisk) {
 		mtk_ramdisk_off = lseek(fd, 0, SEEK_CUR);
-		restore_buf(fd, ramdisk, 512);
+		// Skip header
+		write_zero(fd, 512);
 		memcpy(&mtk_ramdisk_hdr, ramdisk, sizeof(mtk_ramdisk_hdr));
 	}
 	if (access(RAMDISK_FILE, R_OK) == 0) {
 		// If we found raw cpio, compress to original format
-
-		// Before we start, clean up previous compressed files
-		for (int i = 0; SUP_EXT_LIST[i]; ++i) {
-			sprintf(name, "%s.%s", RAMDISK_FILE, SUP_EXT_LIST[i]);
-			unlink(name);
-		}
-
 		size_t cpio_size;
 		void *cpio;
 		mmap_ro(RAMDISK_FILE, &cpio, &cpio_size);
 
-		if (comp(ramdisk_type, RAMDISK_FILE, cpio, cpio_size))
+		long long ramdisk_size = comp(ramdisk_type, fd, cpio, cpio_size);
+
+		if (ramdisk_size < 0) {
 			LOGE(1, "Unsupported ramdisk format!\n");
+		} else {
+			hdr.ramdisk_size = ramdisk_size;
+		}
 
 		munmap(cpio, cpio_size);
-	}
+	} else {
-
+		// Find compressed ramdisk
-	int found = 0;
+		char name[PATH_MAX];
-	for (int i = 0; SUP_EXT_LIST[i]; ++i) {
+		int found = 0;
-		sprintf(name, "%s.%s", RAMDISK_FILE, SUP_EXT_LIST[i]);
+		for (int i = 0; SUP_EXT_LIST[i]; ++i) {
-		if (access(name, R_OK) == 0) {
+			sprintf(name, "%s.%s", RAMDISK_FILE, SUP_EXT_LIST[i]);
-			ramdisk_type = SUP_TYPE_LIST[i];
+			if (access(name, R_OK) == 0) {
-			found = 1;
+				found = 1;
-			break;
+				break;
+			}
 		}
+		if (!found)
+			LOGE(1, "No ramdisk exists!\n");
+		hdr.ramdisk_size = restore(name, fd);
 	}
-	if (!found)
-		LOGE(1, "No ramdisk exists!\n");
-	hdr.ramdisk_size = restore(name, fd);
 	file_align(fd, hdr.page_size, 1);
 
 	// Restore second
@@ -266,7 +266,7 @@ void repack(const char* orig_image, const char* out_image) {
 		}
 	}
 
-	// Write headers back
+	// Write MTK headers back
 	if (mtk_kernel) {
 		lseek(fd, mtk_kernel_off, SEEK_SET);
 		mtk_kernel_hdr.size = hdr.kernel_size;

jni/magiskboot/compress.c

@@ -15,30 +15,12 @@
 #define LZ4_FOOTER_SIZE 4
 #define LZ4_LEGACY_BLOCKSIZE  0x800000
 
-static void write_file(const int fd, const void *buf, const size_t size, const char *filename) {
-	xwrite(fd, buf, size);
-}
-
-static void report(const int mode, const char* filename) {
-	switch(mode) {
-		case 0:
-			fprintf(stderr, "Decompressing to [%s]\n\n", filename);
-			break;
-		default:
-			fprintf(stderr, "Compressing to [%s]\n\n", filename);
-			break;
-	}
-}
-
 // Mode: 0 = decode; 1 = encode
-void gzip(int mode, const char* filename, const void *buf, size_t size) {
+size_t gzip(int mode, int fd, const void *buf, size_t size) {
-	size_t ret = 0, flush, have, pos = 0;
+	size_t ret = 0, flush, have, pos = 0, total = 0;
 	z_stream strm;
 	unsigned char out[CHUNK];
 
-	report(mode, filename);
-	int fd = open_new(filename);
-
 	strm.zalloc = Z_NULL;
 	strm.zfree = Z_NULL;
 	strm.opaque = Z_NULL;
@@ -83,7 +65,7 @@ void gzip(int mode, const char* filename, const void *buf, size_t size) {
 				LOGE(1, "Error when running gzip\n");
 
 			have = CHUNK - strm.avail_out;
-			write_file(fd, out, have, filename);
+			total += xwrite(fd, out, have);
 
 		} while (strm.avail_out == 0);
 
@@ -97,22 +79,19 @@ void gzip(int mode, const char* filename, const void *buf, size_t size) {
 			deflateEnd(&strm);
 			break;
 	}
-	close(fd);
+	return total;
 }
 
 
 // Mode: 0 = decode xz/lzma; 1 = encode xz; 2 = encode lzma
-void lzma(int mode, const char* filename, const void *buf, size_t size) {
+size_t lzma(int mode, int fd, const void *buf, size_t size) {
-	size_t have, pos = 0;
+	size_t have, pos = 0, total = 0;
 	lzma_ret ret = 0;
 	lzma_stream strm = LZMA_STREAM_INIT;
 	lzma_options_lzma opt;
 	lzma_action action;
 	unsigned char out[BUFSIZ];
 
-	report(mode, filename);
-	int fd = open_new(filename);
-
 	// Initialize preset
 	lzma_lzma_preset(&opt, LZMA_PRESET_DEFAULT);
 	lzma_filter filters[] = {
@@ -154,7 +133,7 @@ void lzma(int mode, const char* filename, const void *buf, size_t size) {
 			strm.next_out = out;
 			ret = lzma_code(&strm, action);
 			have = BUFSIZ - strm.avail_out;
-			write_file(fd, out, have, filename);
+			total += xwrite(fd, out, have);
 		} while (strm.avail_out == 0 && ret == LZMA_OK);
 
 		if (ret != LZMA_OK && ret != LZMA_STREAM_END)
@@ -163,22 +142,19 @@ void lzma(int mode, const char* filename, const void *buf, size_t size) {
 	} while (pos < size);
 
 	lzma_end(&strm);
-	close(fd);
+	return total;
 }
 
 // Mode: 0 = decode; 1 = encode
-void lz4(int mode, const char* filename, const void *buf, size_t size) {
+size_t lz4(int mode, int fd, const void *buf, size_t size) {
 	LZ4F_decompressionContext_t dctx;
 	LZ4F_compressionContext_t cctx;
 	LZ4F_frameInfo_t info;
 
-	size_t outCapacity, avail_in, ret = 0, pos = 0;
+	size_t outCapacity, avail_in, ret = 0, pos = 0, total = 0;
 	size_t have, read;
 	void *out = NULL;
 
-	report(mode, filename);
-	int fd = open_new(filename);
-
 	// Initialize context
 	switch(mode) {
 		case 0:
@@ -225,7 +201,7 @@ void lz4(int mode, const char* filename, const void *buf, size_t size) {
 		have = ret = LZ4F_compressBegin(cctx, out, size, NULL);
 		if (LZ4F_isError(ret))
 			LOGE(1, "Failed to start compression: error %s\n", LZ4F_getErrorName(ret));
-		write_file(fd, out, have, filename);
+		total += xwrite(fd, out, have);
 	}
 
 	do {
@@ -249,7 +225,7 @@ void lz4(int mode, const char* filename, const void *buf, size_t size) {
 			if (LZ4F_isError(ret))
 				LOGE(1, "LZ4 coding error: %s\n", LZ4F_getErrorName(ret));
 
-			write_file(fd, out, have, filename);
+			total += xwrite(fd, out, have);
 			// Update status
 			pos += read;
 			avail_in -= read;
@@ -266,25 +242,22 @@ void lz4(int mode, const char* filename, const void *buf, size_t size) {
 			if (LZ4F_isError(ret))
 				LOGE(1, "Failed to end compression: error %s\n", LZ4F_getErrorName(ret));
 
-			write_file(fd, out, have, filename);
+			total += xwrite(fd, out, have);
 
 			LZ4F_freeCompressionContext(cctx);
 			break;
 	}
 
 	free(out);
-	close(fd);
+	return total;
 }
 
 // Mode: 0 = decode; 1 = encode
-void bzip2(int mode, const char* filename, const void* buf, size_t size) {
+size_t bzip2(int mode, int fd, const void* buf, size_t size) {
-	size_t ret = 0, action, have, pos = 0;
+	size_t ret = 0, action, have, pos = 0, total = 0;
 	bz_stream strm;
 	char out[CHUNK];
 
-	report(mode, filename);
-	int fd = open_new(filename);
-
 	strm.bzalloc = NULL;
 	strm.bzfree = NULL;
 	strm.opaque = NULL;
@@ -327,7 +300,7 @@ void bzip2(int mode, const char* filename, const void* buf, size_t size) {
 			}
 
 			have = CHUNK - strm.avail_out;
-			write_file(fd, out, have, filename);
+			total += xwrite(fd, out, have);
 
 		} while (strm.avail_out == 0);
 
@@ -341,21 +314,17 @@ void bzip2(int mode, const char* filename, const void* buf, size_t size) {
 			BZ2_bzCompressEnd(&strm);
 			break;
 	}
-	close(fd);
+	return total;
 }
 
 // Mode: 0 = decode; 1 = encode
-void lz4_legacy(int mode, const char* filename, const void* buf, size_t size) {
+size_t lz4_legacy(int mode, int fd, const void* buf, size_t size) {
-	size_t pos = 0;
+	size_t pos = 0, total = 0;
 	int have;
 	char *out;
 	unsigned block_size, insize;
 	unsigned char block_size_le[4];
 
-
-	report(mode, filename);
-	int fd = open_new(filename);
-
 	switch(mode) {
 		case 0:
 			out = xmalloc(LZ4_LEGACY_BLOCKSIZE);
@@ -365,7 +334,7 @@ void lz4_legacy(int mode, const char* filename, const void* buf, size_t size) {
 		case 1:
 			out = xmalloc(LZ4_COMPRESSBOUND(LZ4_LEGACY_BLOCKSIZE));
 			// Write magic
-			write_file(fd, "\x02\x21\x4c\x18", 4, filename);
+			total += xwrite(fd, "\x02\x21\x4c\x18", 4);
 			break;
 		default:
 			LOGE(1, "Unsupported lz4_legacy mode!\n");
@@ -400,88 +369,62 @@ void lz4_legacy(int mode, const char* filename, const void* buf, size_t size) {
 				block_size_le[1] = (have >> 8) & 0xff;
 				block_size_le[2] = (have >> 16) & 0xff;
 				block_size_le[3] = (have >> 24) & 0xff;
-				write_file(fd, block_size_le, 4, filename);
+				total += xwrite(fd, block_size_le, 4);
 				break;
 		}
 		// Write main data
-		write_file(fd, out, have, filename);
+		total += xwrite(fd, out, have);
 	} while(pos < size);
 
 	free(out);
-	close(fd);
+	return total;
 }
 
-int decomp(file_t type, const char *to, const void *from, size_t size) {
+long long decomp(file_t type, int to, const void *from, size_t size) {
 	switch (type) {
 		case GZIP:
-			gzip(0, to, from, size);
+			return gzip(0, to, from, size);
-			break;
 		case XZ:
-			lzma(0, to, from, size);
+			return lzma(0, to, from, size);
-			break;
 		case LZMA:
-			lzma(0, to, from, size);
+			return lzma(0, to, from, size);
-			break;
 		case BZIP2:
-			bzip2(0, to, from, size);
+			return bzip2(0, to, from, size);
-			break;
 		case LZ4:
-			lz4(0, to, from, size);
+			return lz4(0, to, from, size);
-			break;
 		case LZ4_LEGACY:
-			lz4_legacy(0, to, from, size);
+			return lz4_legacy(0, to, from, size);
-			break;
 		default:
 			// Unsupported
-			return 1;
+			return -1;
 	}
-	return 0;
 }
 
 // Output will be to.ext
-int comp(file_t type, const char *to, const void *from, size_t size) {
+long long comp(file_t type, int to, const void *from, size_t size) {
-	char name[PATH_MAX];
-	const char *ext = strrchr(to, '.');
-	if (ext == NULL) ext = to;
-	strcpy(name, to);
 	switch (type) {
 		case GZIP:
-			if (strcmp(ext, ".gz") != 0)
+			return gzip(1, to, from, size);
-				sprintf(name, "%s.%s", to, "gz");
-			gzip(1, name, from, size);
-			break;
 		case XZ:
-			if (strcmp(ext, ".xz") != 0)
+			return lzma(1, to, from, size);
-				sprintf(name, "%s.%s", to, "xz");
-			lzma(1, name, from, size);
-			break;
 		case LZMA:
-			if (strcmp(ext, ".lzma") != 0)
+			return lzma(2, to, from, size);
-				sprintf(name, "%s.%s", to, "lzma");
-			lzma(2, name, from, size);
-			break;
 		case BZIP2:
-			if (strcmp(ext, ".bz2") != 0)
+			return bzip2(1, to, from, size);
-				sprintf(name, "%s.%s", to, "bz2");
-			bzip2(1, name, from, size);
-			break;
 		case LZ4:
-			if (strcmp(ext, ".lz4") != 0)
+			return lz4(1, to, from, size);
-				sprintf(name, "%s.%s", to, "lz4");
-			lz4(1, name, from, size);
-			break;
 		case LZ4_LEGACY:
-			if (strcmp(ext, ".lz4") != 0)
+			return lz4_legacy(1, to, from, size);
-				sprintf(name, "%s.%s", to, "lz4");
-			lz4_legacy(1, name, from, size);
-			break;
 		default:
 			// Unsupported
-			return 1;
+			return -1;
 	}
-	return 0;
 }
 
+/*
+ * Below are utility functions for commandline
+ */
+
 void decomp_file(char *from, const char *to) {
 	int ok = 1;
 	void *file;
@@ -525,7 +468,10 @@ void decomp_file(char *from, const char *to) {
 			*ext = '\0';
 			to = from;
 		}
-		decomp(type, to, file, size);
+		int fd = open_new(to);
+		fprintf(stderr, "Decompressing to [%s]\n\n", to);
+		decomp(type, fd, file, size);
+		close(fd);
 		if (to == from) {
 			*ext = '.';
 			unlink(from);
@@ -538,18 +484,25 @@ void decomp_file(char *from, const char *to) {
 
 void comp_file(const char *method, const char *from, const char *to) {
 	file_t type;
+	char *ext, dest[PATH_MAX];
 	if (strcmp(method, "gzip") == 0) {
 		type = GZIP;
+		ext = "gz";
 	} else if (strcmp(method, "xz") == 0) {
 		type = XZ;
+		ext = "xz";
 	} else if (strcmp(method, "lzma") == 0) {
 		type = LZMA;
+		ext = "lzma";
 	} else if (strcmp(method, "lz4") == 0) {
 		type = LZ4;
+		ext = "lz4";
 	} else if (strcmp(method, "lz4_legacy") == 0) {
 		type = LZ4_LEGACY;
+		ext = "lz4";
 	} else if (strcmp(method, "bzip2") == 0) {
 		type = BZIP2;
+		ext = "bz2";
 	} else {
 		fprintf(stderr, "Only support following methods: ");
 		for (int i = 0; SUP_LIST[i]; ++i)
@@ -560,9 +513,14 @@ void comp_file(const char *method, const char *from, const char *to) {
 	void *file;
 	size_t size;
 	mmap_ro(from, &file, &size);
-	if (!to)
+	if (!to) {
-		to = from;
+		snprintf(dest, sizeof(dest), "%s.%s", from, ext);
-	comp(type, to, file, size);
+		to = dest;
+	}
+	fprintf(stderr, "Compressing to [%s]\n\n", to);
+	int fd = open_new(to);
+	comp(type, fd, file, size);
+	close(fd);
 	munmap(file, size);
 	if (to == from)
 		unlink(from);

jni/magiskboot/magiskboot.h

@@ -37,18 +37,18 @@ void repack(const char* orig_image, const char* out_image);
 void hexpatch(const char *image, const char *from, const char *to);
 int parse_img(void *orig, size_t size);
 int cpio_commands(const char *command, int argc, char *argv[]);
+void comp_file(const char *method, const char *from, const char *to);
+void decomp_file(char *from, const char *to);
 void cleanup();
 
 // Compressions
-void gzip(int mode, const char* filename, const void *buf, size_t size);
+size_t gzip(int mode, int fd, const void *buf, size_t size);
-void lzma(int mode, const char* filename, const void *buf, size_t size);
+size_t lzma(int mode, int fd, const void *buf, size_t size);
-void lz4(int mode, const char* filename, const void *buf, size_t size);
+size_t lz4(int mode, int fd, const void *buf, size_t size);
-void bzip2(int mode, const char* filename, const void *buf, size_t size);
+size_t bzip2(int mode, int fd, const void *buf, size_t size);
-void lz4_legacy(int mode, const char* filename, const void *buf, size_t size);
+size_t lz4_legacy(int mode, int fd, const void *buf, size_t size);
-int comp(file_t type, const char *to, const void *from, size_t size);
+long long comp(file_t type, int to, const void *from, size_t size);
-void comp_file(const char *method, const char *from, const char *to);
+long long decomp(file_t type, int to, const void *from, size_t size);
-int decomp(file_t type, const char *to, const void *from, size_t size);
-void decomp_file(char *from, const char *to);
 
 // Utils
 void mmap_ro(const char *filename, void **buf, size_t *size);