Replace PQconnectdb() by PQconnectdbParams() (#384)

Currently pg_repack calls `PQconnectdb()` to connect to a database. The drawback is that it is necessary to pass a connection, and all unexpected characters should be escaped. pg_repack has separate options for connection parameters (host, port, dbname, etc) and it isn't necessary to use `PQconnectdb()`. More robust and solid approach is to use `PQconnectdbParams()`, which accepts connection parameters as two arrays of values and it doesn't require escaping value strings.
reorg · Apr 14, 2024 · 90da0a2 · 90da0a2
1 parent 7da5302
commit 90da0a2
Show file tree

Hide file tree

Showing 3 changed files with 59 additions and 71 deletions.
diff --git a/bin/pgut/pgut-fe.c b/bin/pgut/pgut-fe.c
@@ -44,25 +44,30 @@ static bool parse_pair(const char buffer[], char key[], char value[]);
 void
 setup_workers(int num_workers)
 {
- 	StringInfoData	buf;
  	int i;
  	PGconn *conn;
 
  	elog(DEBUG2, "In setup_workers(), target num_workers = %d", num_workers);
 
  	if (num_workers > 1 && num_workers > workers.num_workers)
  	{
- 		initStringInfo(&buf);
- 		if (dbname && dbname[0])
- 			appendStringInfo(&buf, "dbname=%s ", dbname);
- 		if (host && host[0])
- 			appendStringInfo(&buf, "host=%s ", host);
- 		if (port && port[0])
- 			appendStringInfo(&buf, "port=%s ", port);
- 		if (username && username[0])
- 			appendStringInfo(&buf, "user=%s ", username);
- 		if (password && password[0])
- 			appendStringInfo(&buf, "password=%s ", password);
+#define PARAMS_ARRAY_SIZE	6
+
+		const char *keywords[PARAMS_ARRAY_SIZE];
+		const char *values[PARAMS_ARRAY_SIZE];
+
+		keywords[0] = "host";
+		values[0] = host;
+		keywords[1] = "port";
+		values[1] = port;
+		keywords[2] = "user";
+		values[2] = username;
+		keywords[3] = "password";
+		values[3] = password;
+		keywords[4] = "dbname";
+		values[4] = dbname;
+		keywords[5] = NULL;
+		values[5] = NULL;
 
  		if (workers.conns == NULL)
  		{
@@ -86,11 +91,11 @@ setup_workers(int num_workers)
 			 * XXX: could use PQconnectStart() and PQconnectPoll() to
 			 * open these connections in non-blocking manner.
 			 */
- 			conn = PQconnectdb(buf.data);
- 			if (PQstatus(conn) == CONNECTION_OK)
- 			{
- 				workers.conns[i] = conn;
- 			}
+			conn = PQconnectdbParams(keywords, values, true);
+			if (PQstatus(conn) == CONNECTION_OK)
+			{
+				workers.conns[i] = conn;
+			}
 			else
 			{
 				elog(WARNING, "Unable to set up worker conn #%d: %s", i,
@@ -113,10 +118,8 @@ setup_workers(int num_workers)
 		/* In case we bailed out of setting up all workers, record
 		 * how many successful worker conns we actually have.
 		 */
- 		workers.num_workers = i;
-
- 		termStringInfo(&buf);
- 	}
+		workers.num_workers = i;
+	}
 }
 
 /* Disconnect all our worker conns. */
@@ -154,24 +157,12 @@ void disconnect_workers(void)
 void
 reconnect(int elevel)
 {
-	StringInfoData	buf;
 	char		   *new_password;
 
 	disconnect();
-	initStringInfo(&buf);
-	if (dbname && dbname[0])
-		appendStringInfo(&buf, "dbname=%s ", dbname);
-	if (host && host[0])
-		appendStringInfo(&buf, "host=%s ", host);
-	if (port && port[0])
-		appendStringInfo(&buf, "port=%s ", port);
-	if (username && username[0])
-		appendStringInfo(&buf, "user=%s ", username);
-	if (password && password[0])
-		appendStringInfo(&buf, "password=%s ", password);
-
-	connection = pgut_connect(buf.data, prompt_password, elevel);
-	conn2      = pgut_connect(buf.data, prompt_password, elevel);
+
+	connection = pgut_connect(dbname, host, port, username, password, prompt_password, elevel);
+	conn2      = pgut_connect(dbname, host, port, username, password, prompt_password, elevel);
 
 	/* update password */
 	if (connection)
@@ -184,8 +175,6 @@ reconnect(int elevel)
 			password = pgut_strdup(new_password);
 		}
 	}
-
-	termStringInfo(&buf);
 }
 
 void

diff --git a/bin/pgut/pgut.c b/bin/pgut/pgut.c
@@ -496,34 +496,40 @@ prompt_for_password(void)
 
 
 PGconn *
-pgut_connect(const char *info, YesNo prompt, int elevel)
+pgut_connect(const char *dbname, const char *host, const char *port,
+			 const char *username, const char *password,
+			 YesNo prompt, int elevel)
 {
-	char	   *passwd;
-	StringInfoData add_pass;
+	char	   *new_password = NULL;
 
 	if (prompt == YES)
-	{
-		passwd = prompt_for_password();
-		initStringInfo(&add_pass);
-		appendStringInfoString(&add_pass, info);
-		appendStringInfo(&add_pass, " password=%s ", passwd);
-	}
-	else
-	{
-		passwd = NULL;
-		add_pass.data = NULL;
-	}
+		new_password = prompt_for_password();
 
 	/* Start the connection. Loop until we have a password if requested by backend. */
 	for (;;)
 	{
+#define PARAMS_ARRAY_SIZE	6
+
+		const char *keywords[PARAMS_ARRAY_SIZE];
+		const char *values[PARAMS_ARRAY_SIZE];
 		PGconn	   *conn;
+
 		CHECK_FOR_INTERRUPTS();
 
-		if (!passwd)
-			conn = PQconnectdb(info);
-		else
-			conn = PQconnectdb(add_pass.data);
+		keywords[0] = "host";
+		values[0] = host;
+		keywords[1] = "port";
+		values[1] = port;
+		keywords[2] = "user";
+		values[2] = username;
+		keywords[3] = "password";
+		values[3] = (new_password != NULL) ? new_password : password;
+		keywords[4] = "dbname";
+		values[4] = dbname;
+		keywords[5] = NULL;
+		values[5] = NULL;
+
+		conn = PQconnectdbParams(keywords, values, true);
 
 		if (PQstatus(conn) == CONNECTION_OK)
 		{
@@ -538,32 +544,23 @@ pgut_connect(const char *info, YesNo prompt, int elevel)
 			pgut_connections = c;
 			pgut_conn_unlock();
 
-			if (add_pass.data != NULL)
-				termStringInfo(&add_pass);
-			free(passwd);
+			free(new_password);
 
 			/* Hardcode a search path to avoid injections into public or pg_temp */
 			pgut_command(conn, "SET search_path TO pg_catalog, pg_temp, public", 0, NULL);
 
 			return conn;
 		}
 
-		if (conn && PQconnectionNeedsPassword(conn) && !passwd && prompt != NO)
+		if (conn && PQconnectionNeedsPassword(conn) && !new_password && prompt != NO)
 		{
 			PQfinish(conn);
-			passwd = prompt_for_password();
-			if (add_pass.data != NULL)
-	 			resetStringInfo(&add_pass);
-			else
-	 			initStringInfo(&add_pass);
-			appendStringInfoString(&add_pass, info);
-			appendStringInfo(&add_pass, " password=%s ", passwd);
+			new_password = prompt_for_password();
+
 			continue;
 		}
 
-		if (add_pass.data != NULL)
-			termStringInfo(&add_pass);
-		free(passwd);
+		free(new_password);
 
 		ereport(elevel,
 			(errcode(E_PG_CONNECT),

diff --git a/bin/pgut/pgut.h b/bin/pgut/pgut.h
@@ -61,7 +61,9 @@ extern void pgut_putenv(const char *key, const char *value);
 /*
  * Database connections
  */
-extern PGconn *pgut_connect(const char *info, YesNo prompt, int elevel);
+extern PGconn *pgut_connect(const char *dbname, const char *host, const char *port,
+							const char *username, const char *password,
+							YesNo prompt, int elevel);
 extern void pgut_disconnect(PGconn *conn);
 extern void pgut_disconnect_all(void);
 extern PGresult *pgut_execute(PGconn* conn, const char *query, int nParams, const char **params);