Add the ability to hook into the data writer block writing

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

6 files changed, 152 insertions, 17 deletions

diff --git a/include/sqfs/data_writer.h b/include/sqfs/data_writer.h
index 3f1ad02..cc99700 100644
--- a/include/sqfs/data_writer.h
+++ b/include/sqfs/data_writer.h
@@ -42,6 +42,91 @@
  * is called for each one.
  */
 
+/**
+ * @struct sqfs_block_hooks_t
+ *
+ * @brief A set of hooks for tapping into the data writer.
+ *
+ * This structure can be registered with an @ref sqfs_data_writer_t and
+ * contains function pointers that will be called during various stages
+ * when writing data to disk.
+ *
+ * The callbacks can not only be used for accounting but may also write extra
+ * data to the output file or make modifications to the blocks before they are
+ * writtien.
+ *
+ * The callbacks can be individually set to NULL to disable them.
+ */
+struct sqfs_block_hooks_t {
+	/**
+	 * @brief Gets called before writing a block to disk.
+	 *
+	 * If this is not NULL, it gets called before a block is written to
+	 * disk. If the block has the @ref SQFS_BLK_ALLIGN flag set, the
+	 * function is called before padding the file.
+	 *
+	 * The function may modify the block itself or write data to the file.
+	 * which is taken into account when padding the file.
+	 *
+	 * @param user A user pointer.
+	 * @param block The block that is about to be written.
+	 * @param file The file that the block will be written to.
+	 */
+	void (*pre_block_write)(void *user, sqfs_block_t *block,
+				sqfs_file_t *file);
+
+	/**
+	 * @brief Gets called after writing a block to disk.
+	 *
+	 * If this is not NULL, it gets called after a block is written to
+	 * disk. If the block has the @ref SQFS_BLK_ALLIGN flag set, the
+	 * function is called before padding the file.
+	 *
+	 * Modifying the block is rather pointless, but the function may
+	 * write data to the file which is taken into account when padding
+	 * the file.
+	 *
+	 * @param user A user pointer.
+	 * @param block The block that is about to be written.
+	 * @param file The file that the block was written to.
+	 */
+	void (*post_block_write)(void *user, const sqfs_block_t *block,
+				 sqfs_file_t *file);
+
+	/**
+	 * @brief Gets called before storing a fragment in a fragment block.
+	 *
+	 * The function can modify the block before it is stored.
+	 *
+	 * @param user A user pointer.
+	 * @param block The data chunk that is about to be merged into the
+	 *              fragment block.
+	 */
+	void (*pre_fragment_store)(void *user, sqfs_block_t *block);
+
+	/**
+	 * @brief Gets called if block deduplication managed to get
+	 *        rid of the data blocks of a file.
+	 *
+	 * @param user A user pointer.
+	 * @param count The number of blocks that have been erased.
+	 * @param bytes The number of bytes that have been erased. Includes
+	 *              potential padding before and after the end.
+	 */
+	void (*notify_blocks_erased)(void *user, size_t count,
+				     uint64_t bytes);
+
+	/**
+	 * @brief Gets called before throwing away a fragment that turned out
+	 *        to be a duplicate.
+	 *
+	 * @param user A user pointer.
+	 * @param block The data chunk that is about to be merged into the
+	 *              fragment block.
+	 */
+	void (*notify_fragment_discard)(void *user, const sqfs_block_t *block);
+};
+
 #ifdef __cplusplus
 extern "C" {
 #endif
@@ -134,6 +219,19 @@ SQFS_API
 int sqfs_data_writer_write_fragment_table(sqfs_data_writer_t *proc,
 					  sqfs_super_t *super);
 
+/**
+ * @brief Register a set of hooks to be invoked when writing blocks to disk.
+ *
+ * @memberof sqfs_data_writer_t
+ *
+ * @param proc A pointer to a data writer object.
+ * @param user_ptr A user pointer to pass to the callbacks.
+ * @param hooks A structure containing the hooks.
+ */
+SQFS_API
+void sqfs_data_writer_set_hooks(sqfs_data_writer_t *proc, void *user_ptr,
+				const sqfs_block_hooks_t *hooks);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/include/sqfs/predef.h b/include/sqfs/predef.h
index c94ca29..8a554dd 100644
--- a/include/sqfs/predef.h
+++ b/include/sqfs/predef.h
@@ -72,6 +72,7 @@ typedef struct sqfs_file_t sqfs_file_t;
 typedef struct sqfs_sparse_map_t sqfs_sparse_map_t;
 typedef struct sqfs_tree_node_t sqfs_tree_node_t;
 typedef struct sqfs_data_reader_t sqfs_data_reader_t;
+typedef struct sqfs_block_hooks_t sqfs_block_hooks_t;
 
 typedef struct sqfs_fragment_t sqfs_fragment_t;
 typedef struct sqfs_dir_header_t sqfs_dir_header_t;
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;