5 files changed, 257 insertions, 24 deletions

diff --git a/lib/sqfs/Makemodule.am b/lib/sqfs/Makemodule.am
index c37301b..06ffea1 100644
--- a/lib/sqfs/Makemodule.am
+++ b/lib/sqfs/Makemodule.am
@@ -33,7 +33,7 @@ libsquashfs_la_SOURCES += lib/sqfs/block_processor/block_processor.c
 libsquashfs_la_SOURCES += lib/sqfs/block_processor/backend.c
 libsquashfs_la_SOURCES += lib/sqfs/frag_table.c include/sqfs/frag_table.h
 libsquashfs_la_SOURCES += lib/sqfs/block_writer.c include/sqfs/block_writer.h
-libsquashfs_la_SOURCES += lib/sqfs/misc.c
+libsquashfs_la_SOURCES += lib/sqfs/misc.c lib/sqfs/dir_reader/dcache.c
 libsquashfs_la_CPPFLAGS = $(AM_CPPFLAGS)
 libsquashfs_la_LDFLAGS = $(AM_LDFLAGS) -version-info $(LIBSQUASHFS_SO_VERSION)
 libsquashfs_la_CFLAGS = $(AM_CFLAGS) $(PTHREAD_CFLAGS) $(ZLIB_CFLAGS)
diff --git a/lib/sqfs/dir_reader/dcache.c b/lib/sqfs/dir_reader/dcache.c
new file mode 100644
index 0000000..47fbf73
--- /dev/null
+++ b/lib/sqfs/dir_reader/dcache.c
@@ -0,0 +1,73 @@
+/* SPDX-License-Identifier: LGPL-3.0-or-later */
+/*
+ * dcache.c
+ *
+ * Copyright (C) 2019 David Oberhollenzer <goliath@infraroot.at>
+ */
+#define SQFS_BUILDING_DLL
+#include "internal.h"
+
+static int dcache_key_compare(const void *ctx, const void *l, const void *r)
+{
+	sqfs_u32 lhs = *((const sqfs_u32 *)l), rhs = *((const sqfs_u32 *)r);
+	(void)ctx;
+
+	return lhs < rhs ? -1 : (lhs > rhs ? 1 : 0);
+}
+
+int sqfs_dir_reader_dcache_init(sqfs_dir_reader_t *rd, sqfs_u32 flags)
+{
+	if (!(flags & SQFS_DIR_READER_DOT_ENTRIES))
+		return 0;
+
+	return rbtree_init(&rd->dcache, sizeof(sqfs_u32), sizeof(sqfs_u64),
+			   dcache_key_compare);
+}
+
+int sqfs_dir_reader_dcache_init_copy(sqfs_dir_reader_t *copy,
+				     const sqfs_dir_reader_t *rd)
+{
+	if (!(rd->flags & SQFS_DIR_READER_DOT_ENTRIES))
+		return 0;
+
+	return rbtree_copy(&rd->dcache, &copy->dcache);
+}
+
+void sqfs_dir_reader_dcache_cleanup(sqfs_dir_reader_t *rd)
+{
+	if (!(rd->flags & SQFS_DIR_READER_DOT_ENTRIES))
+		return;
+
+	rbtree_cleanup(&rd->dcache);
+}
+
+int sqfs_dir_reader_dcache_add(sqfs_dir_reader_t *rd,
+			       sqfs_u32 inode, sqfs_u64 ref)
+{
+	rbtree_node_t *node;
+
+	if (!(rd->flags & SQFS_DIR_READER_DOT_ENTRIES))
+		return 0;
+
+	node = rbtree_lookup(&rd->dcache, &inode);
+	if (node != NULL)
+		return 0;
+
+	return rbtree_insert(&rd->dcache, &inode, &ref);
+}
+
+int sqfs_dir_reader_dcache_find(sqfs_dir_reader_t *rd,
+				sqfs_u32 inode, sqfs_u64 *ref)
+{
+	rbtree_node_t *node;
+
+	if (!(rd->flags & SQFS_DIR_READER_DOT_ENTRIES))
+		return SQFS_ERROR_NO_ENTRY;
+
+	node = rbtree_lookup(&rd->dcache, &inode);
+	if (node == NULL)
+		return SQFS_ERROR_NO_ENTRY;
+
+	*ref = *((sqfs_u64 *)rbtree_node_value(node));
+	return 0;
+}
diff --git a/lib/sqfs/dir_reader/dir_reader.c b/lib/sqfs/dir_reader/dir_reader.c
index 969b71d..8a074aa 100644
--- a/lib/sqfs/dir_reader/dir_reader.c
+++ b/lib/sqfs/dir_reader/dir_reader.c
@@ -11,6 +11,7 @@ static void dir_reader_destroy(sqfs_object_t *obj)
 {
 	sqfs_dir_reader_t *rd = (sqfs_dir_reader_t *)obj;
 
+	sqfs_dir_reader_dcache_cleanup(rd);
 	sqfs_destroy(rd->meta_inode);
 	sqfs_destroy(rd->meta_dir);
 	free(rd);
@@ -26,6 +27,9 @@ static sqfs_object_t *dir_reader_copy(const sqfs_object_t *obj)
 
 	memcpy(copy, rd, sizeof(*copy));
 
+	if (sqfs_dir_reader_dcache_init_copy(copy, rd))
+		goto fail_cache;
+
 	copy->meta_inode = sqfs_copy(rd->meta_inode);
 	if (copy->meta_inode == NULL)
 		goto fail_mino;
@@ -38,6 +42,8 @@ static sqfs_object_t *dir_reader_copy(const sqfs_object_t *obj)
 fail_mdir:
 	sqfs_destroy(copy->meta_inode);
 fail_mino:
+	sqfs_dir_reader_dcache_cleanup(copy);
+fail_cache:
 	free(copy);
 	return NULL;
 }
