4 files changed, 55 insertions, 41 deletions

diff --git a/lib/sqfs/data_writer.c b/lib/sqfs/data_writer.c
index 74701fd..ae8257c 100644
--- a/lib/sqfs/data_writer.c
+++ b/lib/sqfs/data_writer.c
@@ -295,6 +295,8 @@ void data_writer_destroy(data_writer_t *data)
 int data_writer_write_fragment_table(data_writer_t *data)
 {
 	uint64_t start;
+	size_t size;
+	int ret;
 
 	if (data->num_fragments == 0) {
 		data->super->fragment_entry_count = 0;
@@ -302,14 +304,13 @@ int data_writer_write_fragment_table(data_writer_t *data)
 		return 0;
 	}
 
-	data->super->fragment_entry_count = data->num_fragments;
-
-	if (sqfs_write_table(data->outfd, data->super, data->fragments,
-			     sizeof(data->fragments[0]), data->num_fragments,
-			     &start, data->cmp)) {
+	size = sizeof(data->fragments[0]) * data->num_fragments;
+	ret = sqfs_write_table(data->outfd, data->super, data->cmp,
+			       data->fragments, size, &start);
+	if (ret)
 		return -1;
-	}
 
+	data->super->fragment_entry_count = data->num_fragments;
 	data->super->fragment_table_start = start;
 	return 0;
 }
diff --git a/lib/sqfs/id_table_write.c b/lib/sqfs/id_table_write.c
index c2ed353..3bbe302 100644
--- a/lib/sqfs/id_table_write.c
+++ b/lib/sqfs/id_table_write.c
@@ -15,8 +15,8 @@ int id_table_write(id_table_t *tbl, int outfd, sqfs_super_t *super,
 
 	super->id_count = tbl->num_ids;
 
-	ret = sqfs_write_table(outfd, super, tbl->ids, sizeof(tbl->ids[0]),
-			       tbl->num_ids, &start, cmp);
+	ret = sqfs_write_table(outfd, super, cmp, tbl->ids,
+			       sizeof(tbl->ids[0]) * tbl->num_ids, &start);
 
 	super->id_table_start = start;
 
diff --git a/lib/sqfs/write_export_table.c b/lib/sqfs/write_export_table.c
index a0d81eb..8794609 100644
--- a/lib/sqfs/write_export_table.c
+++ b/lib/sqfs/write_export_table.c
@@ -10,10 +10,11 @@ int write_export_table(int outfd, fstree_t *fs, sqfs_super_t *super,
 		       compressor_t *cmp)
 {
 	uint64_t *table, start;
-	size_t i;
+	size_t i, size;
 	int ret;
 
-	table = malloc(sizeof(uint64_t) * (fs->inode_tbl_size - 1));
+	size = sizeof(uint64_t) * (fs->inode_tbl_size - 1);
+	table = malloc(size);
 
 	if (table == NULL) {
 		perror("Allocating NFS export table");
@@ -28,8 +29,7 @@ int write_export_table(int outfd, fstree_t *fs, sqfs_super_t *super,
 		}
 	}
 
-	ret = sqfs_write_table(outfd, super, table, sizeof(table[0]),
-			       fs->inode_tbl_size - 1, &start, cmp);
+	ret = sqfs_write_table(outfd, super, cmp, table, size, &start);
 
 	super->export_table_start = start;
 	super->flags |= SQFS_FLAG_EXPORTABLE;
diff --git a/lib/sqfs/write_table.c b/lib/sqfs/write_table.c
index fc90a3f..23ba5a8 100644
--- a/lib/sqfs/write_table.c
+++ b/lib/sqfs/write_table.c
@@ -6,56 +6,69 @@
 #include "util.h"
 
 #include <endian.h>
+#include <stdlib.h>
 #include <stdio.h>
 
-int sqfs_write_table(int outfd, sqfs_super_t *super, const void *data,
-		     size_t entsize, size_t count, uint64_t *startblock,
-		     compressor_t *cmp)
+int sqfs_write_table(int outfd, sqfs_super_t *super, compressor_t *cmp,
+		     const void *data, size_t table_size, uint64_t *start)
 {
-	size_t ent_per_blocks = SQFS_META_BLOCK_SIZE / entsize;
-	uint64_t blocks[count / ent_per_blocks + 1], block;
-	size_t i, blkidx = 0, tblsize;
+	size_t block_count, list_size, diff, blkidx = 0;
+	uint64_t block, *locations;
 	meta_writer_t *m;
 	uint32_t offset;
+	int ret = -1;
 
-	/* Write actual data. Whenever we cross a block boundary, remember
-	   the block start offset */
+	block_count = table_size / SQFS_META_BLOCK_SIZE;
+	if ((table_size % SQFS_META_BLOCK_SIZE) != 0)
+		++block_count;
+
+	list_size = sizeof(uint64_t) * block_count;
+	locations = malloc(list_size);
+
+	if (locations == NULL) {
+		perror("writing table");
+		return -1;
+	}
+
+	/* Write actual data */
 	m = meta_writer_create(outfd, cmp, false);
 	if (m == NULL)
-		return -1;
+		goto out_idx;
 
-	for (i = 0; i < count; ++i) {
+	while (table_size > 0) {
 		meta_writer_get_position(m, &block, &offset);
+		locations[blkidx++] = htole64(super->bytes_used + block);
 
-		if (blkidx == 0 || block > blocks[blkidx - 1])
-			blocks[blkidx++] = block;
+		diff = SQFS_META_BLOCK_SIZE;
+		if (diff > table_size)
+			diff = table_size;
 
-		if (meta_writer_append(m, data, entsize))
-			goto fail;
+		if (meta_writer_append(m, data, diff))
+			goto out;
 
-		data = (const char *)data + entsize;
+		data = (const char *)data + diff;
+		table_size -= diff;
 	}
 
 	if (meta_writer_flush(m))
-		goto fail;
-
-	for (i = 0; i < blkidx; ++i)
-		blocks[i] = htole64(blocks[i] + super->bytes_used);
+		goto out;
 
 	meta_writer_get_position(m, &block, &offset);
 	super->bytes_used += block;
-	meta_writer_destroy(m);
 
-	/* write new index table */
-	*startblock = super->bytes_used;
-	tblsize = sizeof(blocks[0]) * blkidx;
+	/* write location list */
+	*start = super->bytes_used;
 
-	if (write_data("writing table index", outfd, blocks, tblsize))
-		return -1;
+	if (write_data("writing table locations", outfd, locations, list_size))
+		goto out;
+
+	super->bytes_used += list_size;
 
-	super->bytes_used += tblsize;
-	return 0;
-fail:
+	/* cleanup */
+	ret = 0;
+out:
 	meta_writer_destroy(m);
-	return -1;
+out_idx:
+	free(locations);
+	return ret;
 }