From: Kirill Isakov <bootctl@gmail.com>
Date: Tue, 26 Apr 2022 10:57:18 +0000 (+0600)
Subject: Add colors to tincd logger
X-Git-Url: https://git.tinc-vpn.org/git/browse?a=commitdiff_plain;h=48dd371b2ec69f02870aa5e2a67fc8adb4617ff1;p=tinc

Add colors to tincd logger
---

diff --git a/src/connection.h b/src/connection.h
index 6489ef4f..23bde776 100644
--- a/src/connection.h
+++ b/src/connection.h
@@ -49,6 +49,7 @@ typedef union connection_status_t {
 		bool control: 1;                /* 1 if this is a control connection */
 		bool pcap: 1;                   /* 1 if this is a control connection requesting packet capture */
 		bool log: 1;                    /* 1 if this is a control connection requesting log dump */
+		bool log_color: 1;              /* 1 if this connection supports ANSI escape codes */
 		bool invitation: 1;             /* 1 if this is an invitation */
 		bool invitation_used: 1;        /* 1 if the invitation has been consumed */
 		bool tarpit: 1;                 /* 1 if the connection should be added to the tarpit */
diff --git a/src/console.c b/src/console.c
new file mode 100644
index 00000000..503b3e65
--- /dev/null
+++ b/src/console.c
@@ -0,0 +1,28 @@
+#include "system.h"
+
+#include "console.h"
+
+bool use_ansi_escapes(FILE *out) {
+	bool is_tty = isatty(fileno(out));
+
+#ifndef HAVE_WINDOWS
+	const char *term = getenv("TERM");
+	return is_tty && term && strcmp(term, "dumb");
+#else
+	HANDLE console;
+
+	if(out == stdout) {
+		console = GetStdHandle(STD_OUTPUT_HANDLE);
+	} else if(out == stderr) {
+		console = GetStdHandle(STD_ERROR_HANDLE);
+	} else {
+		return false;
+	}
+
+	DWORD mode = 0;
+	return is_tty &&
+	       console != INVALID_HANDLE_VALUE &&
+	       GetConsoleMode(console, &mode) &&
+	       SetConsoleMode(console, mode | ENABLE_VIRTUAL_TERMINAL_PROCESSING);
+#endif
+}
diff --git a/src/console.h b/src/console.h
new file mode 100644
index 00000000..b11227a7
--- /dev/null
+++ b/src/console.h
@@ -0,0 +1,9 @@
+#ifndef TINC_CONSOLE_H
+#define TINC_CONSOLE_H
+
+#include "system.h"
+
+// true if stderr supports ANSI escape sequences.
+extern bool use_ansi_escapes(FILE *out);
+
+#endif // TINC_CONSOLE_H
diff --git a/src/control.c b/src/control.c
index 6e8580a4..550384f2 100644
--- a/src/control.c
+++ b/src/control.c
@@ -131,10 +131,11 @@ bool control_h(connection_t *c, const char *request) {
 		return true;
 
 	case REQ_LOG: {
-		int level = 0;
-		sscanf(request, "%*d %*d %d", &level);
+		int level = 0, colorize = 0;
+		sscanf(request, "%*d %*d %d %d", &level, &colorize);
 		c->log_level = CLAMP(level, DEBUG_UNSET, DEBUG_SCARY_THINGS);
 		c->status.log = true;
+		c->status.log_color = colorize;
 		logcontrol = true;
 		return true;
 	}
diff --git a/src/dropin.h b/src/dropin.h
index f161c13d..f020f215 100644
--- a/src/dropin.h
+++ b/src/dropin.h
@@ -78,6 +78,8 @@ extern int gettimeofday(struct timeval *, void *);
 #define PATH_MAX MAX_PATH
 #define strcasecmp _stricmp
 #define strncasecmp _strnicmp
+#define isatty _isatty
+#define fileno _fileno
 #define close CloseHandle
 #define __const const
 
diff --git a/src/logger.c b/src/logger.c
index 026a1202..3e630f23 100644
--- a/src/logger.c
+++ b/src/logger.c
@@ -28,7 +28,7 @@
 #include "control_common.h"
 #include "process.h"
 #include "sptps.h"
-#include "compression.h"
+#include "console.h"
 
 debug_t debug_level = DEBUG_NOTHING;
 static logmode_t logmode = LOGMODE_STDERR;
@@ -38,15 +38,108 @@ static FILE *logfile = NULL;
 static HANDLE loghandle = NULL;
 #endif
 static const char *logident = NULL;
+static bool colorize_stderr = false;
 bool logcontrol = false; // controlled by REQ_LOG <level>
 int umbilical = 0;
+bool umbilical_colorize = false;
+
+#define SGR(s) ("\x1b[" s "m")
+
+typedef enum color_t {
+	RED, GREEN, YELLOW, BLUE, MAGENTA, CYAN, WHITE, GRAY,
+	RESET, // not really a color
+} color_t;
+
+typedef struct priority_disp_t {
+	const char *name;
+	color_t color;
+} priority_disp_t;
+
+static const priority_disp_t priorities[] = {
+	[LOG_EMERG]   = {"EMERGENCY", MAGENTA},
+	[LOG_ALERT]   = {"ALERT",     MAGENTA},
+	[LOG_CRIT]    = {"CRITICAL",  MAGENTA},
+	[LOG_ERR]     = {"ERROR",     RED},
+	[LOG_WARNING] = {"WARNING",   YELLOW},
+	[LOG_NOTICE]  = {"NOTICE",    CYAN},
+	[LOG_INFO]    = {"INFO",      GREEN},
+	[LOG_DEBUG]   = {"DEBUG",     BLUE},
+};
+
+static const char *ansi_codes[] = {
+	[RED]     = SGR("31;1"),
+	[GREEN]   = SGR("32;1"),
+	[YELLOW]  = SGR("33;1"),
+	[BLUE]    = SGR("34;1"),
+	[MAGENTA] = SGR("35;1"),
+	[CYAN]    = SGR("36;1"),
+	[WHITE]   = SGR("37;1"),
+	[GRAY]    = SGR("90"),
+	[RESET]   = SGR("0"),
+};
+
+static priority_disp_t get_priority(int priority) {
+	static const priority_disp_t unknown = {"UNKNOWN", WHITE};
+
+	if(priority >= LOG_EMERG && priority <= LOG_DEBUG) {
+		return priorities[priority];
+	} else {
+		return unknown;
+	}
+}
+
+// Formats current time to the second.
+// Reuses result so repeated calls within the same second are more efficient.
+static const char *current_time_str(void) {
+	static char timestr[sizeof("2000-12-31 12:34:56")] = "";
+	static time_t last_time = 0;
+
+	if(!now.tv_sec) {
+		gettimeofday(&now, NULL);
+	}
+
+	time_t now_sec = now.tv_sec;
+
+	if(!*timestr || now_sec != last_time) {
+		last_time = now_sec;
+		strftime(timestr, sizeof(timestr), "%Y-%m-%d %H:%M:%S", localtime(&now_sec));
+	}
+
+	return timestr;
+}
+
+// Format log entry with time, log level, and (possibly) colors, remembering if it was colorized.
+// Returns true if buffer has been changed.
+static bool format_pretty(char *buf, size_t buflen, int prio, const char *message, bool colorize, bool *colorized) {
+	// If we already wrote to buffer, and its colorization matches, we're done here
+	if(*buf && colorize == *colorized) {
+		return false;
+	}
+
+	// Otherwise, remember current colorization for future comparisons
+	*colorized = colorize;
+
+	priority_disp_t priority = get_priority(prio);
+	const char *color = "", *reset = "", *timecol = "";
+
+	if(colorize) {
+		color = ansi_codes[priority.color];
+		reset = ansi_codes[RESET];
+		timecol = ansi_codes[GRAY];
+	}
+
+	const char *timestr = current_time_str();
+	snprintf(buf, buflen, "%s%s %s%-7s%s %s", timecol, timestr, color, priority.name, reset, message);
+	return true;
+}
 
 static bool should_log(debug_t level) {
 	return (level <= debug_level && logmode != LOGMODE_NULL) || logcontrol;
 }
 
 static void real_logger(debug_t level, int priority, const char *message) {
-	char timestr[32] = "";
+	char pretty[1024] = "";
+	bool pretty_colorized = false;
 	static bool suppress = false;
 
 	if(suppress) {
@@ -56,20 +149,17 @@ static void real_logger(debug_t level, int priority, const char *message) {
 	if(level <= debug_level) {
 		switch(logmode) {
 		case LOGMODE_STDERR:
-			fprintf(stderr, "%s\n", message);
+			format_pretty(pretty, sizeof(pretty), priority, message, colorize_stderr, &pretty_colorized);
+			fprintf(stderr, "%s\n", pretty);
 			fflush(stderr);
 			break;
 
-		case LOGMODE_FILE:
-			if(!now.tv_sec) {
-				gettimeofday(&now, NULL);
-			}
-
-			time_t now_sec = now.tv_sec;
-			strftime(timestr, sizeof(timestr), "%Y-%m-%d %H:%M:%S", localtime(&now_sec));
+		case LOGMODE_FILE: {
+			const char *timestr = current_time_str();
 			fprintf(logfile, "%s %s[%ld]: %s\n", timestr, logident, (long)logpid, message);
 			fflush(logfile);
 			break;
+		}
 
 		case LOGMODE_SYSLOG:
 #ifdef HAVE_WINDOWS
@@ -91,9 +181,9 @@ static void real_logger(debug_t level, int priority, const char *message) {
 		}
 
 		if(umbilical && do_detach) {
-			size_t len = strlen(message);
+			format_pretty(pretty, sizeof(pretty), priority, message, umbilical_colorize, &pretty_colorized);
 
-			if(write(umbilical, message, len) != (ssize_t)len || write(umbilical, "\n", 1) != 1) {
+			if(write(umbilical, pretty, strlen(pretty)) == -1 || write(umbilical, "\n", 1) == -1) {
 				// Other end broken, nothing we can do about it.
 			}
 		}
@@ -103,6 +193,8 @@ static void real_logger(debug_t level, int priority, const char *message) {
 		suppress = true;
 		logcontrol = false;
 
+		size_t msglen = strlen(pretty);
+
 		for list_each(connection_t, c, &connection_list) {
 			if(!c->status.log) {
 				continue;
@@ -114,10 +206,12 @@ static void real_logger(debug_t level, int priority, const char *message) {
 				continue;
 			}
 
-			size_t len = strlen(message);
+			if(format_pretty(pretty, sizeof(pretty), priority, message, c->status.log_color, &pretty_colorized)) {
+				msglen = strlen(pretty);
+			}
 
-			if(send_request(c, "%d %d %lu", CONTROL, REQ_LOG, (unsigned long)len)) {
-				send_meta(c, message, len);
+			if(send_request(c, "%d %d %lu", CONTROL, REQ_LOG, (unsigned long)msglen)) {
+				send_meta(c, pretty, msglen);
 			}
 		}
 
@@ -181,6 +275,7 @@ void openlogger(const char *ident, logmode_t mode) {
 	switch(mode) {
 	case LOGMODE_STDERR:
 		logpid = getpid();
+		colorize_stderr = use_ansi_escapes(stderr);
 		break;
 
 	case LOGMODE_FILE:
diff --git a/src/logger.h b/src/logger.h
index 038f23ea..52f8283d 100644
--- a/src/logger.h
+++ b/src/logger.h
@@ -73,6 +73,7 @@ enum {
 extern debug_t debug_level;
 extern bool logcontrol;
 extern int umbilical;
+extern bool umbilical_colorize;
 extern void openlogger(const char *ident, logmode_t mode);
 extern void reopenlogger(void);
 extern void logger(debug_t level, int priority, const char *format, ...) ATTR_FORMAT(printf, 3, 4);
diff --git a/src/meson.build b/src/meson.build
index c780ba11..e8c3e53f 100644
--- a/src/meson.build
+++ b/src/meson.build
@@ -106,6 +106,7 @@ subdir('chacha-poly1305')
 
 src_lib_common = [
   'conf.c',
+  'console.c',
   'dropin.c',
   'keys.c',
   'list.c',
diff --git a/src/tincctl.c b/src/tincctl.c
index 9b39f2ce..a882a668 100644
--- a/src/tincctl.c
+++ b/src/tincctl.c
@@ -42,6 +42,7 @@
 #include "keys.h"
 #include "random.h"
 #include "pidfile.h"
+#include "console.h"
 
 #ifndef MSG_NOSIGNAL
 #define MSG_NOSIGNAL 0
@@ -612,8 +613,9 @@ static void pcap(int fd, FILE *out, uint32_t snaplen) {
 	}
 }
 
-static void log_control(int fd, FILE *out, int level) {
-	sendline(fd, "%d %d %d", CONTROL, REQ_LOG, level);
+static void log_control(int fd, FILE *out, int level, bool use_color) {
+	sendline(fd, "%d %d %d %d", CONTROL, REQ_LOG, level, use_color);
+
 	char data[1024];
 	char line[32];
 
@@ -955,7 +957,7 @@ static int cmd_start(int argc, char *argv[]) {
 	if(!pid) {
 		close(pfd[0]);
 		char buf[100];
-		snprintf(buf, sizeof(buf), "%d", pfd[1]);
+		snprintf(buf, sizeof(buf), "%d %d", pfd[1], use_ansi_escapes(stderr));
 		setenv("TINC_UMBILICAL", buf, true);
 		exit(execvp(c, nargv));
 	} else {
@@ -1518,7 +1520,8 @@ static int cmd_log(int argc, char *argv[]) {
 	signal(SIGINT, sigint_handler);
 #endif
 
-	log_control(fd, stdout, argc > 1 ? atoi(argv[1]) : DEBUG_UNSET);
+	bool use_color = use_ansi_escapes(stdout);
+	log_control(fd, stdout, argc > 1 ? atoi(argv[1]) : DEBUG_UNSET, use_color);
 
 #ifdef SIGINT
 	signal(SIGINT, SIG_DFL);
diff --git a/src/tincd.c b/src/tincd.c
index 947e7b3e..231e01fa 100644
--- a/src/tincd.c
+++ b/src/tincd.c
@@ -458,7 +458,9 @@ int main(int argc, char **argv) {
 	char *umbstr = getenv("TINC_UMBILICAL");
 
 	if(umbstr) {
-		umbilical = atoi(umbstr);
+		int colorize = 0;
+		sscanf(umbstr, "%d %d", &umbilical, &colorize);
+		umbilical_colorize = colorize;
 
 		if(fcntl(umbilical, F_GETFL) < 0) {
 			umbilical = 0;