ctx: Try to reset TTY state when terminating abnormally

Fixes: https://github.com/tavianator/bfs/issues/138

3 files changed, 105 insertions, 51 deletions

diff --git a/src/color.c b/src/color.c
index f004bf2..326291d 100644
--- a/src/color.c
+++ b/src/color.c
@@ -161,8 +161,13 @@ static struct esc_seq **get_esc(const struct colors *colors, const char *name) {
 	return leaf ? leaf->value : NULL;
 }
 
+/** Append an escape sequence to a string. */
+static int cat_esc(dchar **dstr, const struct esc_seq *seq) {
+	return dstrxcat(dstr, seq->seq, seq->len);
+}
+
 /** Set a named escape sequence. */
-static int set_esc(struct colors *colors, const char *name, char *value) {
+static int set_esc(struct colors *colors, const char *name, dchar *value) {
 	struct esc_seq **field = get_esc(colors, name);
 	if (!field) {
 		return 0;
@@ -587,8 +592,8 @@ static int parse_gnu_ls_colors(struct colors *colors, const char *ls_colors) {
 
 			// All-zero values should be treated like NULL, to fall
 			// back on any other relevant coloring for that file
-			char *esc = value;
-			if (strspn(value, "0") == strlen(value)
+			dchar *esc = value;
+			if (strspn(value, "0") == dstrlen(value)
 			    && strcmp(key, "rs") != 0
 			    && strcmp(key, "lc") != 0
 			    && strcmp(key, "rc") != 0
@@ -693,6 +698,20 @@ struct colors *parse_colors(void) {
 		colors->link->len = 0;
 	}
 
+	// Pre-compute the reset escape sequence
+	if (!colors->endcode) {
+		dchar *ec = dstralloc(0);
+		if (!ec
+		    || cat_esc(&ec, colors->leftcode) != 0
+		    || cat_esc(&ec, colors->reset) != 0
+		    || cat_esc(&ec, colors->rightcode) != 0
+		    || set_esc(colors, "ec", ec) != 0) {
+			dstrfree(ec);
+			goto fail;
+		}
+		dstrfree(ec);
+	}
+
 	return colors;
 
 fail:
@@ -727,10 +746,11 @@ CFILE *cfwrap(FILE *file, const struct colors *colors, bool close) {
 	}
 
 	cfile->file = file;
+	cfile->fd = fileno(file);
 	cfile->need_reset = false;
 	cfile->close = close;
 
-	if (isatty(fileno(file))) {
+	if (isatty(cfile->fd)) {
 		cfile->colors = colors;
 	} else {
 		cfile->colors = NULL;
@@ -874,7 +894,7 @@ error:
 
 /** Print an escape sequence chunk. */
 static int print_esc_chunk(CFILE *cfile, const struct esc_seq *esc) {
-	return dstrxcat(&cfile->buffer, esc->seq, esc->len);
+	return cat_esc(&cfile->buffer, esc);
 }
 
 /** Print an ANSI escape sequence. */
@@ -908,12 +928,7 @@ static int print_reset(CFILE *cfile) {
 	}
 	cfile->need_reset = false;
 
-	const struct colors *colors = cfile->colors;
-	if (colors->endcode) {
-		return print_esc_chunk(cfile, colors->endcode);
-	} else {
-		return print_esc(cfile, colors->reset);
-	}
+	return print_esc_chunk(cfile, cfile->colors->endcode);
 }
 
 /** Print a shell-escaped string. */
@@ -1390,3 +1405,14 @@ int cfprintf(CFILE *cfile, const char *format, ...) {
 	va_end(args);
 	return ret;
 }
+
+int cfreset(CFILE *cfile) {
+	const struct colors *colors = cfile->colors;
+	if (!colors) {
+		return 0;
+	}
+
+	const struct esc_seq *esc = colors->endcode;
+	size_t ret = xwrite(cfile->fd, esc->seq, esc->len);
+	return ret == esc->len ? 0 : -1;
+}
diff --git a/src/color.h b/src/color.h
index 3278cd6..3550888 100644
--- a/src/color.h
+++ b/src/color.h
@@ -42,6 +42,8 @@ typedef struct CFILE {
 	const struct colors *colors;
 	/** A buffer for colored formatting. */
 	dchar *buffer;
+	/** Cached file descriptor number. */
+	int fd;
 	/** Whether the next ${rs} is actually necessary. */
 	bool need_reset;
 	/** Whether to close the underlying stream. */
@@ -108,4 +110,9 @@ int cfprintf(CFILE *cfile, const char *format, ...);
 attr(printf(2, 0))
 int cvfprintf(CFILE *cfile, const char *format, va_list args);
 
+/**
+ * Reset the TTY state when terminating abnormally (async-signal-safe).
+ */
+int cfreset(CFILE *cfile);
+
 #endif // BFS_COLOR_H
diff --git a/src/ctx.c b/src/ctx.c
index aa73b35..356adbc 100644
--- a/src/ctx.c
+++ b/src/ctx.c
@@ -9,11 +9,13 @@
 #include "list.h"
 #include "mtab.h"
 #include "pwcache.h"
+#include "sighook.h"
 #include "stat.h"
 #include "trie.h"
 #include "xtime.h"
 #include <errno.h>
 #include <limits.h>
+#include <signal.h>
 #include <stdio.h>
 #include <stdlib.h>
 #include <unistd.h>
@@ -98,13 +100,20 @@ struct bfs_ctx_file {
 	CFILE *cfile;
 	/** The path to the file (for diagnostics). */
 	const char *path;
+	/** Signal hook to send a reset escape sequence. */
+	struct sighook *hook;
 	/** Remembers I/O errors, to propagate them to the exit status. */
 	int error;
 };
 
+/** Call cfreset() on a tracked file. */
+static void cfreset_hook(int sig, siginfo_t *info, void *arg) {
+	cfreset(arg);
+}
+
 CFILE *bfs_ctx_dedup(struct bfs_ctx *ctx, CFILE *cfile, const char *path) {
 	struct bfs_stat sb;
-	if (bfs_stat(fileno(cfile->file), NULL, 0, &sb) != 0) {
+	if (bfs_stat(cfile->fd, NULL, 0, &sb) != 0) {
 		return NULL;
 	}
 
@@ -124,19 +133,31 @@ CFILE *bfs_ctx_dedup(struct bfs_ctx *ctx, CFILE *cfile, const char *path) {
 
 	leaf->value = ctx_file = ALLOC(struct bfs_ctx_file);
 	if (!ctx_file) {
-		trie_remove(&ctx->files, leaf);
-		return NULL;
+		goto fail;
 	}
 
 	ctx_file->cfile = cfile;
 	ctx_file->path = path;
 	ctx_file->error = 0;
+	ctx_file->hook = NULL;
+
+	if (cfile->colors) {
+		ctx_file->hook = atsigexit(cfreset_hook, cfile);
+		if (!ctx_file->hook) {
+			goto fail;
+		}
+	}
 
 	if (cfile != ctx->cout && cfile != ctx->cerr) {
 		++ctx->nfiles;
 	}
 
 	return cfile;
+
+fail:
+	trie_remove(&ctx->files, leaf);
+	free(ctx_file);
+	return NULL;
 }
 
 void bfs_ctx_flush(const struct bfs_ctx *ctx) {
@@ -156,7 +177,7 @@ void bfs_ctx_flush(const struct bfs_ctx *ctx) {
 
 		const char *path = ctx_file->path;
 		if (path) {
-			bfs_error(ctx, "'%s': %m.\n", path);
+			bfs_error(ctx, "%pq: %m.\n", path);
 		} else if (cfile == ctx->cout) {
 			bfs_error(ctx, "(standard output): %m.\n");
 		}
@@ -188,30 +209,49 @@ static int bfs_ctx_fflush(CFILE *cfile) {
 static int bfs_ctx_fclose(struct bfs_ctx *ctx, struct bfs_ctx_file *ctx_file) {
 	CFILE *cfile = ctx_file->cfile;
 
-	if (cfile == ctx->cout) {
-		// Will be checked later
-		return 0;
-	} else if (cfile == ctx->cerr) {
-		// Writes to stderr are allowed to fail silently, unless the same file was used by
-		// -fprint, -fls, etc.
-		if (ctx_file->path) {
-			return bfs_ctx_fflush(cfile);
-		} else {
-			return 0;
-		}
-	}
-
+	// Writes to stderr are allowed to fail silently, unless the same file
+	// was used by -fprint, -fls, etc.
+	bool silent = cfile == ctx->cerr && !ctx_file->path;
 	int ret = 0, error = 0;
-	if (ferror(cfile->file)) {
+
+	if (ctx_file->error) {
+		// An error was previously reported during bfs_ctx_flush()
 		ret = -1;
-		error = EIO;
+		error = ctx_file->error;
 	}
-	if (cfclose(cfile) != 0) {
+
+	// Flush the file just before we remove the hook, to maximize the chance
+	// we leave the TTY in a good state
+	if (bfs_ctx_fflush(cfile) != 0) {
 		ret = -1;
 		error = errno;
 	}
 
-	errno = error;
+	if (ctx_file->hook) {
+		sigunhook(ctx_file->hook);
+	}
+
+	// Close the CFILE, except for stdio streams, which are closed later
+	if (cfile != ctx->cout && cfile != ctx->cerr) {
+		if (cfclose(cfile) != 0) {
+			ret = -1;
+			error = errno;
+		}
+	}
+
+	if (silent) {
+		ret = 0;
+	}
+
+	if (ret != 0 && ctx->cerr) {
+		if (ctx_file->path) {
+			bfs_error(ctx, "%pq: %s.\n", ctx_file->path, xstrerror(error));
+		} else if (cfile == ctx->cout) {
+			bfs_error(ctx, "(standard output): %s.\n", xstrerror(error));
+		}
+	}
+
+	free(ctx_file);
 	return ret;
 }
 
@@ -229,33 +269,14 @@ int bfs_ctx_free(struct bfs_ctx *ctx) {
 
 		for_trie (leaf, &ctx->files) {
 			struct bfs_ctx_file *ctx_file = leaf->value;
-
-			if (ctx_file->error) {
-				// An error was previously reported during bfs_ctx_flush()
-				ret = -1;
-			}
-
 			if (bfs_ctx_fclose(ctx, ctx_file) != 0) {
-				if (cerr) {
-					bfs_error(ctx, "%pq: %m.\n", ctx_file->path);
-				}
 				ret = -1;
 			}
-
-			free(ctx_file);
 		}
 		trie_destroy(&ctx->files);
 
-		if (cout && bfs_ctx_fflush(cout) != 0) {
-			if (cerr) {
-				bfs_error(ctx, "(standard output): %m.\n");
-			}
-			ret = -1;
-		}
-
 		cfclose(cout);
 		cfclose(cerr);
-
 		free_colors(ctx->colors);
 
 		for_slist (struct bfs_expr, expr, &ctx->expr_list, freelist) {