@@ -50,22 +56,22 @@ sqfs_dir_reader_t *sqfs_dir_reader_create(const sqfs_super_t *super,
 	sqfs_dir_reader_t *rd;
 	sqfs_u64 start, limit;
 
-	if (flags != 0)
+	if (flags & ~SQFS_DIR_READER_ALL_FLAGS)
 		return NULL;
 
 	rd = calloc(1, sizeof(*rd));
 	if (rd == NULL)
 		return NULL;
 
+	if (sqfs_dir_reader_dcache_init(rd, flags))
+		goto fail_dcache;
+
 	start = super->inode_table_start;
 	limit = super->directory_table_start;
 
 	rd->meta_inode = sqfs_meta_reader_create(file, cmp, start, limit);
-
-	if (rd->meta_inode == NULL) {
-		free(rd);
-		return NULL;
-	}
+	if (rd->meta_inode == NULL)
+		goto fail_mino;
 
 	start = super->directory_table_start;
 	limit = super->id_table_start;
@@ -77,17 +83,22 @@ sqfs_dir_reader_t *sqfs_dir_reader_create(const sqfs_super_t *super,
 		limit = super->export_table_start;
 
 	rd->meta_dir = sqfs_meta_reader_create(file, cmp, start, limit);
-
-	if (rd->meta_dir == NULL) {
-		sqfs_destroy(rd->meta_inode);
-		free(rd);
-		return NULL;
-	}
+	if (rd->meta_dir == NULL)
+		goto fail_mdir;
 
 	((sqfs_object_t *)rd)->destroy = dir_reader_destroy;
 	((sqfs_object_t *)rd)->copy = dir_reader_copy;
 	rd->super = super;
+	rd->flags = flags;
+	rd->state = DIR_STATE_NONE;
 	return rd;
+fail_mdir:
+	sqfs_destroy(rd->meta_inode);
+fail_mino:
+	sqfs_dir_reader_dcache_cleanup(rd);
+fail_dcache:
+	free(rd);
+	return NULL;
 }
 
 int sqfs_dir_reader_open_dir(sqfs_dir_reader_t *rd,
@@ -96,15 +107,18 @@ int sqfs_dir_reader_open_dir(sqfs_dir_reader_t *rd,
 {
 	sqfs_u64 block_start;
 	size_t size, offset;
+	sqfs_u32 parent;
 
-	if (flags != 0)
+	if (flags & (~SQFS_DIR_OPEN_ALL_FLAGS))
 		return SQFS_ERROR_UNSUPPORTED;
 
 	if (inode->base.type == SQFS_INODE_DIR) {
+		parent = inode->data.dir.parent_inode;
 		size = inode->data.dir.size;
 		offset = inode->data.dir.offset;
 		block_start = inode->data.dir.start_block;
 	} else if (inode->base.type == SQFS_INODE_EXT_DIR) {
+		parent = inode->data.dir_ext.parent_inode;
 		size = inode->data.dir_ext.size;
 		offset = inode->data.dir_ext.offset;
 		block_start = inode->data.dir_ext.start_block;
@@ -112,6 +126,27 @@ int sqfs_dir_reader_open_dir(sqfs_dir_reader_t *rd,
 		return SQFS_ERROR_NOT_DIR;
 	}
 
+	if ((rd->flags & SQFS_DIR_READER_DOT_ENTRIES) &&
+	    !(flags & SQFS_DIR_OPEN_NO_DOT_ENTRIES)) {
+		if (sqfs_dir_reader_dcache_find(rd, inode->base.inode_number,
+						&rd->cur_ref)) {
+			return SQFS_ERROR_NO_ENTRY;
+		}
+
+		if (rd->cur_ref == rd->super->root_inode_ref) {
+			rd->parent_ref = rd->cur_ref;
+		} else if (sqfs_dir_reader_dcache_find(rd, parent,
+						       &rd->parent_ref)) {
+			return SQFS_ERROR_NO_ENTRY;
+		}
+
+		rd->state = DIR_STATE_OPENED;
+	} else {
+		rd->state = DIR_STATE_ENTRIES;
+	}
+
+	rd->start_state = rd->state;
+
 	memset(&rd->hdr, 0, sizeof(rd->hdr));
 	rd->size = size;
 	rd->entries = 0;
@@ -128,12 +163,50 @@ int sqfs_dir_reader_open_dir(sqfs_dir_reader_t *rd,
 	return sqfs_meta_reader_seek(rd->meta_dir, block_start, offset);
 }
 
+static int mk_dummy_entry(const char *str, sqfs_dir_entry_t **out)
+{
+	size_t len = strlen(str);
+	sqfs_dir_entry_t *ent;
+
+	ent = calloc(1, sizeof(sqfs_dir_entry_t) + len + 1);
+	if (ent == NULL)
+		return SQFS_ERROR_ALLOC;
+
+	ent->type = SQFS_INODE_DIR;
+	ent->size = len - 1;
+
+	strcpy((char *)ent->name, str);
+
+	*out = ent;
+	return 0;
+}
+
 int sqfs_dir_reader_read(sqfs_dir_reader_t *rd, sqfs_dir_entry_t **out)
 {
 	sqfs_dir_entry_t *ent;
 	size_t count;
 	int err;
 
+	switch (rd->state) {
+	case DIR_STATE_OPENED:
+		err = mk_dummy_entry(".", out);
+		if (err == 0)
+			rd->state = DIR_STATE_DOT;
+		return err;
+	case DIR_STATE_DOT:
+		err = mk_dummy_entry("..", out);
+		if (err == 0)
+			rd->state = DIR_STATE_DOT_DOT;
+		return err;
+	case DIR_STATE_DOT_DOT:
+		rd->state = DIR_STATE_ENTRIES;
+		break;
+	case DIR_STATE_ENTRIES:
+		break;
+	default:
+		return SQFS_ERROR_SEQUENCE;
+	}
+
 	if (!rd->entries) {
 		if (rd->size <= sizeof(rd->hdr))
 			return 1;
@@ -173,9 +246,13 @@ int sqfs_dir_reader_read(sqfs_dir_reader_t *rd, sqfs_dir_entry_t **out)
 
 int sqfs_dir_reader_rewind(sqfs_dir_reader_t *rd)
 {
+	if (rd->state == DIR_STATE_NONE)
+		return SQFS_ERROR_SEQUENCE;
+
 	memset(&rd->hdr, 0, sizeof(rd->hdr));
 	rd->size = rd->start_size;
 	rd->entries = 0;
+	rd->state = rd->start_state;
 
 	if (rd->size <= sizeof(rd->hdr))
 		return 0;
@@ -189,7 +266,10 @@ int sqfs_dir_reader_find(sqfs_dir_reader_t *rd, const char *name)
 	sqfs_dir_entry_t *ent;
 	int ret;
 
-	if (rd->size != rd->start_size) {
+	if (rd->state == DIR_STATE_NONE)
+		return SQFS_ERROR_SEQUENCE;
+
+	if (rd->size != rd->start_size || rd->state != rd->start_state) {
 		ret = sqfs_dir_reader_rewind(rd);
 		if (ret)
 			return ret;
@@ -213,12 +293,42 @@ int sqfs_dir_reader_get_inode(sqfs_dir_reader_t *rd,
 			      sqfs_inode_generic_t **inode)
 {
 	sqfs_u64 block_start;
+	sqfs_u16 offset;
+	int ret;
+
+	switch (rd->state) {
+	case DIR_STATE_DOT:
+		block_start = rd->cur_ref >> 16;
+		offset = rd->cur_ref & 0x0FFFF;
+		break;
+	case DIR_STATE_DOT_DOT:
+		block_start = rd->parent_ref >> 16;
+		offset = rd->parent_ref & 0x0FFFF;
+		break;
+	case DIR_STATE_ENTRIES:
+		block_start = rd->hdr.start_block;
+		offset = rd->inode_offset;
+		break;
+	default:
+		return SQFS_ERROR_SEQUENCE;
+	}
 
-	block_start = rd->hdr.start_block;
+	ret = sqfs_meta_reader_read_inode(rd->meta_inode, rd->super,
+					  block_start, offset, inode);
+	if (ret != 0)
+		return ret;
 
-	return sqfs_meta_reader_read_inode(rd->meta_inode, rd->super,
-					   block_start, rd->inode_offset,
-					   inode);
+	if ((*inode)->base.type == SQFS_INODE_DIR ||
+	    (*inode)->base.type == SQFS_INODE_EXT_DIR) {
+		sqfs_u32 inum = (*inode)->base.inode_number;
+		sqfs_u64 ref = (block_start << 16) | rd->inode_offset;
+
+		ret = sqfs_dir_reader_dcache_add(rd, inum, ref);
+		if (ret != 0)
+			return ret;
+	}
+
+	return 0;
 }
 
 int sqfs_dir_reader_get_root_inode(sqfs_dir_reader_t *rd,
@@ -226,7 +336,23 @@ int sqfs_dir_reader_get_root_inode(sqfs_dir_reader_t *rd,
 {
 	sqfs_u64 block_start = rd->super->root_inode_ref >> 16;
 	sqfs_u16 offset = rd->super->root_inode_ref & 0xFFFF;
+	int ret;
+
+	ret = sqfs_meta_reader_read_inode(rd->meta_inode, rd->super,
+					  block_start, offset, inode);
 
-	return sqfs_meta_reader_read_inode(rd->meta_inode, rd->super,
-					   block_start, offset, inode);
+	if (ret != 0)
+		return ret;
+
+	if ((*inode)->base.type == SQFS_INODE_DIR ||
+	    (*inode)->base.type == SQFS_INODE_EXT_DIR) {
+		sqfs_u32 inum = (*inode)->base.inode_number;
+		sqfs_u64 ref = rd->super->root_inode_ref;
+
+		ret = sqfs_dir_reader_dcache_add(rd, inum, ref);
+		if (ret != 0)
+			return ret;
+	}
+
+	return 0;
 }
diff --git a/lib/sqfs/dir_reader/internal.h b/lib/sqfs/dir_reader/internal.h
index ff162ff..cd20b69 100644
--- a/lib/sqfs/dir_reader/internal.h
+++ b/lib/sqfs/dir_reader/internal.h
@@ -17,11 +17,20 @@
 #include "sqfs/inode.h"
 #include "sqfs/error.h"
 #include "sqfs/dir.h"
+#include "rbtree.h"
 #include "util.h"
 
 #include <string.h>
 #include <stdlib.h>
 
+enum {
+	DIR_STATE_NONE = 0,
+	DIR_STATE_OPENED = 1,
+	DIR_STATE_DOT = 2,
+	DIR_STATE_DOT_DOT = 3,
+	DIR_STATE_ENTRIES = 4,
+};
+
 struct sqfs_dir_reader_t {
 	sqfs_object_t base;
 
@@ -37,6 +46,28 @@ struct sqfs_dir_reader_t {
 	size_t start_size;
 	sqfs_u16 dir_offset;
 	sqfs_u16 inode_offset;
+
+	sqfs_u32 flags;
+
+	int start_state;
+	int state;
+	sqfs_u64 parent_ref;
+	sqfs_u64 cur_ref;
+	rbtree_t dcache;
 };
 
+SQFS_INTERNAL int sqfs_dir_reader_dcache_init(sqfs_dir_reader_t *rd,
+					      sqfs_u32 flags);
+
+SQFS_INTERNAL int sqfs_dir_reader_dcache_init_copy(sqfs_dir_reader_t *copy,
+						   const sqfs_dir_reader_t *rd);
+
+SQFS_INTERNAL int sqfs_dir_reader_dcache_add(sqfs_dir_reader_t *rd,
+					     sqfs_u32 inode, sqfs_u64 ref);
+
+SQFS_INTERNAL int sqfs_dir_reader_dcache_find(sqfs_dir_reader_t *rd,
+					      sqfs_u32 inode, sqfs_u64 *ref);
+
+SQFS_INTERNAL void sqfs_dir_reader_dcache_cleanup(sqfs_dir_reader_t *rd);
+
 #endif /* DIR_READER_INTERNAL_H */
diff --git a/lib/sqfs/dir_reader/read_tree.c b/lib/sqfs/dir_reader/read_tree.c
index d173ef7..7fa944a 100644
--- a/lib/sqfs/dir_reader/read_tree.c
+++ b/lib/sqfs/dir_reader/read_tree.c
@@ -113,7 +113,8 @@ static int fill_dir(sqfs_dir_reader_t *dr, sqfs_tree_node_t *root,
 		if (n->inode->base.type == SQFS_INODE_DIR ||
 		    n->inode->base.type == SQFS_INODE_EXT_DIR) {
 			if (!(flags & SQFS_TREE_NO_RECURSE)) {
-				err = sqfs_dir_reader_open_dir(dr, n->inode, 0);
+				err = sqfs_dir_reader_open_dir(dr, n->inode,
+					       SQFS_DIR_OPEN_NO_DOT_ENTRIES);
 				if (err)
 					return err;
 
@@ -212,7 +213,8 @@ int sqfs_dir_reader_get_full_hierarchy(sqfs_dir_reader_t *rd,
 			continue;
 		}
 
-		ret = sqfs_dir_reader_open_dir(rd, tail->inode, 0);
+		ret = sqfs_dir_reader_open_dir(rd, tail->inode,
+					       SQFS_DIR_OPEN_NO_DOT_ENTRIES);
 		if (ret)
 			goto fail;
 
@@ -271,7 +273,8 @@ int sqfs_dir_reader_get_full_hierarchy(sqfs_dir_reader_t *rd,
 
 	if (tail->inode->base.type == SQFS_INODE_DIR ||
 	    tail->inode->base.type == SQFS_INODE_EXT_DIR) {
-		ret = sqfs_dir_reader_open_dir(rd, tail->inode, 0);
+		ret = sqfs_dir_reader_open_dir(rd, tail->inode,
+					       SQFS_DIR_OPEN_NO_DOT_ENTRIES);
 		if (ret)
 			goto fail;