Add support for writing a gzip compressed image

Change-Id: Ica2ff90060f6a4ced2c671084205b58eede66cdf
diff --git a/ext4_utils/Android.mk b/ext4_utils/Android.mk
index c8bda90..819e634 100644
--- a/ext4_utils/Android.mk
+++ b/ext4_utils/Android.mk
@@ -8,6 +8,7 @@
         ext_utils.c \
         allocate.c \
         backed_block.c \
+        output_file.c \
         contents.c \
         extent.c \
         indirect.c \
@@ -17,6 +18,8 @@
 LOCAL_SRC_FILES := $(make_ext4fs_src_files)
 LOCAL_MODULE := make_ext4fs
 LOCAL_MODULE_TAGS := optional
+LOCAL_C_INCLUDES += external/zlib
+LOCAL_SHARED_LIBRARIES += libz
 
 include $(BUILD_EXECUTABLE)
 
@@ -24,5 +27,6 @@
 
 LOCAL_SRC_FILES := $(make_ext4fs_src_files)
 LOCAL_MODULE := make_ext4fs
+LOCAL_STATIC_LIBRARIES += libz
 
 include $(BUILD_HOST_EXECUTABLE)
diff --git a/ext4_utils/backed_block.c b/ext4_utils/backed_block.c
index e438498..6ec778a 100644
--- a/ext4_utils/backed_block.c
+++ b/ext4_utils/backed_block.c
@@ -102,7 +102,7 @@
 /* Iterates over the queued data blocks, calling data_func for each contiguous
    data block, and file_func for each contiguous file block */
 void for_each_data_block(data_block_callback_t data_func,
-        data_block_file_callback_t file_func, void *priv)
+        data_block_file_callback_t file_func, struct output_file *out)
 {
 	struct data_block *db;
 	u32 last_block = 0;
@@ -113,9 +113,9 @@
 		last_block = db->block + DIV_ROUND_UP(db->len, info.block_size) - 1;
 
 		if (db->filename)
-			file_func(priv, db->block, db->filename, db->offset, db->len);
+			file_func(out, db->block * info.block_size, db->filename, db->offset, db->len);
 		else
-			data_func(priv, db->block, db->data, db->len);
+			data_func(out, db->block * info.block_size, db->data, db->len);
 	}
 }
 
diff --git a/ext4_utils/backed_block.h b/ext4_utils/backed_block.h
index fc76229..9bcfd44 100644
--- a/ext4_utils/backed_block.h
+++ b/ext4_utils/backed_block.h
@@ -20,17 +20,19 @@
 #include <sys/types.h>
 #include <unistd.h>
 #include "ext4_utils.h"
+#include "output_file.h"
 
-typedef void (*data_block_callback_t)(void *priv, u32 block, u8 *data,
-	int len);
-typedef void (*data_block_file_callback_t)(void *priv, u32 block,
-	const char *file, off_t offset, int len);
+typedef void (*data_block_callback_t)(struct output_file *out, u64 off,
+				      u8 *data, int len);
+typedef void (*data_block_file_callback_t)(struct output_file *out, u64 off,
+					   const char *file, off_t offset,
+					   int len);
 
 void queue_data_block(u8 *data, u32 len, u32 block);
 void queue_data_file(const char *filename, off_t offset, u32 len,
         u32 block);
 void for_each_data_block(data_block_callback_t data_func,
-        data_block_file_callback_t file_func, void *priv);
+        data_block_file_callback_t file_func, struct output_file *out);
 void free_data_blocks();
 
 #endif
diff --git a/ext4_utils/make_ext4fs.c b/ext4_utils/make_ext4fs.c
index c96048b..5b0a3af 100644
--- a/ext4_utils/make_ext4fs.c
+++ b/ext4_utils/make_ext4fs.c
@@ -19,11 +19,10 @@
 #include <sys/types.h>
 #include <sys/ioctl.h>
 #include <sys/stat.h>
