From e2b226fe991f43cc2a5596dfab08b7bffb96ef17 Mon Sep 17 00:00:00 2001
From: Luca Ferrari <fluca1978@gmail.com>
Date: Tue, 28 Jun 2022 04:34:04 -0400
Subject: [PATCH] [#215] Process title length

This is a brute force approach to set the process title with the
appropriate length.
The idea is to compute the size of the process title as it has been
requested, and set it to the computed size.

This works on Rocky Linux release 8.6 (Green Obsidian).

It could break some aggressive security-enforced environments.

This commit introduces the function
`pgagroal_set_connection_proc_title()` that builds a process title
with the form:

      user@host:port/database

using in turn the `pgagroal_set_proc_title` function.
The `pgagroal_set_connection_proc_title` function is used only by the
`worker` process stuff and required to identify the primary server (to
get the hostname and port).

Close #215.
---
 src/include/utils.h      | 14 +++++++++
 src/libpgagroal/utils.c  | 67 ++++++++++++++++++++++++----------------
 src/libpgagroal/worker.c |  2 +-
 3 files changed, 55 insertions(+), 28 deletions(-)

diff --git a/src/include/utils.h b/src/include/utils.h
index b48c451e..ad0b0660 100644
--- a/src/include/utils.h
+++ b/src/include/utils.h
@@ -254,6 +254,20 @@ pgagroal_base64_decode(char* encoded, size_t encoded_length, char** raw, int* ra
 void
 pgagroal_set_proc_title(int argc, char** argv, char* s1, char* s2);
 
+/**
+ * Sets the process title for a given connection.
+ *
+ * Uses `pgagroal_set_proc_title` to build an information string
+ * with the form
+ *    user@host:port/database
+ *
+ * @param argc the number of arguments
+ * @param argv command line arguments
+ * @param connection the struct connection pointer for the established connection.
+ */
+void
+pgagroal_set_connection_proc_title(int argc, char** argv, struct connection* connection);
+
 #ifdef DEBUG
 
 /**
diff --git a/src/libpgagroal/utils.c b/src/libpgagroal/utils.c
index 6d2b1ecd..d369c53d 100644
--- a/src/libpgagroal/utils.c
+++ b/src/libpgagroal/utils.c
@@ -30,6 +30,7 @@
 #include <pgagroal.h>
 #include <logging.h>
 #include <utils.h>
+#include <server.h>
 
 /* system */
 #include <ev.h>
@@ -55,7 +56,7 @@
 extern char** environ;
 #ifdef HAVE_LINUX
 static bool env_changed = false;
-static int max_process_title_size = -1;
+static int max_process_title_size = 0;
 #endif
 
 int32_t
@@ -754,50 +755,62 @@ pgagroal_set_proc_title(int argc, char** argv, char* s1, char* s2)
       env_changed = true;
    }
 
-   if (max_process_title_size == -1)
+   // compute how long was the command line
+   // when the application was started
+   if (max_process_title_size == 0)
    {
-      int m = 0;
-
       for (int i = 0; i < argc; i++)
       {
-         m += strlen(argv[i]);
-         m += 1;
+         max_process_title_size += strlen(argv[i]) + 1;
       }
-
-      max_process_title_size = m;
-      memset(*argv, 0, max_process_title_size);
    }
 
+   // compose the new title
    memset(&title, 0, sizeof(title));
+   snprintf(title, sizeof(title) - 1, "pgagroal: %s%s%s",
+            s1 != NULL ? s1 : "",
+            s1 != NULL && s2 != NULL ? "/" : "",
+            s2 != NULL ? s2 : "");
 
-   if (s1 != NULL && s2 != NULL)
-   {
-      snprintf(title, sizeof(title) - 1, "pgagroal: %s/%s", s1, s2);
-   }
-   else
-   {
-      snprintf(title, sizeof(title) - 1, "pgagroal: %s", s1);
-   }
+   size = strlen(title) + 1;
 
-   size = MIN(max_process_title_size, sizeof(title));
-   size = MIN(size, strlen(title) + 1);
+   // nuke the command line info
+   memset(*argv, 0, max_process_title_size);
 
    memcpy(*argv, title, size);
    memset(*argv + size, 0, 1);
 
+   // keep track of how long is now the title
+   max_process_title_size = size;
+
 #else
-   if (s1 != NULL && s2 != NULL)
-   {
-      setproctitle("-pgagroal: %s/%s", s1, s2);
-   }
-   else
-   {
-      setproctitle("-pgagroal: %s", s1);
-   }
+   setproctitle("-pgagroal: %s%s%s",
+                s1 != NULL ? s1 : "",
+                s1 != NULL && s2 != NULL ? "/" : "",
+                s2 != NULL ? s2 : "");
 
 #endif
 }
 
+void
+pgagroal_set_connection_proc_title(int argc, char** argv, struct connection* connection)
+{
+   struct configuration* config;
+   int primary;
+   char info[MISC_LENGTH];
+
+   config = (struct configuration*)shmem;
+
+   pgagroal_get_primary(&primary);
+
+   snprintf(info, MISC_LENGTH, "%s@%s:%d",
+            connection->username,
+            config->servers[primary].host,
+            config->servers[primary].port);
+
+   pgagroal_set_proc_title(argc, argv, info, connection->database);
+}
+
 #ifdef DEBUG
 
 int
diff --git a/src/libpgagroal/worker.c b/src/libpgagroal/worker.c
index ab2b2e60..614d9d66 100644
--- a/src/libpgagroal/worker.c
+++ b/src/libpgagroal/worker.c
@@ -106,7 +106,7 @@ pgagroal_worker(int client_fd, char* address, char** argv)
       pgagroal_prometheus_client_active_add();
 
       pgagroal_pool_status();
-      pgagroal_set_proc_title(1, argv, config->connections[slot].username, config->connections[slot].database);
+      pgagroal_set_connection_proc_title(1, argv, &config->connections[slot]);
 
       if (config->pipeline == PIPELINE_PERFORMANCE)
       {