Integrate the data reader into libsquashfs

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

2 files changed, 374 insertions, 2 deletions

diff --git a/lib/sqfs/Makemodule.am b/lib/sqfs/Makemodule.am
index 66d49b6..cb003c0 100644
--- a/lib/sqfs/Makemodule.am
+++ b/lib/sqfs/Makemodule.am
@@ -5,7 +5,8 @@ LIBSQFS_HEARDS = include/sqfs/data.h include/sqfs/meta_writer.h \
 		include/sqfs/dir.h include/sqfs/xattr.h \
 		include/sqfs/table.h include/sqfs/predef.h \
 		include/sqfs/error.h include/sqfs/dir_reader.h \
-		include/sqfs/dir_writer.h include/sqfs/io.h
+		include/sqfs/dir_writer.h include/sqfs/io.h \
+		include/sqfs/data_reader.h
 
 libsquashfs_la_SOURCES = $(LIBSQFS_HEARDS) lib/sqfs/id_table.c lib/sqfs/super.c
 libsquashfs_la_SOURCES += lib/sqfs/readdir.c lib/sqfs/io_file.c lib/sqfs/xattr.c
@@ -17,7 +18,7 @@ libsquashfs_la_SOURCES += lib/sqfs/read_table.c lib/sqfs/comp/compressor.c
 libsquashfs_la_SOURCES += lib/sqfs/io_stdin.c lib/sqfs/comp/internal.h
 libsquashfs_la_SOURCES += lib/sqfs/dir_reader.c lib/sqfs/read_tree.c
 libsquashfs_la_SOURCES += lib/sqfs/blk_proc/process_block.c lib/sqfs/io.c
-libsquashfs_la_SOURCES += lib/sqfs/blk_proc/internal.h
+libsquashfs_la_SOURCES += lib/sqfs/blk_proc/internal.h lib/sqfs/data_reader.c
 libsquashfs_la_CPPFLAGS = $(AM_CPPFLAGS)
 libsquashfs_la_CFLAGS = $(AM_CFLAGS) $(PTHREAD_CFLAGS) $(ZLIB_CFLAGS)
 libsquashfs_la_CFLAGS += $(XZ_CFLAGS) $(LZO_CFLAGS) $(LZ4_CFLAGS)
