Pack file data and fragments

Signed-off-by: David Oberhollenzer <david.oberhollenzer@sigma-star.at>

4 files changed, 233 insertions, 3 deletions

diff --git a/mkfs/Makemodule.am b/mkfs/Makemodule.am
index 82ac765..a9ee7fa 100644
--- a/mkfs/Makemodule.am
+++ b/mkfs/Makemodule.am
@@ -1,4 +1,4 @@
-mksquashfs_SOURCES = mkfs/mksquashfs.c mkfs/mksquashfs.h
+mksquashfs_SOURCES = mkfs/mksquashfs.c mkfs/mksquashfs.h mkfs/block.c
 mksquashfs_SOURCES += mkfs/options.c mkfs/meta_writer.c mkfs/super.c
 mksquashfs_SOURCES += include/squashfs.h
 mksquashfs_LDADD = libfstree.a libcompress.a libutil.a
diff --git a/mkfs/block.c b/mkfs/block.c
new file mode 100644
index 0000000..d68c0c7
--- /dev/null
+++ b/mkfs/block.c
@@ -0,0 +1,214 @@
+/* SPDX-License-Identifier: GPL-3.0-or-later */
+#include "mksquashfs.h"
+#include "util.h"
+
+static int write_block(file_info_t *fi, sqfs_info_t *info)
+{
+	size_t idx, bs;
+	ssize_t ret;
+
+	idx = info->file_block_count++;
+	bs = info->super.block_size;
+
+	ret = info->cmp->do_block(info->cmp, info->block, bs);
+	if (ret < 0)
+		return -1;
+
+	if (ret > 0) {
+		bs = ret;
+		fi->blocksizes[idx] = bs;
+	} else {
+		fi->blocksizes[idx] = bs | (1 << 24);
+	}
+
+	ret = write_retry(info->outfd, info->block, bs);
+	if (ret < 0) {
+		perror("writing to output file");
+		return -1;
+	}
+
+	if ((size_t)ret < bs) {
+		fputs("write to output file truncated\n", stderr);
+		return -1;
+	}
+
+	info->super.bytes_used += bs;
+	return 0;
+}
+
+static int flush_fragments(sqfs_info_t *info)
+{
+	size_t newsz, size;
+	file_info_t *fi;
+	uint64_t offset;
+	ssize_t ret;
+	void *new;
+
+	if (info->num_fragments == info->max_fragments) {
+		newsz = info->max_fragments ? info->max_fragments * 2 : 16;
+		new = realloc(info->fragments,
+			      sizeof(info->fragments[0]) * newsz);
+
+		if (new == NULL) {
+			perror("appending to fragment table");
+			return -1;
+		}
+
+		info->max_fragments = newsz;
+		info->fragments = new;
+	}
+
+	offset = info->super.bytes_used;
+	size = info->frag_offset;
+
+	for (fi = info->frag_list; fi != NULL; fi = fi->frag_next)
+		fi->fragment = info->num_fragments;
+
+	ret = info->cmp->do_block(info->cmp, info->fragment, size);
+	if (ret < 0)
+		return -1;
+
+	info->fragments[info->num_fragments].start_offset = htole64(offset);
+	info->fragments[info->num_fragments].pad0 = 0;
+
+	if (ret > 0) {
+		size = ret;
+		info->fragments[info->num_fragments].size = htole32(size);
+	} else {
+		info->fragments[info->num_fragments].size =
+			htole32(size | (1 << 24));
+	}
+
+	info->num_fragments += 1;
+
+	ret = write_retry(info->outfd, info->fragment, size);
+	if (ret < 0) {
+		perror("writing to output file");
+		return -1;
+	}
+
+	if ((size_t)ret < size) {
+		fputs("write to output file truncated\n", stderr);
+		return -1;
+	}
+
+	memset(info->fragment, 0, info->super.block_size);
+
+	info->super.bytes_used += size;
+	info->frag_offset = 0;
+	info->frag_list = NULL;
+
+	info->super.flags &= ~SQFS_FLAG_NO_FRAGMENTS;
+	info->super.flags |= SQFS_FLAG_ALWAYS_FRAGMENTS;
+	return 0;
+}
+
+static int add_fragment(file_info_t *fi, sqfs_info_t *info, size_t size)
+{
+	if (info->frag_offset + size > info->super.block_size) {
+		if (flush_fragments(info))
+			return -1;
+	}
+
+	fi->fragment_offset = info->frag_offset;
+	fi->frag_next = info->frag_list;
+	info->frag_list = fi;
+
+	memcpy((char *)info->fragment + info->frag_offset, info->block, size);
+	info->frag_offset += size;
+	return 0;
+}
+
+static int process_file(sqfs_info_t *info, file_info_t *fi)
+{
+	uint64_t count = fi->size;
+	int infd, ret;
+	size_t diff;
+
+	infd = open(fi->input_file, O_RDONLY);
+	if (infd < 0) {
+		perror(fi->input_file);
+		return -1;
+	}
+
+	fi->startblock = info->super.bytes_used;
+	info->file_block_count = 0;
+
+	while (count != 0) {
+		diff = count > (uint64_t)info->super.block_size ?
+			info->super.block_size : count;
+
+		ret = read_retry(infd, info->block, diff);
+		if (ret < 0)
+			goto fail_read;
+		if ((size_t)ret < diff)
+			goto fail_trunc;
+
+		if (diff < info->super.block_size) {
+			if (add_fragment(fi, info, diff))
+				goto fail;
+		} else {
+			if (write_block(fi, info))
+				goto fail;
+		}
+
+		count -= diff;
+	}
+
+	close(infd);
+	return 0;
+fail:
+	close(infd);
+	return -1;
+fail_read:
+	fprintf(stderr, "read from %s: %s\n", fi->input_file, strerror(errno));
+	goto fail;
+fail_trunc:
+	fprintf(stderr, "%s: truncated read\n", fi->input_file);
+	goto fail;
+}
+
+static int find_and_process_files(sqfs_info_t *info, tree_node_t *n)
+{
+	if (S_ISDIR(n->mode)) {
+		for (n = n->data.dir->children; n != NULL; n = n->next) {
+			if (find_and_process_files(info, n))
+				return -1;
+		}
+		return 0;
+	}
+
+	if (S_ISREG(n->mode))
+		return process_file(info, n->data.file);
+
+	return 0;
+}
+
+int write_data_to_image(sqfs_info_t *info)
+{
+	int ret;
+
+	info->block = malloc(info->super.block_size);
+
+	if (info->block == NULL) {
+		perror("allocating data block buffer");
+		return -1;
+	}
+
+	info->fragment = malloc(info->super.block_size);
+
+	if (info->fragment == NULL) {
+		perror("allocating fragment buffer");
+		free(info->block);
+		return -1;
+	}
+
+	ret = find_and_process_files(info, info->fs.root);
+
+	free(info->block);
+	free(info->fragment);
+
+	info->block = NULL;
+	info->fragment = NULL;
+	return ret;
+}
diff --git a/mkfs/mksquashfs.c b/mkfs/mksquashfs.c
index 65cd552..418d669 100644
--- a/mkfs/mksquashfs.c
+++ b/mkfs/mksquashfs.c
@@ -68,13 +68,25 @@ int main(int argc, char **argv)
 
 	print_tree(0, info.fs.root);
 
