5 files changed, 235 insertions, 86 deletions

diff --git a/bin/gensquashfs/src/fstree_from_dir.c b/bin/gensquashfs/src/fstree_from_dir.c
index c097304..b8b229b 100644
--- a/bin/gensquashfs/src/fstree_from_dir.c
+++ b/bin/gensquashfs/src/fstree_from_dir.c
@@ -37,6 +37,15 @@ static bool should_skip(const char *name, sqfs_u16 mode, unsigned int flags)
 	return false;
 }
 
+static sqfs_u32 clamp_timestamp(sqfs_s64 ts)
+{
+	if (ts < 0)
+		return 0;
+	if (ts > 0x0FFFFFFFFLL)
+		return 0xFFFFFFFF;
+	return ts;
+}
+
 #if defined(_WIN32) || defined(__WINDOWS__)
 static int add_node(fstree_t *fs, tree_node_t *root,
 		    scan_node_callback cb, void *user,
@@ -69,13 +78,7 @@ static int add_node(fstree_t *fs, tree_node_t *root,
 	}
 
 	if (flags & DIR_SCAN_KEEP_TIME) {
-		if (entry->mtime < 0) {
-			n->mod_time = 0;
-		} else if (entry->mtime > 0x0FFFFFFFFLL) {
-			n->mod_time = 0xFFFFFFFF;
-		} else {
-			n->mod_time = entry->mtime;
-		}
+		n->mod_time = clamp_timestamp(entry->mtime);
 	} else {
 		n->mod_time = fs->defaults.mtime;
 	}
@@ -136,116 +139,80 @@ static void discard_node(tree_node_t *root, tree_node_t *n)
 	free(n);
 }
 