diff --git a/lib/sqfs/data_reader.c b/lib/sqfs/data_reader.c
new file mode 100644
index 0000000..3de59ff
--- /dev/null
+++ b/lib/sqfs/data_reader.c
@@ -0,0 +1,371 @@
+/* SPDX-License-Identifier: LGPL-3.0-or-later */
+/*
+ * data_reader.c
+ *
+ * Copyright (C) 2019 David Oberhollenzer <goliath@infraroot.at>
+ */
+#define SQFS_BUILDING_DLL
+#include "config.h"
+
+#include "sqfs/block_processor.h"
+#include "sqfs/data_reader.h"
+#include "sqfs/compress.h"
+#include "sqfs/error.h"
+#include "sqfs/table.h"
+#include "sqfs/inode.h"
+#include "sqfs/data.h"
+#include "sqfs/io.h"
+#include "util.h"
+
+#include <stdlib.h>
+#include <string.h>
+
+struct sqfs_data_reader_t {
+	sqfs_fragment_t *frag;
+	sqfs_compressor_t *cmp;
+	sqfs_block_t *data_block;
+	sqfs_block_t *frag_block;
+
+	uint64_t current_block;
+
+	sqfs_file_t *file;
+
+	uint32_t num_fragments;
+	uint32_t current_frag_index;
+	uint32_t block_size;
+
+	uint8_t scratch[];
+};
+
+static int get_block(sqfs_data_reader_t *data, uint64_t off, uint32_t size,
+		     size_t unpacked_size, sqfs_block_t **out)
+{
+	sqfs_block_t *blk = alloc_flex(sizeof(*blk), 1, unpacked_size);
+	size_t on_disk_size;
+	ssize_t ret;
+	int err;
+
+	if (blk == NULL)
+		return SQFS_ERROR_ALLOC;
+
+	blk->size = unpacked_size;
+
+	if (SQFS_IS_SPARSE_BLOCK(size)) {
+		*out = blk;
+		return 0;
+	}
+
+	on_disk_size = SQFS_ON_DISK_BLOCK_SIZE(size);
+
+	if (on_disk_size > unpacked_size)
+		return SQFS_ERROR_OVERFLOW;
+
+	if (SQFS_IS_BLOCK_COMPRESSED(size)) {
+		err = data->file->read_at(data->file, off,
+					  data->scratch, on_disk_size);
+		if (err) {
+			free(blk);
+			return err;
+		}
+
+		ret = data->cmp->do_block(data->cmp, data->scratch,
+					  on_disk_size, blk->data, blk->size);
+		if (ret <= 0)
+			err = ret < 0 ? ret : SQFS_ERROR_OVERFLOW;
+	} else {
+		err = data->file->read_at(data->file, off,
+					  blk->data, on_disk_size);
+	}
+
+	if (err) {
+		free(blk);
+		return err;
+	}
+
+	*out = blk;
+	return 0;
+}
+
+static int precache_data_block(sqfs_data_reader_t *data, uint64_t location,
+			       uint32_t size)
+{
+	int ret;
+
+	if (data->data_block != NULL && data->current_block == location)
+		return 0;
+
+	free(data->data_block);
+
+	ret = get_block(data, location, size, data->block_size,
+			&data->data_block);
+
+	if (ret < 0) {
+		data->data_block = NULL;
+		return -1;
+	}
+
+	data->current_block = location;
+	return 0;
+}
+
+static int precache_fragment_block(sqfs_data_reader_t *data, size_t idx)
+{
+	int ret;
+
+	if (data->frag_block != NULL && idx == data->current_frag_index)
+		return 0;
+
+	if (idx >= data->num_fragments)
+		return SQFS_ERROR_OUT_OF_BOUNDS;
+
+	free(data->frag_block);
+
+	ret = get_block(data, data->frag[idx].start_offset,
+			data->frag[idx].size, data->block_size,
+			&data->frag_block);
+	if (ret < 0)
+		return -1;
+
+	data->current_frag_index = idx;
+	return 0;
+}
+
+sqfs_data_reader_t *sqfs_data_reader_create(sqfs_file_t *file,
+					    size_t block_size,
+					    sqfs_compressor_t *cmp)
+{
+	sqfs_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 sqfs_data_reader_load_fragment_table(sqfs_data_reader_t *data,
+					 const sqfs_super_t *super)
+{
+	void *raw_frag;
+	size_t size;
+	uint32_t i;
+	int ret;
+
+	free(data->frag_block);
+	free(data->frag);
+
+	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 0;
+	}
+
+	if (super->fragment_table_start >= super->bytes_used)
+		return SQFS_ERROR_OUT_OF_BOUNDS;
+
+	if (SZ_MUL_OV(sizeof(data->frag[0]), super->fragment_entry_count,
+		      &size)) {
+		return SQFS_ERROR_OVERFLOW;
+	}
+
+	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;
+
+	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) {
+		data->frag[i].size = le32toh(data->frag[i].size);
+		data->frag[i].start_offset =
+			le64toh(data->frag[i].start_offset);
+	}
+
+	return 0;
+}
+
+void sqfs_data_reader_destroy(sqfs_data_reader_t *data)
+{
+	free(data->data_block);
+	free(data->frag_block);
+	free(data->frag);
+	free(data);
+}
+
+int sqfs_data_reader_get_block(sqfs_data_reader_t *data,
+			       const sqfs_inode_generic_t *inode,
+			       size_t index, sqfs_block_t **out)
+{
+	size_t i, unpacked_size;
+	uint64_t off, filesz;
+
+	if (inode->base.type == SQFS_INODE_FILE) {
+		off = inode->data.file.blocks_start;
+		filesz = inode->data.file.file_size;
+	} else if (inode->base.type == SQFS_INODE_EXT_FILE) {
+		off = inode->data.file_ext.blocks_start;
+		filesz = inode->data.file_ext.file_size;
+	} else {
+		return SQFS_ERROR_NOT_FILE;
+	}
+
+	if (index >= inode->num_file_blocks)
+		return SQFS_ERROR_OUT_OF_BOUNDS;
+
+	for (i = 0; i < index; ++i) {
+		off += SQFS_ON_DISK_BLOCK_SIZE(inode->block_sizes[i]);
+		filesz -= data->block_size;
+	}
+
+	unpacked_size = filesz < data->block_size ? filesz : data->block_size;
+
+	return get_block(data, off, inode->block_sizes[index],
+			 unpacked_size, out);
+}
+
+int sqfs_data_reader_get_fragment(sqfs_data_reader_t *data,
+				  const sqfs_inode_generic_t *inode,
+				  sqfs_block_t **out)
+{
+	uint32_t frag_idx, frag_off, frag_sz;
+	sqfs_block_t *blk;
+	uint64_t filesz;
+
+	if (inode->base.type == SQFS_INODE_EXT_FILE) {
+		filesz = inode->data.file_ext.file_size;
+		frag_idx = inode->data.file_ext.fragment_idx;
+		frag_off = inode->data.file_ext.fragment_offset;
+	} else if (inode->base.type == SQFS_INODE_FILE) {
+		filesz = inode->data.file.file_size;
+		frag_idx = inode->data.file.fragment_index;
+		frag_off = inode->data.file.fragment_offset;
+	} else {
+		return -1;
+	}
+
+	if (inode->num_file_blocks * data->block_size >= filesz) {
+		*out = NULL;
+		return 0;
+	}
+
+	frag_sz = filesz % data->block_size;
+
+	if (precache_fragment_block(data, frag_idx))
+		return -1;
+
+	if (frag_off + frag_sz > data->block_size)
+		return -1;
+
+	blk = alloc_flex(sizeof(*blk), 1, frag_sz);
+	if (blk == NULL)
+		return -1;
+
+	blk->size = frag_sz;
+	memcpy(blk->data, (char *)data->frag_block->data + frag_off, frag_sz);
+
+	*out = blk;
+	return 0;
+}
+
+ssize_t sqfs_data_reader_read(sqfs_data_reader_t *data,
+			      const sqfs_inode_generic_t *inode,
+			      uint64_t offset, void *buffer, size_t size)
+{
+	uint32_t frag_idx, frag_off;
+	size_t i, diff, total = 0;
+	uint64_t off, filesz;
+	char *ptr;
+
+	/* work out file location and size */
+	if (inode->base.type == SQFS_INODE_EXT_FILE) {
+		off = inode->data.file_ext.blocks_start;
+		filesz = inode->data.file_ext.file_size;
+		frag_idx = inode->data.file_ext.fragment_idx;
+		frag_off = inode->data.file_ext.fragment_offset;
+	} else {
+		off = inode->data.file.blocks_start;
+		filesz = inode->data.file.file_size;
+		frag_idx = inode->data.file.fragment_index;
+		frag_off = inode->data.file.fragment_offset;
+	}
+
+	/* find location of the first block */
+	i = 0;
+
+	while (offset > data->block_size && i < inode->num_file_blocks) {
+		off += SQFS_ON_DISK_BLOCK_SIZE(inode->block_sizes[i++]);
+		offset -= data->block_size;
+
+		if (filesz >= data->block_size) {
+			filesz -= data->block_size;
+		} else {
+			filesz = 0;
+		}
+	}
+
+	/* copy data from blocks */
+	while (i < inode->num_file_blocks && size > 0 && filesz > 0) {
+		diff = data->block_size - offset;
+		if (size < diff)
+			diff = size;
+
+		if (SQFS_IS_SPARSE_BLOCK(inode->block_sizes[i])) {
+			memset(buffer, 0, diff);
+		} else {
+			if (precache_data_block(data, off,
+						inode->block_sizes[i])) {
+				return -1;
+			}
+
+			memcpy(buffer, (char *)data->data_block->data + offset,
+			       diff);
+			off += SQFS_ON_DISK_BLOCK_SIZE(inode->block_sizes[i]);
+		}
+
+		if (filesz >= data->block_size) {
+			filesz -= data->block_size;
+		} else {
+			filesz = 0;
+		}
+
+		++i;
+		offset = 0;
+		size -= diff;
+		total += diff;
+		buffer = (char *)buffer + diff;
+	}
+
+	/* copy from fragment */
+	if (i == inode->num_file_blocks && size > 0 && filesz > 0) {
+		if (precache_fragment_block(data, frag_idx))
+			return -1;
+
+		if (frag_off + filesz > data->block_size)
+			return SQFS_ERROR_OUT_OF_BOUNDS;
+
+		if (offset >= filesz)
+			return total;
+
+		if (offset + size > filesz)
+			size = filesz - offset;
+
+		if (size == 0)
+			return total;
+
+		ptr = (char *)data->frag_block->data + frag_off + offset;
+		memcpy(buffer, ptr, size);
+		total += size;
+	}
+
+	return total;
+}