ptyfwd: optionally prefix window title with colored dot

in uid0/systemd-run/nspawn we already set a window title with a colorful
unicode dot indicating the changed privileges/execution context. This typically
gets overriden by the shell inside the environment however.

Let's tweak this a bit: when we see the window title OSC ANSI sequence
passing through, let's patch in the unicode dot as a prefix to the
title.

This is super pretty, since it makes sure root sessions via 0ad are
really easily recognizable as such, because the window title carries an
🔴 red dot as prefix then.

src/nspawn/nspawn.c

@@ -4443,6 +4443,9 @@ static void set_window_title(PTYForward *f) {
                 (void) pty_forward_set_titlef(f, "%sContainer %s on %s", strempty(dot), arg_machine, hn);
         else
                 (void) pty_forward_set_titlef(f, "%sContainer %s", strempty(dot), arg_machine);
+
+        if (dot)
+                (void) pty_forward_set_title_prefix(f, dot);
 }
 
 static int merge_settings(Settings *settings, const char *path) {

src/run/run.c

@@ -1610,6 +1610,8 @@ static void set_window_title(PTYForward *f) {
                 (void) pty_forward_set_titlef(f, "%s%s on %s", strempty(dot), cl, arg_host ?: hn);
         else
                 (void) pty_forward_set_titlef(f, "%s%s", strempty(dot), cl);
+
+        (void) pty_forward_set_title_prefix(f, dot);
 }
 
 static int start_transient_service(sd_bus *bus) {

src/shared/ptyfwd.c

@@ -33,6 +33,7 @@ typedef enum AnsiColorState  {
         ANSI_COLOR_STATE_TEXT,
         ANSI_COLOR_STATE_ESC,
         ANSI_COLOR_STATE_CSI_SEQUENCE,
+        ANSI_COLOR_STATE_OSC_SEQUENCE,
         ANSI_COLOR_STATE_NEWLINE,
         ANSI_COLOR_STATE_CARRIAGE_RETURN,
         _ANSI_COLOR_STATE_MAX,
@@ -92,8 +93,10 @@ struct PTYForward {
         char *background_color;
         AnsiColorState ansi_color_state;
         char *csi_sequence;
+        char *osc_sequence;
 
-        char *title;
+        char *title;           /* Window title to show by default */
+        char *title_prefix;    /* If terminal client overrides window title, prefix this string */
 };
 
 #define ESCAPE_USEC (1*USEC_PER_SEC)
@@ -145,6 +148,7 @@ static void pty_forward_disconnect(PTYForward *f) {
         f->in_buffer_full = 0;
 
         f->csi_sequence = mfree(f->csi_sequence);
+        f->osc_sequence = mfree(f->osc_sequence);
         f->ansi_color_state = _ANSI_COLOR_STATE_INVALID;
 }
 
@@ -367,7 +371,9 @@ static int is_csi_background_reset_sequence(const char *seq) {
 
 static int insert_background_fix(PTYForward *f, size_t offset) {
         assert(f);
-        assert(f->background_color);
+
+        if (!f->background_color)
+                return 0;
 
         if (!is_csi_background_reset_sequence(strempty(f->csi_sequence)))
                 return 0;
@@ -380,13 +386,33 @@ static int insert_background_fix(PTYForward *f, size_t offset) {
         return insert_string(f, offset, s);
 }
 
+static int insert_window_title_fix(PTYForward *f, size_t offset) {
+        assert(f);
+
+        if (!f->title_prefix)
+                return 0;
+
+        if (!f->osc_sequence)
+                return 0;
+
+        const char *t = startswith(f->osc_sequence, "0;"); /* Set window title OSC sequence*/
+        if (!t)
+                return 0;
+
+        _cleanup_free_ char *joined = strjoin("\x1b]0;", f->title_prefix, t, "\a");
+        if (!joined)
+                return -ENOMEM;
+
+        return insert_string(f, offset, joined);
+}
+
 static int pty_forward_ansi_process(PTYForward *f, size_t offset) {
         int r;
 
         assert(f);
         assert(offset <= f->out_buffer_full);
 
-        if (!f->background_color)
+        if (!f->background_color && !f->title_prefix)
                 return 0;
 
         if (FLAGS_SET(f->flags, PTY_FORWARD_DUMB_TERMINAL))
@@ -427,6 +453,9 @@ static int pty_forward_ansi_process(PTYForward *f, size_t offset) {
                         if (c == '[') {
                                 f->ansi_color_state = ANSI_COLOR_STATE_CSI_SEQUENCE;
                                 continue;
+                        } else if (c == ']') {
+                                f->ansi_color_state = ANSI_COLOR_STATE_OSC_SEQUENCE;
+                                continue;
                         }
 
                         break;
@@ -464,6 +493,33 @@ static int pty_forward_ansi_process(PTYForward *f, size_t offset) {
                         continue;
                 }
 
+                case ANSI_COLOR_STATE_OSC_SEQUENCE: {
+
+                        if ((uint8_t) c >= ' ') {
+                                if (strlen_ptr(f->osc_sequence) >= 64) {
+                                        /* Safety check: lets not accept unbounded OSC sequences */
+                                        f->osc_sequence = mfree(f->osc_sequence);
+                                        break;
+                                } else if (!strextend(&f->osc_sequence, CHAR_TO_STR(c)))
+                                        return -ENOMEM;
+                        } else {
+                                /* Otherwise, the OSC sequence is over */
+
+                                if (c == '\x07') {
+                                        r = insert_window_title_fix(f, i+1);
+                                        if (r < 0)
+                                                return r;
+
+                                        i += r;
+                                }
+
+                                f->osc_sequence = mfree(f->osc_sequence);
+                                f->ansi_color_state = ANSI_COLOR_STATE_TEXT;
+                        }
+
+                        continue;
+                }
+
                 default:
                         assert_not_reached();
                 }
@@ -888,6 +944,7 @@ PTYForward *pty_forward_free(PTYForward *f) {
         pty_forward_disconnect(f);
         free(f->background_color);
         free(f->title);
+        free(f->title_prefix);
         return mfree(f);
 }
 
@@ -1060,3 +1117,9 @@ int pty_forward_set_titlef(PTYForward *f, const char *format, ...) {
 
         return free_and_replace(f->title, title);
 }
+
+int pty_forward_set_title_prefix(PTYForward *f, const char *title_prefix) {
+        assert(f);
+
+        return free_and_strdup(&f->title_prefix, title_prefix);
+}

src/shared/ptyfwd.h

@@ -44,7 +44,10 @@ int pty_forward_set_priority(PTYForward *f, int64_t priority);
 int pty_forward_set_width_height(PTYForward *f, unsigned width, unsigned height);
 
 int pty_forward_set_background_color(PTYForward *f, const char *color);
+
 int pty_forward_set_title(PTYForward *f, const char *title);
 int pty_forward_set_titlef(PTYForward *f, const char *format, ...) _printf_(2,3);
 
+int pty_forward_set_title_prefix(PTYForward *f, const char *prefix);
+
 DEFINE_TRIVIAL_CLEANUP_FUNC(PTYForward*, pty_forward_free);