Update chgrp so -R works, tweaking DIRTREE_COMEAGAIN design along the way.

2 files changed, 47 insertions, 63 deletions

diff --git a/lib/dirtree.c b/lib/dirtree.c
index 25248f5d..418e8eb1 100644
--- a/lib/dirtree.c
+++ b/lib/dirtree.c
@@ -82,27 +82,33 @@ int dirtree_notdotdot(struct dirtree *catch)
 //
 // By default, allocates a tree of struct dirtree, not following symlinks
 // If callback==NULL, or callback always returns 0, allocate tree of struct
-// dirtree and return root of tree.  Otherwise call callback(node) on each hit, free
-// structures after use, and return NULL.
+// dirtree and return root of tree.  Otherwise call callback(node) on each
+// hit, free structures after use, and return NULL.
 //
 
 struct dirtree *handle_callback(struct dirtree *new,
 					int (*callback)(struct dirtree *node))
 {
-	int flags;
+	int flags, dir = S_ISDIR(new->st.st_mode);
 
 	if (!callback) callback = dirtree_notdotdot;
 
+	// Directory always has filehandle for examining contents. Whether or
+	// not we'll recurse into it gets decided later.
+
+	if (dir)
+		new->data = openat(new->parent ? new->parent->data : AT_FDCWD,
+			new->name, 0);
+
 	flags = callback(new);
-	if (S_ISDIR(new->st.st_mode)) {
-		if (flags & DIRTREE_RECURSE) {
-			new->data = openat (new->parent ? new->parent->data : AT_FDCWD,
-				new->name, 0);
+
+	if (dir) {
+		if (flags & (DIRTREE_RECURSE|DIRTREE_COMEAGAIN)) {
 			dirtree_recurse(new, callback);
-		}
-		new->data = -1;
-		if (flags & DIRTREE_COMEAGAIN) flags = callback(new);
+			if (flags & DIRTREE_COMEAGAIN) flags = callback(new);
+		} else close(new->data);
 	}
+
 	// If this had children, it was callback's job to free them already.
 	if (!(flags & DIRTREE_SAVE)) {
 		free(new);
@@ -128,7 +134,7 @@ void dirtree_recurse(struct dirtree *node,
 		free(path);
 		close(node->data);
 
-        return;
+		return;
 	}
 
 	// according to the fddir() man page, the filehandle in the DIR * can still
@@ -146,7 +152,9 @@ void dirtree_recurse(struct dirtree *node,
 		}
 	}
 
+	// This closes filehandle as well, so note it
 	closedir(dir);
+	node->data = -1;
 }
 
 // Create dirtree from path, using callback to filter nodes.
@@ -157,5 +165,5 @@ struct dirtree *dirtree_read(char *path, int (*callback)(struct dirtree *node))
 {
 	struct dirtree *root = dirtree_add_node(AT_FDCWD, path);
 
-	return handle_callback(root, callback);
+	return root ? handle_callback(root, callback) : DIRTREE_ABORTVAL;
 }
diff --git a/toys/chgrp.c b/toys/chgrp.c
index 2ad8a025..5cf43896 100644
--- a/toys/chgrp.c
+++ b/toys/chgrp.c
@@ -15,7 +15,7 @@ USE_CHGRP(NEWTOY(chgrp, "<2Rfv", TOYFLAG_BIN))
 
 config CHGRP
 	bool "chgrp"
-	default n
+	default y
 	help
 	  usage: chgrp [-R] [-f] [-v] group file...
 	  Change group ownership of one or more files.
@@ -38,31 +38,32 @@ DEFINE_GLOBALS(
 
 #define TT this.chgrp
 
-static int do_chgrp(const char *path) {
-	int ret = chown(path, -1, TT.group);
-	if (toys.optflags & FLAG_v)
-		xprintf("chgrp(%s, %s)\n", TT.group_name, path);
-	if (ret == -1 && !(toys.optflags & FLAG_f))
-		perror_msg("changing group of '%s' to '%s'", path, TT.group_name);
-	toys.exitval |= ret;
-	return ret;
-}
-
-// Copied from toys/cp.c:cp_node()
-int chgrp_node(char *path, struct dirtree *node)
+static int do_chgrp(struct dirtree *node)
 {
-	char *s = path + strlen(path);
-	struct dirtree *n = node;
-
-	for ( ; ; n = n->parent) {
-		while (s!=path) {
-			if (*(--s) == '/') break;
-		}
-		if (!n) break;
-	}
-	if (s != path) s++;
+	int fd, ret = 1, flags = toys.optflags;
+
+	if (!dirtree_notdotdot(node)) return 0;
+
+	// Handle recursion, and make it depth first
+	if (S_ISDIR(node->st.st_mode)) {
+		if (!node->extra) node->extra = dup(node->data);
+		if ((flags & FLAG_R) && node->data != -1) return DIRTREE_COMEAGAIN;
+		fd = node->extra;
+	} else fd = openat(node->parent ? node->parent->data : AT_FDCWD,
+		node->name, 0);
 
-	do_chgrp(s);
+	if (fd != -1) ret = fchown(fd, -1, TT.group);
+
+	if (ret || (flags & FLAG_v)) {
+		char *path = dirtree_path(node, 0);
+		if (flags & FLAG_v)
+			xprintf("chgrp(%s, %s)\n", TT.group_name, path);
+		if (ret == -1 && !(toys.optflags & FLAG_f))
+			perror_msg("changing group of '%s' to '%s'", path, TT.group_name);
+		free(path);
+	}
+	close(fd);
+	toys.exitval |= ret;
 
 	return 0;
 }
@@ -74,33 +75,8 @@ void chgrp_main(void)
 
 	TT.group_name = *toys.optargs;
 	group = getgrnam(TT.group_name);
-	if (!group) {
-		error_msg("invalid group '%s'", TT.group_name);
-		toys.exitval = 1;
-		return;
-	}
+	if (!group) error_exit("no group '%s'", TT.group_name);
 	TT.group = group->gr_gid;
 
-	if (toys.optflags & FLAG_R) {
-		// Recurse into subdirectories
-		for (s=toys.optargs + 1; *s; s++) {
-			struct stat sb;
-			if (stat(*s, &sb) == -1) {
-				if (!(toys.optflags & FLAG_f))
-					perror_msg("stat '%s'", *s);
-				continue;
-			}
-			do_chgrp(*s);
-			if (S_ISDIR(sb.st_mode)) {
-				strncpy(toybuf, *s, sizeof(toybuf) - 1);
-				toybuf[sizeof(toybuf) - 1] = 0;
-				dirtree_read(toybuf, NULL, chgrp_node);
-			}
-		}
-	} else {
-		// Do not recurse
-		for (s=toys.optargs + 1; *s; s++) {
-			do_chgrp(*s);
-		}
-	}
+	for (s=toys.optargs+1; *s; s++) dirtree_read(*s, do_chgrp);
 }