Add application_name parameter to clientauth (#291)

ajrajkumar · web-flow · commit f6dcc39e07eb · 2025-03-18T11:24:23.000-04:00
diff --git a/docs/04_hooks.md b/docs/04_hooks.md
@@ -209,14 +209,15 @@ A `clientauth` hook function takes the following arguments and returns either `t
 
 clientauth_hook(port pgtle.clientauth_port_subset, status integer)
 
-* `port` (`pgtle.clientauth_port_subset`) - an object containing the following fields. These are a subset of the Port object that client authentication hook passes to internal C functions.
+* `port` (`pgtle.clientauth_port_subset`) - an object containing the following fields. These are a subset of the [Port](https://doxygen.postgresql.org/structPort.html) object that client authentication hook passes to internal C functions.
   * `noblock` (`bool`)
   * `remote_host` (`text`)
   * `remote_hostname` (`text`)
   * `remote_hostname_resolv` (`integer`)
   * `remote_hostname_errcode` (`integer`)
   * `database_name` (`text`)
   * `user_name` (`text`)
+  * `application_name` (`text`)
 * `status` (`integer`) - connection status code. This can be one of the following:
   * 0, representing successful connection
   * -1, representing a connection error
diff --git a/include/constants.h b/include/constants.h
@@ -36,6 +36,8 @@
 #define TLE_INPUT_FUNC_STR       "input"
 #define TLE_OUTPUT_FUNC_STR      "output"
 
+#define TLE_CLIENTAUTH_PORT_SUBSET_TYPE	"clientauth_port_subset"
+
 /*
  * TLE_BASE_TYPE_SIZE_LIMIT is the maximum allowed size of pg_tle type.
  *
diff --git a/pg_tle--1.4.1--1.5.0.sql b/pg_tle--1.4.1--1.5.0.sql
@@ -81,3 +81,6 @@ CREATE FUNCTION pgtle.available_extensions
 RETURNS SETOF RECORD
 AS 'MODULE_PATHNAME', 'pg_tle_available_extensions'
 LANGUAGE C STABLE STRICT;
+
+ALTER TYPE pgtle.clientauth_port_subset
+ADD ATTRIBUTE application_name text;
diff --git a/src/clientauth.c b/src/clientauth.c
@@ -50,6 +50,7 @@
 #include "commands/extension.h"
 #include "commands/user.h"
 #include "executor/spi.h"
+#include "funcapi.h"
 #include "libpq/auth.h"
 #include "nodes/pg_list.h"
 #include "postmaster/bgworker_internals.h"
@@ -96,51 +97,6 @@
  */
 #define CLIENT_AUTH_USER_ERROR_MAX_STRLEN 256
 
-static const char *clientauth_shmem_name = "pgtle_clientauth";
-static const char *clientauth_feature = "clientauth";
-static const char *clientauth_worker_name = "pg_tle_clientauth worker";
-
-/* Background worker main entry function */
-PGDLLEXPORT void clientauth_launcher_main(Datum arg);
-
-/* Set up our hooks */
-static ClientAuthentication_hook_type prev_clientauth_hook = NULL;
-static void clientauth_hook(Port *port, int status);
-
-static shmem_startup_hook_type prev_shmem_startup_hook = NULL;
-static void clientauth_shmem_startup(void);
-
-#if (PG_VERSION_NUM >= 150000)
-static shmem_request_hook_type prev_shmem_request_hook = NULL;
-static void clientauth_shmem_request(void);
-#endif
-
-/* Helper functions */
-static Size clientauth_shared_memsize(void);
-static void clientauth_sighup(SIGNAL_ARGS);
-
-void		clientauth_init(void);
-static bool can_allow_without_executing(void);
-static bool can_reject_without_executing(void);
-
-/* GUC that determines whether clientauth is enabled */
-static int	enable_clientauth_feature = FEATURE_OFF;
-
-/* GUC that determines which database SPI_exec runs against */
-static char *clientauth_database_name = "postgres";
-
-/* GUC that determines the number of background workers */
-static int	clientauth_num_parallel_workers = 1;
-
-/* GUC that determines users that clientauth feature skips */
-static char *clientauth_users_to_skip = "";
-
-/* GUC that determines databases that clientauth feature skips */
-static char *clientauth_databases_to_skip = "";
-
-/* Global flags */
-static bool clientauth_reload_config = false;
-
 /*
  * Fixed-length subset of Port, passed to user function. A corresponding SQL
  * base type is defined. Shared memory structs are required to be fixed-size,
@@ -163,6 +119,7 @@ typedef struct PortSubset
 
 	char		database_name[CLIENT_AUTH_PORT_SUBSET_MAX_STRLEN];
 	char		user_name[CLIENT_AUTH_PORT_SUBSET_MAX_STRLEN];
+	char		application_name[CLIENT_AUTH_PORT_SUBSET_MAX_STRLEN];
 }			PortSubset;
 
 /* Represents a pending connection */
@@ -233,6 +190,53 @@ typedef struct ClientAuthBgwShmemSharedState
 	ClientAuthStatusEntry requests[CLIENT_AUTH_MAX_PENDING_ENTRIES];
 }			ClientAuthBgwShmemSharedState;
 
+static const char *clientauth_shmem_name = "pgtle_clientauth";
+static const char *clientauth_feature = "clientauth";
+static const char *clientauth_worker_name = "pg_tle_clientauth worker";
+
+/* Background worker main entry function */
+PGDLLEXPORT void clientauth_launcher_main(Datum arg);
+
+/* Set up our hooks */
+static ClientAuthentication_hook_type prev_clientauth_hook = NULL;
+static void clientauth_hook(Port *port, int status);
+
+static shmem_startup_hook_type prev_shmem_startup_hook = NULL;
+static void clientauth_shmem_startup(void);
+
+#if (PG_VERSION_NUM >= 150000)
+static shmem_request_hook_type prev_shmem_request_hook = NULL;
+static void clientauth_shmem_request(void);
+#endif
+
+/* Helper functions */
+static Size clientauth_shared_memsize(void);
+static void clientauth_sighup(SIGNAL_ARGS);
+
+void		clientauth_init(void);
+static bool can_allow_without_executing(void);
+static bool can_reject_without_executing(void);
+
+static char *create_port_subset_str(PortSubset * port);
+
+/* GUC that determines whether clientauth is enabled */
+static int	enable_clientauth_feature = FEATURE_OFF;
+
+/* GUC that determines which database SPI_exec runs against */
+static char *clientauth_database_name = "postgres";
+
+/* GUC that determines the number of background workers */
+static int	clientauth_num_parallel_workers = 1;
+
+/* GUC that determines users that clientauth feature skips */
+static char *clientauth_users_to_skip = "";
+
+/* GUC that determines databases that clientauth feature skips */
+static char *clientauth_databases_to_skip = "";
+
+/* Global flags */
+static bool clientauth_reload_config = false;
+
 static ClientAuthBgwShmemSharedState * clientauth_ss = NULL;
 
 void		clientauth_launcher_run_user_functions(bool *error, char (*error_msg)[CLIENT_AUTH_USER_ERROR_MAX_STRLEN], PortSubset * port, int *status);
@@ -587,14 +591,7 @@ clientauth_launcher_run_user_functions(bool *error, char (*error_msg)[CLIENT_AUT
 						 func_name,
 						 quote_identifier(PG_TLE_NSPNAME));
 
-		port_subset_str = psprintf("(%d,%s,%s,%d,%d,%s,%s)",
-								   port->noblock,
-								   quote_identifier(port->remote_host),
-								   quote_identifier(port->remote_hostname),
-								   port->remote_hostname_resolv,
-								   port->remote_hostname_errcode,
-								   quote_identifier(port->database_name),
-								   quote_identifier(port->user_name));
+		port_subset_str = create_port_subset_str(port);
 
 		hookargs[0] = CStringGetTextDatum(port_subset_str);
 		hookargs[1] = Int32GetDatum(*status);
@@ -718,6 +715,10 @@ clientauth_hook(Port *port, int status)
 			 CLIENT_AUTH_PORT_SUBSET_MAX_STRLEN,
 			 "%s",
 			 port->user_name == NULL ? "" : port->user_name);
+	snprintf(clientauth_ss->requests[idx].port_info.application_name,
+			 CLIENT_AUTH_PORT_SUBSET_MAX_STRLEN,
+			 "%s",
+			 port->application_name == NULL ? "" : port->application_name);
 	clientauth_ss->requests[idx].port_info.noblock = port->noblock;
 	clientauth_ss->requests[idx].port_info.remote_hostname_resolv = port->remote_hostname_resolv;
 	clientauth_ss->requests[idx].port_info.remote_hostname_errcode = port->remote_hostname_errcode;
@@ -894,3 +895,49 @@ can_reject_without_executing()
 
 	return false;
 }
+
+/*
+ * Constructs a string representation of the PortSubset composite type and
+ * returns it in an allocated string.  The signature of the PortSubset type
+ * varies based on the pg_tle version, so check the signature first.
+ */
+static char *
+create_port_subset_str(PortSubset * port)
+{
+	TupleDesc	tupdesc =
+		RelationNameGetTupleDesc(PG_TLE_NSPNAME "."
+								 TLE_CLIENTAUTH_PORT_SUBSET_TYPE);
+	char	   *port_subset_str;
+
+	if (tupdesc->natts == 7)
+		port_subset_str = psprintf("(%d,%s,%s,%d,%d,%s,%s)",
+								   port->noblock,
+								   quote_identifier(port->remote_host),
+								   quote_identifier(port->remote_hostname),
+								   port->remote_hostname_resolv,
+								   port->remote_hostname_errcode,
+								   quote_identifier(port->database_name),
+								   quote_identifier(port->user_name));
+	else if (tupdesc->natts == 8)
+		port_subset_str = psprintf("(%d,%s,%s,%d,%d,%s,%s,%s)",
+								   port->noblock,
+								   quote_identifier(port->remote_host),
+								   quote_identifier(port->remote_hostname),
+								   port->remote_hostname_resolv,
+								   port->remote_hostname_errcode,
+								   quote_identifier(port->database_name),
+								   quote_identifier(port->user_name),
+								   quote_identifier(port->application_name));
+	else
+
+		/*
+		 * Should be unreachable.  If we add more fields in the future, we
+		 * need to modify the logic above.
+		 */
+		ereport(ERROR,
+				errmsg("\"%s.clientauth\" feature encountered an unexpected number of fields in the \"%s.%s\" composite type: %d",
+					   PG_TLE_NSPNAME, PG_TLE_NSPNAME,
+					   TLE_CLIENTAUTH_PORT_SUBSET_TYPE, tupdesc->natts));
+
+	return port_subset_str;
+}
diff --git a/test/t/004_pg_tle_clientauth.pl b/test/t/004_pg_tle_clientauth.pl
@@ -29,6 +29,8 @@
 ### 15. Rejects connections when no schema qualified function is found
 ### 16. Database does not come up if clientauth workers fail to start
 ### 17. Malformed strings cannot be used for SQL injection
+### 18. pg_tle can be updated from 1.4.1 to 1.5.0 without affecting clientauth functions
+### 19. application_name field works
 
 use strict;
 use warnings;
@@ -322,5 +324,42 @@
 like($psql_out, qr/^"\) \/\*$/,
     "role with injection payload in name can connect");
 
+### 18. pg_tle can be updated from 1.4.1 to 1.5.0 without affecting clientauth functions
+$node->safe_psql('postgres', 'DROP EXTENSION pg_tle CASCADE');
+$node->safe_psql('postgres', "CREATE EXTENSION pg_tle VERSION '1.4.1'");
+$node->safe_psql('postgres', q[
+    CREATE FUNCTION reject_testuser(port pgtle.clientauth_port_subset, status integer) RETURNS void AS $$
+        BEGIN
+            IF port.user_name = 'testuser' THEN
+                RAISE EXCEPTION 'testuser is not allowed to connect';
+            END IF;
+        END
+    $$ LANGUAGE plpgsql]);
+$node->safe_psql('postgres', qq[SELECT pgtle.register_feature('reject_testuser', 'clientauth')]);
+
+$node->safe_psql('postgres', 'SELECT 1');
+$node->psql('not_excluded', 'select', extra_params => ['-U', 'testuser'], stderr => \$psql_err);
+like($psql_err, qr/FATAL:  testuser is not allowed to connect/,
+    "clientauth function works on pg_tle 1.4.1");
+
+$node->safe_psql('postgres', "ALTER EXTENSION pg_tle UPDATE TO '1.5.0'");
+$node->psql('not_excluded', 'select', extra_params => ['-U', 'testuser'], stderr => \$psql_err);
+like($psql_err, qr/FATAL:  testuser is not allowed to connect/,
+    "clientauth function works on pg_tle 1.5.0");
+
+### 19. application_name field works
+$node->safe_psql('postgres', q[
+    CREATE OR REPLACE FUNCTION reject_testuser(port pgtle.clientauth_port_subset, status integer) RETURNS void AS $$
+        BEGIN
+            IF port.user_name = 'testuser' THEN
+                RAISE EXCEPTION '%', port.application_name;
+            END IF;
+        END
+    $$ LANGUAGE plpgsql]);
+
+$node->psql('not_excluded', 'select', extra_params => ['-U', 'testuser'], stderr => \$psql_err);
+like($psql_err, qr/FATAL:  004_pg_tle_clientauth.pl/,
+    "application_name field works on pg_tle 1.5.0");
+
 $node->stop;
-done_testing();
+done_testing();

Original file line number	Diff line number	Diff line change
`@@ -36,6 +36,8 @@`
`36`	`36`	`#define TLE_INPUT_FUNC_STR "input"`
`37`	`37`	`#define TLE_OUTPUT_FUNC_STR "output"`
`38`	`38`
	`39`	`+#define TLE_CLIENTAUTH_PORT_SUBSET_TYPE "clientauth_port_subset"`
	`40`	`+`
`39`	`41`	`/*`
`40`	`42`	`* TLE_BASE_TYPE_SIZE_LIMIT is the maximum allowed size of pg_tle type.`
`41`	`43`	`*`