Split data reader constructor and fragment table loading

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

5 files changed, 54 insertions, 40 deletions

diff --git a/difftool/sqfsdiff.c b/difftool/sqfsdiff.c
index 6d1a9af..22b8ec6 100644
--- a/difftool/sqfsdiff.c
+++ b/difftool/sqfsdiff.c
@@ -69,14 +69,19 @@ static int open_sfqs(sqfs_state_t *state, const char *path)
 		goto fail_dr;
 	}
 
-	state->data = data_reader_create(state->file, &state->super,
+	state->data = data_reader_create(state->file, state->super.block_size,
 					 state->cmp);
 	if (state->data == NULL) {
 		fprintf(stderr, "%s: error loading file system tree\n", path);
 		goto fail_tree;
 	}
 
+	if (data_reader_load_fragment_table(state->data, &state->super))
+		goto fail_data;
+
 	return 0;
+fail_data:
+	data_reader_destroy(state->data);
 fail_tree:
 	sqfs_dir_tree_destroy(state->root);
 fail_dr:
diff --git a/include/data_reader.h b/include/data_reader.h
index 324cdc8..a8ccabf 100644
--- a/include/data_reader.h
+++ b/include/data_reader.h
@@ -15,17 +15,13 @@
 
 typedef struct data_reader_t data_reader_t;
 
-/*
-  Create a data reader for accessing data blocks in a squashfs image.
-
-  Internally creates a fragment_reader_t (if applicable) to resolve
-  fragment indices.
 
-  Prints error messsages to stderr on failure.
- */
-data_reader_t *data_reader_create(sqfs_file_t *file, sqfs_super_t *super,
+data_reader_t *data_reader_create(sqfs_file_t *file, size_t block_size,
 				  sqfs_compressor_t *cmp);
 
+int data_reader_load_fragment_table(data_reader_t *data,
+				    const sqfs_super_t *super);
+
 void data_reader_destroy(data_reader_t *data);
 
 int data_reader_get_fragment(data_reader_t *data,
diff --git a/lib/sqfshelper/data_reader.c b/lib/sqfshelper/data_reader.c
index 9a31fa4..d054069 100644
--- a/lib/sqfshelper/data_reader.c
+++ b/lib/sqfshelper/data_reader.c
@@ -130,51 +130,58 @@ static int precache_fragment_block(data_reader_t *data, size_t idx)
 	return 0;
 }
 
-data_reader_t *data_reader_create(sqfs_file_t *file, sqfs_super_t *super,
+data_reader_t *data_reader_create(sqfs_file_t *file, size_t block_size,
 				  sqfs_compressor_t *cmp)
 {
-	data_reader_t *data = alloc_flex(sizeof(*data), 1, super->block_size);
-	size_t i, size;
+	data_reader_t *data = alloc_flex(sizeof(*data), 1, block_size);
+
+	if (data != NULL) {
+		data->file = file;
+		data->block_size = block_size;
+		data->cmp = cmp;
+	}
+
+	return data;
+}
+
+int data_reader_load_fragment_table(data_reader_t *data,
+				    const sqfs_super_t *super)
+{
 	void *raw_frag;
+	size_t size;
+	uint32_t i;
 	int ret;
 
-	if (data == NULL) {
-		perror("creating data reader");
-		return data;
-	}
+	free(data->frag_block);
+	free(data->frag);
 
-	data->num_fragments = super->fragment_entry_count;
-	data->current_frag_index = super->fragment_entry_count;
-	data->file = file;
-	data->block_size = super->block_size;
-	data->cmp = cmp;
+	data->frag = NULL;
+	data->frag_block = NULL;
+	data->num_fragments = 0;
+	data->current_frag_index = 0;
 
 	if (super->fragment_entry_count == 0 ||
 	    (super->flags & SQFS_FLAG_NO_FRAGMENTS) != 0) {
-		return data;
+		return 0;
 	}
 
-	if (super->fragment_table_start >= super->bytes_used) {
-		fputs("Fragment table start is past end of file\n", stderr);
-		free(data);
-		return NULL;
-	}
+	if (super->fragment_table_start >= super->bytes_used)
+		return SQFS_ERROR_OUT_OF_BOUNDS;
 
-	if (SZ_MUL_OV(sizeof(data->frag[0]), data->num_fragments, &size)) {
-		fputs("Too many fragments: overflow\n", stderr);
-		free(data);
-		return NULL;
+	if (SZ_MUL_OV(sizeof(data->frag[0]), super->fragment_entry_count,
+		      &size)) {
+		return SQFS_ERROR_OVERFLOW;
 	}
 
-	ret = sqfs_read_table(file, cmp, size, super->fragment_table_start,
+	ret = sqfs_read_table(data->file, data->cmp, size,
+			      super->fragment_table_start,
 			      super->directory_table_start,
 			      super->fragment_table_start, &raw_frag);
+	if (ret)
+		return ret;
 
-	if (ret) {
-		free(data);
-		return NULL;
-	}
-
+	data->num_fragments = super->fragment_entry_count;
+	data->current_frag_index = super->fragment_entry_count;
 	data->frag = raw_frag;
 
 	for (i = 0; i < data->num_fragments; ++i) {
@@ -183,7 +190,7 @@ data_reader_t *data_reader_create(sqfs_file_t *file, sqfs_super_t *super,
 			le64toh(data->frag[i].start_offset);
 	}
 
-	return data;
+	return 0;
 }
 
 void data_reader_destroy(data_reader_t *data)
diff --git a/tar/sqfs2tar.c b/tar/sqfs2tar.c
index c624934..7f388fe 100644
--- a/tar/sqfs2tar.c
+++ b/tar/sqfs2tar.c
@@ -434,8 +434,11 @@ int main(int argc, char **argv)
 		goto out_id;
 	}
 
-	data = data_reader_create(file, &super, cmp);
+	data = data_reader_create(file, super.block_size, cmp);
 	if (data == NULL)
+		goto out_id;
+
+	if (data_reader_load_fragment_table(data, &super))
 		goto out_data;
 
 	dr = sqfs_dir_reader_create(&super, cmp, file);
diff --git a/unpack/rdsquashfs.c b/unpack/rdsquashfs.c
index e8ec3ec..29f9b79 100644
--- a/unpack/rdsquashfs.c
+++ b/unpack/rdsquashfs.c
@@ -80,10 +80,13 @@ int main(int argc, char **argv)
 		goto out_id;
 	}
 
-	data = data_reader_create(file, &super, cmp);
+	data = data_reader_create(file, super.block_size, cmp);
 	if (data == NULL)
 		goto out_dr;
 
+	if (data_reader_load_fragment_table(data, &super))
+		goto out_data;
+
 	ret = sqfs_dir_reader_get_full_hierarchy(dirrd, idtbl, opt.cmdpath,
 						 opt.rdtree_flags, &n);
 	if (ret) {