-#include <sys/mman.h>
+#include <sys/ioctl.h>
 #include <limits.h>
 #include <arpa/inet.h>
 #include <fcntl.h>
-#include <unistd.h>
 #include <stdlib.h>
 #include <strings.h>
 #include <string.h>
@@ -38,6 +37,7 @@
 #endif
 
 #include "make_ext4fs.h"
+#include "output_file.h"
 #include "ext4_utils.h"
 #include "allocate.h"
 #include "ext_utils.h"
@@ -57,8 +57,7 @@
 /* TODO: Not implemented:
    Allocating blocks in the same block group as the file inode
    Hash or binary tree directories
-   Non-extent inodes
-   Special files: symbolic links, sockets, devices, fifos
+   Special files: sockets, devices, fifos
  */
 
 int force = 0;
@@ -66,132 +65,27 @@
 struct fs_info info;
 struct fs_aux_info aux_info;
 
-/* Write a contiguous region of data blocks from a memory buffer */
-static void write_data_block(void *priv, u32 block, u8 *data, int len)
-{
-	int fd = *(int*)priv;
-	off_t off;
-	int ret;
-
-	if (block * info.block_size + len >= info.len) {
-		error("attempted to write block %llu past end of filesystem",
-				block * info.block_size + len - info.len);
-		return;
-	}
-
-	off = (off_t)block * info.block_size;
-	off = lseek(fd, off, SEEK_SET);
-	if (off < 0) {
-		error_errno("lseek");
-		return;
-	}
-
-	ret = write(fd, data, len);
-	if (ret < 0)
-		error_errno("write");
-	else if (ret < len)
-		error("incomplete write");
-}
-
-/* Write a contiguous region of data blocks from a file */
-static void write_data_file(void *priv, u32 block, const char *file,
-        off_t offset, int len)
-{
-	int fd = *(int*)priv;
-	off_t off;
-	int ret;
-
-	if (block * info.block_size + len >= info.len) {
-		error("attempted to write block %llu past end of filesystem",
-				block * info.block_size + len - info.len);
-		return;
-	}
-
-	int file_fd = open(file, O_RDONLY);
-	if (file_fd < 0) {
-		error_errno("open");
-		return;
-	}
-
-	u8 *data = mmap(NULL, len, PROT_READ, MAP_SHARED, file_fd, offset);
-	if (data == MAP_FAILED) {
-		error_errno("mmap");
-		close(fd);
-		return;
-	}
-
-	off = (off_t)block * info.block_size;
-	off = lseek(fd, off, SEEK_SET);
-	if (off < 0) {
-		error_errno("lseek");
-		return;
-	}
-
-	ret = write(fd, data, len);
-	if (ret < 0)
-		error_errno("write");
-	else if (ret < len)
-		error("incomplete write");
-
-	munmap(data, len);
-
-	close(file_fd);
-}
-
 /* Write the filesystem image to a file */
-static void write_ext4_image(const char *filename)
+static void write_ext4_image(const char *filename, int gz)
 {
 	int ret = 0;
-	int fd = open(filename, O_WRONLY | O_CREAT | O_TRUNC, 0644);
+	struct output_file *out = open_output_file(filename, gz);
 	off_t off;
 
-	if (fd < 0) {
-		error_errno("open");
+	if (!out)
 		return;
-	}
 
-	off = lseek(fd, 1024, SEEK_SET);
-	if (off < 0) {
-		error_errno("lseek");
-		return;
-	}
+	write_data_block(out, 1024, (u8*)aux_info.sb, 1024);
 
-	ret = write(fd, aux_info.sb, 1024);
-	if (ret < 0)
-		error_errno("write");
-	else if (ret < 1024)
-		error("incomplete write");
+	write_data_block(out, (aux_info.first_data_block + 1) * info.block_size,
+			 (u8*)aux_info.bg_desc, 
+			 aux_info.bg_desc_blocks * info.block_size);
 
-	off = (aux_info.first_data_block + 1) * info.block_size;
-	off = lseek(fd, off, SEEK_SET);
-	if (off < 0) {
-		error_errno("lseek");
-		return;
-	}
+	for_each_data_block(write_data_block, write_data_file, out);
 
-	ret = write(fd, aux_info.bg_desc,
-			aux_info.bg_desc_blocks * info.block_size);
-	if (ret < 0)
-		error_errno("write");
-	else if (ret < (int)(aux_info.bg_desc_blocks * info.block_size))
-		error("incomplete write");
+	write_data_block(out, info.len - 1, (u8*)"", 1);
 
-	for_each_data_block(write_data_block, write_data_file, &fd);
-
-	off = info.len - 1;
-	off = lseek(fd, off, SEEK_SET);
-	if (off < 0) {
-		error_errno("lseek");
-		return;
-	}
-
-	ret = write(fd, "\0", 1);
-	if (ret < 0)
-		error_errno("write");
-	else if (ret < 1)
-		error("incomplete write");
-
-	close(fd);
+	close_output_file(out);
 }
 
 /* Compute the rest of the parameters of the filesystem from the basic info */