-static char *read_link(const struct stat *sb, int dir_fd, const char *name)
-{
-	size_t size;
-	char *out;
-
-	if ((sizeof(sb->st_size) > sizeof(size_t)) && sb->st_size > SIZE_MAX) {
-		errno = EOVERFLOW;
-		return NULL;
-	}
-
-	if (SZ_ADD_OV((size_t)sb->st_size, 1, &size)) {
-		errno = EOVERFLOW;
-		return NULL;
-	}
-
-	out = calloc(1, size);
-	if (out == NULL)
-		return NULL;
-
-	if (readlinkat(dir_fd, name, out, (size_t)sb->st_size) < 0) {
-		int temp = errno;
-		free(out);
-		errno = temp;
-		return NULL;
-	}
-
-	out[sb->st_size] = '\0';
-	return out;
-}
-
 static int scan_dir(fstree_t *fs, tree_node_t *root, const char *path,
 		    scan_node_callback cb, void *user, unsigned int flags)
 {
-	char *extra = NULL;
-	dev_t devstart = 0;
-	struct dirent *ent;
-	struct stat sb;
-	tree_node_t *n;
-	DIR *dir;
-	int ret;
+	dir_iterator_t *dir;
 
-	dir = opendir(path);
-	if (dir == NULL) {
-		perror(path);
+	dir = dir_iterator_create(path);
+	if (dir == NULL)
 		return -1;
-	}
-
-	if (flags & DIR_SCAN_ONE_FILESYSTEM) {
-		if (fstat(dirfd(dir), &sb)) {
-			perror(path);
-			goto fail;
-		}
-		devstart = sb.st_dev;
-	}
 
 	for (;;) {
-		errno = 0;
-		ent = readdir(dir);
+		dir_entry_t *ent = NULL;
+		tree_node_t *n = NULL;
+		char *extra = NULL;
 
-		if (ent == NULL) {
-			if (errno) {
-				perror("readdir");
-				goto fail;
-			}
+		int ret = dir->next(dir, &ent);
+		if (ret > 0)
 			break;
+		if (ret < 0) {
+			sqfs_drop(dir);
+			return -1;
 		}
 
-		if (fstatat(dirfd(dir), ent->d_name,
-			    &sb, AT_SYMLINK_NOFOLLOW)) {
-			perror(ent->d_name);
-			goto fail;
-		}
-
-		if (should_skip(ent->d_name, sb.st_mode, flags))
+		if (should_skip(ent->name, ent->mode, flags)) {
+			free(ent);
 			continue;
+		}
 
-		if ((flags & DIR_SCAN_ONE_FILESYSTEM) && sb.st_dev != devstart)
+		if ((flags & DIR_SCAN_ONE_FILESYSTEM) && ent->dev != dir->dev) {
+			free(ent);
 			continue;
+		}
 
-		if (S_ISLNK(sb.st_mode)) {
-			extra = read_link(&sb, dirfd(dir), ent->d_name);
-			if (extra == NULL) {
-				perror("readlink");
-				goto fail;
+		if (S_ISLNK(ent->mode)) {
+			if (dir->read_link(dir, &extra)) {
+				free(ent);
+				return -1;
 			}
 		}
 
 		if (!(flags & DIR_SCAN_KEEP_TIME))
-			sb.st_mtime = fs->defaults.mtime;
+			ent->mtime = fs->defaults.mtime;
 
-		if (S_ISDIR(sb.st_mode) && (flags & DIR_SCAN_NO_DIR)) {
-			n = fstree_get_node_by_path(fs, root, ent->d_name,
+		if (S_ISDIR(ent->mode) && (flags & DIR_SCAN_NO_DIR)) {
+			n = fstree_get_node_by_path(fs, root, ent->name,
 						    false, false);
-			if (n == NULL)
+			if (n == NULL) {
+				free(ent);
 				continue;
+			}
 
 			ret = 0;
 		} else {
-			n = fstree_mknode(root, ent->d_name,
-					  strlen(ent->d_name), extra, &sb);
+			struct stat sb;
+
+			memset(&sb, 0, sizeof(sb));
+			sb.st_uid = ent->uid;
+			sb.st_gid = ent->gid;
+			sb.st_mode = ent->mode;
+			sb.st_mtime = clamp_timestamp(ent->mtime);
+
+			n = fstree_mknode(root, ent->name,
+					  strlen(ent->name), extra, &sb);
 			if (n == NULL) {
 				perror("creating tree node");
+				free(extra);
+				free(ent);
 				goto fail;
 			}
 
 			ret = (cb == NULL) ? 0 : cb(user, fs, n);
 		}
 
+		free(ent);
 		free(extra);
-		extra = NULL;
 
 		if (ret < 0)
 			goto fail;
@@ -254,11 +221,10 @@ static int scan_dir(fstree_t *fs, tree_node_t *root, const char *path,
 			discard_node(root, n);
 	}
 
-	closedir(dir);
+	sqfs_drop(dir);
 	return 0;
 fail:
-	closedir(dir);
-	free(extra);
+	sqfs_drop(dir);
 	return -1;
 }
 #endif
diff --git a/include/util/dir_iterator.h b/include/util/dir_iterator.h
index 66c13f3..77849e7 100644
--- a/include/util/dir_iterator.h
+++ b/include/util/dir_iterator.h
@@ -23,8 +23,12 @@ typedef struct {
 typedef struct dir_iterator_t {
 	sqfs_object_t obj;
 
+	sqfs_u64 dev;
+
 	int (*next)(struct dir_iterator_t *it,
 		    dir_entry_t **out);
+
+	int (*read_link)(struct dir_iterator_t *it, char **out);
 } dir_iterator_t;
 
 dir_iterator_t *dir_iterator_create(const char *path);
diff --git a/lib/util/Makemodule.am b/lib/util/Makemodule.am
index 8e9b01d..2c50594 100644
--- a/lib/util/Makemodule.am
+++ b/lib/util/Makemodule.am
@@ -16,6 +16,8 @@ libutil_a_CPPFLAGS = $(AM_CPPFLAGS)
 if WINDOWS
 libutil_a_CFLAGS += -DWINVER=0x0600 -D_WIN32_WINNT=0x0600
 libutil_a_SOURCES += lib/util/src/w32_dir_iterator.c
+else
+libutil_a_SOURCES += lib/util/src/unix_dir_iterator.c
 endif
 
 if HAVE_PTHREAD
diff --git a/lib/util/src/unix_dir_iterator.c b/lib/util/src/unix_dir_iterator.c
new file mode 100644
index 0000000..e712a45
--- /dev/null
+++ b/lib/util/src/unix_dir_iterator.c
@@ -0,0 +1,169 @@
+/* SPDX-License-Identifier: LGPL-3.0-or-later */
+/*
+ * unix_dir_iterator.c
+ *
+ * Copyright (C) 2023 David Oberhollenzer <goliath@infraroot.at>
+ */
+#include "config.h"
+#include "util/dir_iterator.h"
+#include "util/util.h"
+#include "sqfs/error.h"
+
+#include <sys/stat.h>
+#include <stdlib.h>
+#include <dirent.h>
+#include <string.h>
+#include <errno.h>
+
+typedef struct {
+	dir_iterator_t base;
+
+	struct dirent *ent;
+	struct stat sb;
+	int state;
+	DIR *dir;
+} unix_dir_iterator_t;
+
+static void dir_destroy(sqfs_object_t *obj)
+{
+	unix_dir_iterator_t *it = (unix_dir_iterator_t *)obj;
+
+	closedir(it->dir);
+	free(it);
+}
+
+static int dir_read_link(dir_iterator_t *base, char **out)
+{
+	unix_dir_iterator_t *it = (unix_dir_iterator_t *)base;
+	size_t size;
+	char *str;
+
+	*out = NULL;
+
+	if (it->state != 0 || it->ent == NULL) {
+		fputs("[BUG] no entry loaded, cannot readlink\n", stderr);
+		return SQFS_ERROR_INTERNAL;
+	}
+
+	if (S_ISLNK(it->sb.st_mode)) {
+		fprintf(stderr, "[BUG] %s is not a symlink, cannot readlink\n",
+			it->ent->d_name);
+		it->state = SQFS_ERROR_INTERNAL;
+		return SQFS_ERROR_INTERNAL;
+	}
+
+	if ((sizeof(it->sb.st_size) > sizeof(size_t)) &&
+	    it->sb.st_size > SIZE_MAX) {
+		goto fail_ov;
+	}
+
+	if (SZ_ADD_OV((size_t)it->sb.st_size, 1, &size))
+		goto fail_ov;
+
+	str = calloc(1, size);
+	if (str == NULL) {
+		perror(it->ent->d_name);
+		it->state = SQFS_ERROR_ALLOC;
+		return it->state;
+	}
+
+	if (readlinkat(dirfd(it->dir), it->ent->d_name,
+		       str, (size_t)it->sb.st_size) < 0) {
+		fprintf(stderr, "%s: readlink: %s\n", it->ent->d_name,
+			strerror(errno));
+		free(str);
+		it->state = SQFS_ERROR_INTERNAL;
+		return it->state;
+	}
+
+	str[it->sb.st_size] = '\0';
+
+	*out = str;
+	return it->state;
+fail_ov:
+	fprintf(stderr, "%s: link target too long\n", it->ent->d_name);
+	return SQFS_ERROR_OVERFLOW;
+}
+
+static int dir_next(dir_iterator_t *base, dir_entry_t **out)
+{
+	unix_dir_iterator_t *it = (unix_dir_iterator_t *)base;
+	dir_entry_t *decoded;
+	size_t len;
+
+	if (it->state != 0)
+		return it->state;
+
+	errno = 0;
+	it->ent = readdir(it->dir);
+
+	if (it->ent == NULL) {
+		if (errno != 0) {
+			perror("readdir");
+			it->state = SQFS_ERROR_INTERNAL;
+		} else {
+			it->state = 1;
+		}
+
+		return it->state;
+	}
+
+	if (fstatat(dirfd(it->dir), it->ent->d_name, &it->sb, AT_SYMLINK_NOFOLLOW)) {
+		perror(it->ent->d_name);
+		it->state = SQFS_ERROR_INTERNAL;
+		return it->state;
+	}
+
+	len = strlen(it->ent->d_name);
+
+	decoded = alloc_flex(sizeof(*decoded), 1, len + 1);
+	if (decoded == NULL) {
+		perror(it->ent->d_name);
+		it->state = SQFS_ERROR_ALLOC;
+		return it->state;
+	}
+
+	memcpy(decoded->name, it->ent->d_name, len);
+	decoded->mtime = it->sb.st_mtime;
+	decoded->dev = it->sb.st_dev;
+	decoded->rdev = it->sb.st_rdev;
+	decoded->uid = it->sb.st_uid;
+	decoded->gid = it->sb.st_gid;
+	decoded->mode = it->sb.st_mode;
+
+	*out = decoded;
+	return it->state;
+}
+
+dir_iterator_t *dir_iterator_create(const char *path)
+{
+	unix_dir_iterator_t *it = calloc(1, sizeof(*it));
+
+	if (it == NULL) {
+		perror(path);
+		return NULL;
+	}
+
+	it->state = 0;
+	it->dir = opendir(path);
+
+	if (it->dir == NULL) {
+		perror(path);
+		free(it);
+		return NULL;
+	}
+
+	if (fstat(dirfd(it->dir), &it->sb)) {
+		perror(path);
+		closedir(it->dir);
+		free(it);
+		return NULL;
+	}
+
+	sqfs_object_init(it, dir_destroy, NULL);
+	((dir_iterator_t *)it)->dev = it->sb.st_dev;
+	((dir_iterator_t *)it)->next = dir_next;
+	((dir_iterator_t *)it)->read_link = dir_read_link;
+
+	return (dir_iterator_t *)it;
+}
diff --git a/lib/util/src/w32_dir_iterator.c b/lib/util/src/w32_dir_iterator.c
index 1a330ae..9ec8b18 100644
--- a/lib/util/src/w32_dir_iterator.c
+++ b/lib/util/src/w32_dir_iterator.c
@@ -47,6 +47,13 @@ static sqfs_s64 w32time_to_unix(const FILETIME *ft)
 	return w32ts - UNIX_EPOCH_ON_W32;
 }
 
+static int dir_iterator_read_link(dir_iterator_t *it, char **out)
+{
+	(void)it;
+	*out = NULL;
+	return SQFS_ERROR_UNSUPPORTED;
+}
+
 static int dir_iterator_next(dir_iterator_t *it, dir_entry_t **out)
 {
 	dir_iterator_win32_t *w32 = (dir_iterator_win32_t *)it;
@@ -177,6 +184,7 @@ dir_iterator_t *dir_iterator_create(const char *path)
 	sqfs_object_init(it, dir_iterator_destroy, NULL);
 
 	((dir_iterator_t *)it)->next = dir_iterator_next;
+	((dir_iterator_t *)it)->read_link = dir_iterator_read_link;
 
 	it->state = STATE_HAVE_ENTRY;
 	it->dirhnd = dirhnd;