2 files changed, 51 insertions, 24 deletions

diff --git a/lib/io/src/xattr.c b/lib/io/src/dir_entry.c
index 85b7e53..b7f38f5 100644
--- a/lib/io/src/xattr.c
+++ b/lib/io/src/dir_entry.c
@@ -1,28 +1,61 @@
 /* SPDX-License-Identifier: GPL-3.0-or-later */
 /*
- * xattr.c
+ * dir_entry.c
  *
  * Copyright (C) 2023 David Oberhollenzer <goliath@infraroot.at>
  */
-#include "io/xattr.h"
+#include "io/dir_entry.h"
+#include "compat.h"
 
 #include <stdlib.h>
 #include <string.h>
 
+dir_entry_t *dir_entry_create(const char *name)
+{
+	size_t len, name_len;
+	dir_entry_t *out;
+
+	name_len = strlen(name);
+	if (SZ_ADD_OV(name_len, 1, &name_len))
+		return NULL;
+	if (SZ_ADD_OV(sizeof(*out), name_len, &len))
+		return NULL;
+
+	out = calloc(1, len);
+	if (out == NULL)
+		return NULL;
+
+	memcpy(out->name, name, name_len);
+	return out;
+}
+
 dir_entry_xattr_t *dir_entry_xattr_create(const char *key, const sqfs_u8 *value,
 					  size_t value_len)
 {
-	size_t key_len = strlen(key);
-	dir_entry_xattr_t *out = calloc(1, sizeof(*out) + key_len + 1 +
-					value_len + 1);
+	dir_entry_xattr_t *out;
+	size_t len, key_len;
+
+	/* key_ley = strlen(key) + 1 */
+	key_len = strlen(key);
+	if (SZ_ADD_OV(key_len, 1, &key_len))
+		return NULL;
+
+	/* len = key_len + value_len + 1 + sizeof(*out) */
+	if (SZ_ADD_OV(key_len, value_len, &len))
+		return NULL;
+	if (SZ_ADD_OV(len, 1, &len))
+		return NULL;
+	if (SZ_ADD_OV(len, sizeof(*out), &len))
+		return NULL;
 
+	out = calloc(1, len);
 	if (out != NULL) {
 		out->key = out->data;
-		out->value = (sqfs_u8 *)(out->data + key_len + 1);
+		out->value = (sqfs_u8 *)out->data + key_len;
+		out->value_len = value_len;
 
 		memcpy(out->key, key, key_len);
 		memcpy(out->value, value, value_len);
-		out->value_len = value_len;
 	}
 
 	return out;
diff --git a/lib/io/src/unix/dir_iterator.c b/lib/io/src/unix/dir_iterator.c
index 7154ec9..afbf0c0 100644
--- a/lib/io/src/unix/dir_iterator.c
+++ b/lib/io/src/unix/dir_iterator.c
@@ -76,8 +76,6 @@ static int dir_read_link(dir_iterator_t *base, char **out)
 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;
 
 	*out = NULL;
 	if (it->state != 0)
@@ -102,29 +100,25 @@ static int dir_next(dir_iterator_t *base, dir_entry_t **out)
 		return it->state;
 	}
 
-	len = strlen(it->ent->d_name);
-
-	decoded = alloc_flex(sizeof(*decoded), 1, len + 1);
-	if (decoded == NULL) {
+	*out = dir_entry_create(it->ent->d_name);
+	if ((*out) == NULL) {
 		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)->mtime = it->sb.st_mtime;
+	(*out)->dev = it->sb.st_dev;
+	(*out)->rdev = it->sb.st_rdev;
+	(*out)->uid = it->sb.st_uid;
+	(*out)->gid = it->sb.st_gid;
+	(*out)->mode = it->sb.st_mode;
 
 	if (S_ISREG(it->sb.st_mode))
-		decoded->size = it->sb.st_size;
+		(*out)->size = it->sb.st_size;
 
-	if (decoded->dev != it->device)
-		decoded->flags |= DIR_ENTRY_FLAG_MOUNT_POINT;
+	if ((*out)->dev != it->device)
+		(*out)->flags |= DIR_ENTRY_FLAG_MOUNT_POINT;
 
-	*out = decoded;
 	return it->state;
 }