@@ -664,10 +558,11 @@
 	const char *directory = NULL;
 	char *mountpoint = "";
 	int android = 0;
+	int gzip = 0;
 	u32 root_inode_num;
 	u16 root_mode;
 
-	while ((opt = getopt(argc, argv, "l:j:b:g:i:I:L:a:f")) != -1) {
+	while ((opt = getopt(argc, argv, "l:j:b:g:i:I:L:a:fz")) != -1) {
 		switch (opt) {
 		case 'l':
 			info.len = parse_num(optarg);
@@ -697,6 +592,9 @@
 			android = 1;
 			mountpoint = optarg;
 			break;
+		case 'z':
+			gzip = 1;
+			break;
 		default: /* '?' */
 			usage(argv[0]);
 			exit(EXIT_FAILURE);
@@ -804,7 +702,7 @@
 			aux_info.sb->s_blocks_count_lo - aux_info.sb->s_free_blocks_count_lo,
 			aux_info.sb->s_blocks_count_lo);
 
-	write_ext4_image(filename);
+	write_ext4_image(filename, gzip);
 
 	return 0;
 }
diff --git a/ext4_utils/output_file.c b/ext4_utils/output_file.c
new file mode 100644
index 0000000..fa6af5f
--- /dev/null
+++ b/ext4_utils/output_file.c
@@ -0,0 +1,213 @@
+/*
+ * Copyright (C) 2010 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#define _LARGEFILE64_SOURCE
+
+#include <sys/types.h>
+#include <sys/stat.h>
+#include <sys/types.h>
+#include <sys/mman.h>
+#include <unistd.h>
+#include <fcntl.h>
+
+#include <zlib.h>
+
+#include "ext4_utils.h"
+#include "output_file.h"
+
+struct output_file_ops {
+	int (*seek)(struct output_file *, off_t);
+	int (*write)(struct output_file *, u8 *, int);
+	void (*close)(struct output_file *);
+};
+
+struct output_file {
+	int fd;
+	gzFile gz_fd;
+	struct output_file_ops *ops;
+};
+
+static int file_seek(struct output_file *out, off_t off)
+{
+	off_t ret;
+
+	ret = lseek(out->fd, off, SEEK_SET);
+	if (ret < 0) {
+		error_errno("lseek");
+		return -1;
+	}
+	return 0;
+}
+
+static int file_write(struct output_file *out, u8 *data, int len)
+{
+	int ret;
+	ret = write(out->fd, data, len);
+	if (ret < 0) {
+		error_errno("write");
+		return -1;
+	} else if (ret < len) {
+		error("incomplete write");
+		return -1;
+	}
+
+	return 0;
+}
+
+static void file_close(struct output_file *out)
+{
+	close(out->fd);
+}
+
+
+static struct output_file_ops file_ops = {
+	.seek = file_seek,
+	.write = file_write,
+	.close = file_close,
+};
+
+static int gz_file_seek(struct output_file *out, off_t off)
+{
+	off_t ret;
+
+	ret = gzseek(out->gz_fd, off, SEEK_SET);
+	if (ret < 0) {
+		error_errno("gzseek");
+		return -1;
+	}
+	return 0;
+}
+
+static int gz_file_write(struct output_file *out, u8 *data, int len)
+{
+	int ret;
+	ret = gzwrite(out->gz_fd, data, len);
+	if (ret < 0) {
+		error_errno("gzwrite");
+		return -1;
+	} else if (ret < len) {
+		error("incomplete gzwrite");
+		return -1;
+	}
+
+	return 0;
+}
+
+static void gz_file_close(struct output_file *out)
+{
+	gzclose(out->gz_fd);
+}
+
+static struct output_file_ops gz_file_ops = {
+	.seek = gz_file_seek,
+	.write = gz_file_write,
+	.close = gz_file_close,
+};
+
+void close_output_file(struct output_file *out)
+{
+	out->ops->close(out);
+}
+
+struct output_file *open_output_file(const char *filename, int gz)
+{
+	struct output_file *out = malloc(sizeof(struct output_file));
+	if (!out) {
+		error_errno("malloc");
+		return NULL;
+	}
+
+	if (gz) {
+		out->ops = &gz_file_ops;
+		out->gz_fd = gzopen(filename, "wb9");
+		if (!out->gz_fd) {
+			error_errno("gzopen");
+			free(out);
+			return NULL;
+		}
+	} else {
+		out->ops = &file_ops;
+		out->fd = open(filename, O_WRONLY | O_CREAT | O_TRUNC, 0644);
+		if (out->fd < 0) {
+			error_errno("open");
+			free(out);
+			return NULL;
+		}
+	}
+	return out;
+}
+
+/* Write a contiguous region of data blocks from a memory buffer */
+void write_data_block(struct output_file *out, u64 off, u8 *data, int len)
+{
+	int ret;
+	
+	if (off + len > info.len) {
+		error("attempted to write block %llu past end of filesystem",
+				off + len - info.len);
+		return;
+	}
+
+	ret = out->ops->seek(out, off);
+	if (ret < 0)
+		return;
+
+	ret = out->ops->write(out, data, len);
+	if (ret < 0)
+		return;
+}
+
+/* Write a contiguous region of data blocks from a file */
+void write_data_file(struct output_file *out, u64 off, const char *file,
+		     off_t offset, int len)
+{
+	int ret;
+
+	if (off + len >= info.len) {
+		error("attempted to write block %llu past end of filesystem",
+				off + len - info.len);
+		return;
+	}
+
+	int file_fd = open(file, O_RDONLY);
+	if (file_fd < 0) {
+		error_errno("open");
+		return;
+	}
+
+	u8 *data = mmap(NULL, len, PROT_READ, MAP_SHARED, file_fd, offset);
+	if (data == MAP_FAILED) {
+		error_errno("mmap");
+		close(file_fd);
+		return;
+	}
+
+	ret = out->ops->seek(out, off);
+	if (ret < 0)
+		goto err;
+
+	ret = out->ops->write(out, data, len);
+	if (ret < 0)
+		goto err;
+
+
+	munmap(data, len);
+
+	close(file_fd);
+
+err:
+	munmap(data, len);
+	close(file_fd);
+}
diff --git a/ext4_utils/output_file.h b/ext4_utils/output_file.h
new file mode 100644
index 0000000..1df9e81
--- /dev/null
+++ b/ext4_utils/output_file.h
@@ -0,0 +1,23 @@
+/*
+ * Copyright (C) 2010 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+struct output_file;
+
+struct output_file *open_output_file(const char *filename, int gz);
+void write_data_block(struct output_file *out, u64 off, u8 *data, int len);
+void write_data_file(struct output_file *out, u64 off, const char *file,
+		     off_t offset, int len);
+void close_output_file(struct output_file *out);