Initial commit.

src/coloredstderr.c

@@ -0,0 +1,305 @@
+/*
+ * Hook output functions (like printf(3)) with LD_PRELOAD to color stderr (or
+ * other file descriptors).
+ *
+ * Copyright (C) 2013  Simon Ruderich
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+#include "config.h"
+
+/* Must be loaded before the following headers. */
+#include "ldpreload.h"
+
+/* FIXME: use correct declaration for fcntl() */
+#define fcntl fcntl_ignore
+
+#include <errno.h>
+#include <fcntl.h>
+#include <stdarg.h>
+#include <stdio.h>
+#include <string.h>
+#include <unistd.h>
+
+#undef fcntl
+/* Conflicting declaration in glibc. */
+#undef fwrite_unlocked
+
+
+/* Used by various functions, including debug(). */
+static ssize_t (*real_write)(int, const void *, size_t);
+static int (*real_close)(int);
+static size_t (*real_fwrite)(const void *, size_t, size_t, FILE *);
+
+
+#include "constants.h"
+#ifdef DEBUG
+# include "debug.h"
+#endif
+
+#include "macros.h"
+#include "trackfds.h"
+
+
+
+/* Should the "action" handler be invoked for this file descriptor? */
+static int check_handle_fd(int fd) {
+    /* Never touch anything not going to a terminal. */
+    if (!isatty(fd)) {
+        return 0;
+    }
+
+    /* Load state from environment. Only necessary once per process. */
+    if (!tracked_fds) {
+        init_from_environment();
+    }
+
+    if (tracked_fds_count == 0) {
+        return 0;
+    }
+    return tracked_fds_find(fd);
+}
+
+static void dup_fd(int oldfd, int newfd) {
+#ifdef DEBUG
+    debug("%d -> %d\t\t\t[%d]\n", oldfd, newfd, getpid());
+#endif
+
+    if (!tracked_fds) {
+        init_from_environment();
+    }
+    if (tracked_fds_count == 0) {
+        return;
+    }
+
+    /* We are already tracking this file descriptor, add newfd to the list as
+     * it will reference the same descriptor. */
+    if (tracked_fds_find(oldfd)) {
+        if (!tracked_fds_find(newfd)) {
+            tracked_fds_add(newfd);
+            update_environment();
+        }
+    /* We are not tracking this file descriptor, remove newfd from the list
+     * (if present). */
+    } else {
+        if (tracked_fds_remove(newfd)) {
+            update_environment();
+        }
+    }
+}
+
+static void close_fd(int fd) {
+#ifdef DEBUG
+    debug("%d -> .\t\t\t[%d]\n", fd, getpid());
+#endif
+
+    if (!tracked_fds) {
+        init_from_environment();
+    }
+
+    if (tracked_fds_count == 0) {
+        return;
+    }
+    tracked_fds_remove(fd);
+}
+
+
+/* "Action" handlers called when a file descriptor is matched. */
+
+static char *pre_string;
+static size_t pre_string_size;
+static char *post_string;
+static size_t post_string_size;
+
+/* Load alternative pre/post strings from the environment if available, fall
+ * back to default values. */
+inline static void init_pre_post_string() {
+    pre_string = getenv(ENV_NAME_PRE_STRING);
+    if (!pre_string) {
+        pre_string = DEFAULT_PRE_STRING;
+    }
+    pre_string_size = strlen(pre_string);
+
+    post_string = getenv(ENV_NAME_POST_STRING);
+    if (!post_string) {
+        post_string = DEFAULT_POST_STRING;
+    }
+    post_string_size = strlen(post_string);
+}
+
+static void handle_fd_pre(int fd, int action) {
+    (void)action;
+
+    if (!pre_string || !post_string) {
+        init_pre_post_string();
+    }
+
+    DLSYM_FUNCTION(real_write, "write");
+    real_write(fd, pre_string, pre_string_size);
+}
+static void handle_fd_post(int fd, int action) {
+    (void)action;
+
+    /* write() already loaded above in handle_fd_pre(). */
+    real_write(fd, post_string, post_string_size);
+}
+
+static void handle_file_pre(FILE *stream, int action) {
+    (void)action;
+
+    if (!pre_string || !post_string) {
+        init_pre_post_string();
+    }
+
+    DLSYM_FUNCTION(real_fwrite, "fwrite");
+    real_fwrite(pre_string, pre_string_size, 1, stream);
+}
+static void handle_file_post(FILE *stream, int action) {
+    (void)action;
+
+    /* fwrite() already loaded above in handle_file_pre(). */
+    real_fwrite(post_string, post_string_size, 1, stream);
+}
+
+
+
+/* Hook all important output functions to manipulate their output. */
+
+HOOK_FD3(ssize_t, write, fd,
+         int, fd, const void *, buf, size_t, count)
+HOOK_FILE4(size_t, fwrite, stream,
+           const void *, ptr, size_t, size, size_t, nmemb, FILE *, stream)
+
+/* puts(3) */
+HOOK_FILE2(int, fputs, stream,
+           const char *, s, FILE *, stream)
+HOOK_FILE2(int, fputc, stream,
+           int, c, FILE *, stream)
+HOOK_FILE2(int, putc, stream,
+           int, c, FILE *, stream)
+HOOK_FILE1(int, putchar, stdout,
+           int, c)
+HOOK_FILE1(int, puts, stdout,
+           const char *, s)
+
+/* printf(3), excluding all s*() and vs*() functions (no output) */
+HOOK_VAR_FILE1(int, printf, stdout, vprintf,
+               const char *, format)
+HOOK_VAR_FILE2(int, fprintf, stream, vfprintf,
+               FILE *, stream, const char *, format)
+HOOK_FILE2(int, vprintf, stdout,
+           const char *, format, va_list, ap)
+HOOK_FILE3(int, vfprintf, stream,
+           FILE *, stream, const char *, format, va_list, ap)
+/* Hardening functions (-D_FORTIFY_SOURCE=2). */
+HOOK_VAR_FILE2(int, __printf_chk, stdout, __vprintf_chk,
+               int, flag, const char *, format)
+HOOK_VAR_FILE3(int, __fprintf_chk, fp, __vfprintf_chk,
+               FILE *, fp, int, flag, const char *, format)
+HOOK_FILE3(int, __vprintf_chk, stdout,
+           int, flag, const char *, format, va_list, ap)
+HOOK_FILE4(int, __vfprintf_chk, stream,
+           FILE *, stream, int, flag, const char *, format, va_list, ap)
+
+/* unlocked_stdio(3), only functions from above are hooked */
+HOOK_FILE4(size_t, fwrite_unlocked, stream,
+           const void *, ptr, size_t, size, size_t, nmemb, FILE *, stream)
+HOOK_FILE2(int, fputs_unlocked, stream,
+           const char *, s, FILE *, stream)
+HOOK_FILE2(int, fputc_unlocked, stream,
+           int, c, FILE *, stream)
+HOOK_FILE2(int, putc_unlocked, stream,
+           int, c, FILE *, stream)
+HOOK_FILE1(int, putchar_unlocked, stdout,
+           int, c)
+HOOK_FILE1(int, puts_unlocked, stdout,
+           const char *, s)
+
+
+/* Hook functions which duplicate file descriptors to track them. */
+
+static int (*real_dup)(int);
+static int (*real_dup2)(int, int);
+static int (*real_dup3)(int, int, int);
+int dup(int oldfd) {
+    int newfd;
+
+    DLSYM_FUNCTION(real_dup, "dup");
+
+    newfd = real_dup(oldfd);
+    if (newfd != -1) {
+        int saved_errno = errno;
+        dup_fd(oldfd, newfd);
+        errno = saved_errno;
+    }
+
+    return newfd;
+}
+int dup2(int oldfd, int newfd) {
+    DLSYM_FUNCTION(real_dup2, "dup2");
+
+    newfd = real_dup2(oldfd, newfd);
+    if (newfd != -1) {
+        int saved_errno = errno;
+        dup_fd(oldfd, newfd);
+        errno = saved_errno;
+    }
+
+    return newfd;
+}
+int dup3(int oldfd, int newfd, int flags) {
+    DLSYM_FUNCTION(real_dup3, "dup3");
+
+    newfd = real_dup3(oldfd, newfd, flags);
+    if (newfd != -1) {
+        int saved_errno = errno;
+        dup_fd(oldfd, newfd);
+        errno = saved_errno;
+    }
+
+    return newfd;
+}
+
+static int (*real_fcntl)(int, int, int);
+int fcntl(int fd, int cmd, int arg) {
+    int result;
+
+    DLSYM_FUNCTION(real_fcntl, "fcntl");
+
+    result = real_fcntl(fd, cmd, arg);
+    /* We only care about duping fds. */
+    if (cmd == F_DUPFD) {
+        int saved_errno = errno;
+        dup_fd(fd, result);
+        errno = saved_errno;
+    }
+
+    return result;
+}
+
+static int (*real_close)(int);
+int close(int fd) {
+    DLSYM_FUNCTION(real_close, "close");
+
+    close_fd(fd);
+    return real_close(fd);
+}
+static int (*real_fclose)(FILE *);
+int fclose(FILE *fp) {
+    DLSYM_FUNCTION(real_fclose, "fclose");
+
+    close_fd(fileno(fp));
+    return real_fclose(fp);
+}