Refactor direct SSL handshake logic

This commit refactors the SSL handshake logic in the engine to support protocol-specific handling through the fn_ssl_handshake function pointer. A wrapper function encapsulates the static ProcessSSLStartup logic for non-TDS connections, ensuring the original functionality is preserved. The connection startup process now delegates SSL handshake handling to the protocol-defined function. Task: BABEL-5342 Signed-off-by: Roshan Kanwar <[email protected]>
babelfish-for-postgresql · Dec 16, 2024 · 3bd2029 · 3bd2029
1 parent db4e138
commit 3bd2029
Show file tree

Hide file tree

Showing 5 changed files with 24 additions and 5 deletions.
diff --git a/src/backend/postmaster/postmaster.c b/src/backend/postmaster/postmaster.c
@@ -247,7 +247,8 @@ ProtocolExtensionConfig default_protocol_config = {
 	libpq_end_command,
 	NULL, NULL, NULL, NULL,		/* use libpq defaults for printtup*() */
 	NULL,
-	libpq_report_param_status
+	libpq_report_param_status,
+	libpq_ssl_handshake
 };
 
 /* still more option variables */
@@ -1511,6 +1512,11 @@ libpq_end_command(QueryCompletion *qc, CommandDest dest)
 	EndCommand(qc, dest, false);
 }
 
+int
+libpq_ssl_handshake(struct Port *port)
+{
+	return WrapperProcessSSLStartup(port);
+}
 
 /*
  * on_proc_exit callback to close server's listen sockets

diff --git a/src/backend/tcop/backend_startup.c b/src/backend/tcop/backend_startup.c
@@ -253,15 +253,14 @@ BackendInitialize(ClientSocket *client_sock, CAC_state cac, ProtocolExtensionCon
 	RegisterTimeout(STARTUP_PACKET_TIMEOUT, StartupPacketTimeoutHandler);
 	enable_timeout_after(STARTUP_PACKET_TIMEOUT, AuthenticationTimeout * 1000);
 
-	/* Handle direct SSL handshake for non-TDS connections */
-	if (!port->is_tds_conn)
-	status = ProcessSSLStartup(port);
+	/* Handle protocol-specific SSL handshake */
+	status = port->protocol_config->fn_ssl_handshake(port);
 
 	/*
 	 * Receive the startup packet (which might turn out to be a cancel request
 	 * packet).
 	 */
-	if (port->is_tds_conn || status == STATUS_OK)
+	if (status == STATUS_OK)
 		status = (port->protocol_config->fn_start)(port);
 
 	/*
@@ -889,3 +888,12 @@ StartupPacketTimeoutHandler(void)
 {
 	_exit(1);
 }
+
+/*
+ * Wrapper for ProcessSSLStartup to handle direct SSL handshake
+ */
+int
+WrapperProcessSSLStartup(Port *port)
+{
+    return ProcessSSLStartup(port);
+}
diff --git a/src/include/libpq/libpq-be.h b/src/include/libpq/libpq-be.h
@@ -117,6 +117,9 @@ typedef struct ProtocolExtensionConfig {
 	void	(*fn_printtup_destroy)(DestReceiver *self);
 	int		(*fn_process_command)(void);
 	void	(*fn_report_param_status)(const char *name, char *val);
+
+	/* function pointer for handling direct SSL handshake */
+	int		(*fn_ssl_handshake)(struct Port *port);
 } ProtocolExtensionConfig;
 
 /*

diff --git a/src/include/postmaster/protocol_extension.h b/src/include/postmaster/protocol_extension.h
@@ -42,5 +42,6 @@ extern void	libpq_send_ready_for_query(CommandDest dest);
 extern int	libpq_read_command(StringInfo inBuf);
 extern void	libpq_end_command(QueryCompletion *qc, CommandDest dest);
 extern void	libpq_report_param_status(const char *name, char *val);
+extern int	libpq_ssl_handshake(struct Port *port);
 
 #endif							/* _PROTOCOL_EXTENSION_H */
diff --git a/src/include/tcop/backend_startup.h b/src/include/tcop/backend_startup.h
@@ -39,5 +39,6 @@ typedef struct BackendStartupData
 
 extern void BackendMain(char *startup_data, size_t startup_data_len) pg_attribute_noreturn();
 extern int	ProcessStartupPacket(Port *port, bool ssl_done, bool gss_done);
+extern int	WrapperProcessSSLStartup(Port *port);
 
 #endif							/* BACKEND_STARTUP_H */