-	if (sqfs_super_write(&info))
+	info.cmp = compressor_create(info.super.compression_id, true,
+				     info.super.block_size);
+	if (info.cmp == NULL) {
+		fputs("Error creating compressor\n", stderr);
 		goto out_outfd;
+	}
+
+	if (write_data_to_image(&info))
+		goto out_cmp;
+
+	if (sqfs_super_write(&info))
+		goto out_cmp;
 
 	if (sqfs_padd_file(&info))
-		goto out_fstree;
+		goto out_cmp;
 
 	status = EXIT_SUCCESS;
+out_cmp:
+	info.cmp->destroy(info.cmp);
 out_fstree:
 	fstree_cleanup(&info.fs);
 out_outfd:
diff --git a/mkfs/mksquashfs.h b/mkfs/mksquashfs.h
index e216329..da9c937 100644
--- a/mkfs/mksquashfs.h
+++ b/mkfs/mksquashfs.h
@@ -50,6 +50,8 @@ typedef struct {
 	int file_block_count;
 	file_info_t *frag_list;
 	size_t frag_offset;
+
+	compressor_t *cmp;
 } sqfs_info_t;
 
 void process_command_line(options_t *opt, int argc, char **argv);
@@ -68,4 +70,6 @@ int meta_writer_flush(meta_writer_t *m);
 
 int meta_writer_append(meta_writer_t *m, const void *data, size_t size);
 
+int write_data_to_image(sqfs_info_t *info);
+
 #endif /* MKSQUASHFS_H */