Add the ability to hook into the data writer block writing

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

4 files changed, 53 insertions, 17 deletions

diff --git a/lib/sqfs/data_writer/block.c b/lib/sqfs/data_writer/block.c
index 9461737..9c40793 100644
--- a/lib/sqfs/data_writer/block.c
+++ b/lib/sqfs/data_writer/block.c
@@ -61,11 +61,15 @@ static int allign_file(sqfs_data_writer_t *proc, sqfs_block_t *blk)
 
 int process_completed_block(sqfs_data_writer_t *proc, sqfs_block_t *blk)
 {
+	uint64_t offset, bytes;
 	size_t start, count;
-	uint64_t offset;
 	uint32_t out;
 	int err;
 
+	if (proc->hooks != NULL && proc->hooks->pre_block_write != NULL) {
+		proc->hooks->pre_block_write(proc->user_ptr, blk, proc->file);
+	}
+
 	if (blk->flags & SQFS_BLK_FIRST_BLOCK) {
 		proc->start = proc->file->get_size(proc->file);
 		proc->file_start = proc->num_blocks;
@@ -101,6 +105,10 @@ int process_completed_block(sqfs_data_writer_t *proc, sqfs_block_t *blk)
 			return err;
 	}
 
+	if (proc->hooks != NULL && proc->hooks->post_block_write != NULL) {
+		proc->hooks->post_block_write(proc->user_ptr, blk, proc->file);
+	}
+
 	if (blk->flags & SQFS_BLK_LAST_BLOCK) {
 		err = allign_file(proc, blk);
 		if (err)
@@ -112,19 +120,28 @@ int process_completed_block(sqfs_data_writer_t *proc, sqfs_block_t *blk)
 
 		sqfs_inode_set_file_block_start(blk->inode, offset);
 
-		if (start < proc->file_start) {
-			offset = start + count;
+		if (start >= proc->file_start)
+			return 0;
+
+		offset = start + count;
+		if (offset >= proc->file_start) {
+			count = proc->num_blocks - offset;
+			proc->num_blocks = offset;
+		} else {
+			proc->num_blocks = proc->file_start;
+		}
 
-			if (offset >= proc->file_start) {
-				proc->num_blocks = offset;
-			} else {
-				proc->num_blocks = proc->file_start;
-			}
+		if (proc->hooks != NULL &&
+		    proc->hooks->notify_blocks_erased != NULL) {
+			bytes = proc->file->get_size(proc->file) - proc->start;
 
-			err = proc->file->truncate(proc->file, proc->start);
-			if (err)
-				return err;
+			proc->hooks->notify_blocks_erased(proc->user_ptr,
+							  count, bytes);
 		}
+
+		err = proc->file->truncate(proc->file, proc->start);
+		if (err)
+			return err;
 	}
 
 	return 0;
diff --git a/lib/sqfs/data_writer/common.c b/lib/sqfs/data_writer/common.c
index 51acc1e..c9d5589 100644
--- a/lib/sqfs/data_writer/common.c
+++ b/lib/sqfs/data_writer/common.c
@@ -155,3 +155,10 @@ int sqfs_data_writer_write_fragment_table(sqfs_data_writer_t *proc,
 	super->fragment_table_start = start;
 	return 0;
 }
+
+void sqfs_data_writer_set_hooks(sqfs_data_writer_t *proc, void *user_ptr,
+				const sqfs_block_hooks_t *hooks)
+{
+	proc->hooks = hooks;
+	proc->user_ptr = user_ptr;
+}
diff --git a/lib/sqfs/data_writer/fragment.c b/lib/sqfs/data_writer/fragment.c
index e4fe9b4..9f169b7 100644
--- a/lib/sqfs/data_writer/fragment.c
+++ b/lib/sqfs/data_writer/fragment.c
@@ -64,6 +64,10 @@ static int store_fragment(sqfs_data_writer_t *proc, sqfs_block_t *frag,
 	sqfs_inode_set_frag_location(frag->inode, proc->frag_block->index,
 				     proc->frag_block->size);
 
+	if (proc->hooks != NULL && proc->hooks->pre_fragment_store != NULL) {
+		proc->hooks->pre_fragment_store(proc->user_ptr, frag);
+	}
+
 	memcpy(proc->frag_block->data + proc->frag_block->size,
 	       frag->data, frag->size);
 
@@ -82,12 +86,8 @@ int process_completed_fragment(sqfs_data_writer_t *proc, sqfs_block_t *frag,
 	hash = MK_BLK_HASH(frag->checksum, frag->size);
 
 	for (i = 0; i < proc->frag_list_num; ++i) {
-		if (proc->frag_list[i].hash == hash) {
-			sqfs_inode_set_frag_location(frag->inode,
-						     proc->frag_list[i].index,
-						     proc->frag_list[i].offset);
-			return 0;
-		}
+		if (proc->frag_list[i].hash == hash)
+			goto out_duplicate;
 	}
 
 	if (proc->frag_block != NULL) {
@@ -125,4 +125,13 @@ fail:
 	free(*blk_out);
 	*blk_out = NULL;
 	return err;
+out_duplicate:
+	sqfs_inode_set_frag_location(frag->inode, proc->frag_list[i].index,
+				     proc->frag_list[i].offset);
+
+	if (proc->hooks != NULL &&
+	    proc->hooks->notify_fragment_discard != NULL) {
+		proc->hooks->notify_fragment_discard(proc->user_ptr, frag);
+	}
+	return 0;
 }
diff --git a/lib/sqfs/data_writer/internal.h b/lib/sqfs/data_writer/internal.h
index 69de9ac..53ddc5f 100644
--- a/lib/sqfs/data_writer/internal.h
+++ b/lib/sqfs/data_writer/internal.h
@@ -98,6 +98,9 @@ struct sqfs_data_writer_t {
 	size_t frag_list_num;
 	size_t frag_list_max;
 
+	const sqfs_block_hooks_t *hooks;
+	void *user_ptr;
+
 	/* used only by workers */
 	size_t max_block_size;