Graceful handling of unsupported features

lib/PgSQL_Connection.cpp

@@ -1807,6 +1807,7 @@ PG_ASYNC_ST PgSQL_Connection::handler(short event) {
 				case PGRES_COPY_IN:
 				case PGRES_COPY_BOTH:
 					// NOT IMPLEMENTED
+					proxy_error("COPY not supported\n");
 					assert(0);
 					break;
 				case PGRES_BAD_RESPONSE:

lib/PgSQL_Session.cpp

@@ -959,6 +959,24 @@ bool PgSQL_Session::handler_special_queries(PtrSize_t* pkt) {
 			return true;
 		}
 	}
+	// Unsupported Features:
+	// COPY
+	if (pkt->size > (5 + 5) && strncasecmp((char*)"COPY ", (char*)pkt->ptr + 5, 5) == 0) {
+		client_myds->DSS = STATE_QUERY_SENT_NET;
+		client_myds->myprot.generate_error_packet(true, true, "Feature not supported", PGSQL_ERROR_CODES::ERRCODE_FEATURE_NOT_SUPPORTED,
+			false, true);
+		//client_myds->DSS = STATE_SLEEP;
+		//status = WAITING_CLIENT_DATA;
+		if (mirror == false) {
+			RequestEnd(NULL);
+		} else {
+			client_myds->DSS = STATE_SLEEP;
+			status = WAITING_CLIENT_DATA;
+		}
+		l_free(pkt->size, pkt->ptr);
+		return true;
+	}
+	//
 	if (pkt->size > (5 + 18) && strncasecmp((char*)"PROXYSQL INTERNAL ", (char*)pkt->ptr + 5, 18) == 0) {
 		return_proxysql_internal(pkt);
 		return true;
@@ -2836,17 +2854,24 @@ int PgSQL_Session::get_pkts_from_client(bool& wrong_pass, PtrSize_t& pkt) {
 						c = *((unsigned char*)pkt.ptr + 0);
 						if (c == 'Q') {
 							handler___status_WAITING_CLIENT_DATA___STATE_SLEEP___MYSQL_COM_QUERY___not_mysql(pkt);
-						}
-						else if (c == 'X') {
+						} else if (c == 'X') {
 							//proxy_debug(PROXY_DEBUG_MYSQL_COM, 5, "Got COM_QUIT packet\n");
 							//if (GloPgSQL_Logger) { GloPgSQL_Logger->log_audit_entry(PROXYSQL_MYSQL_AUTH_QUIT, this, NULL); }
 							l_free(pkt.size, pkt.ptr);
 							handler_ret = -1;
 							return handler_ret;
-						}
-						else {
-							proxy_error("Not implemented yet\n");
-							assert(0);
+						} else if (c == 'P' || c == 'B' || c == 'D' || c == 'E') {
+							l_free(pkt.size, pkt.ptr);
+							continue;
+						} else {
+							proxy_error("Not implemented yet. Message type:'%c'\n", c);
+							client_myds->setDSS_STATE_QUERY_SENT_NET();
+							client_myds->myprot.generate_error_packet(true, true, "Feature not supported", PGSQL_ERROR_CODES::ERRCODE_FEATURE_NOT_SUPPORTED,
+								false, true);
+							l_free(pkt.size, pkt.ptr);
+							client_myds->DSS = STATE_SLEEP;
+							//handler_ret = -1;
+							return handler_ret;
 						}
 					}
 					else {
@@ -3035,9 +3060,22 @@ int PgSQL_Session::get_pkts_from_client(bool& wrong_pass, PtrSize_t& pkt) {
 							handler_ret = -1;
 							return handler_ret;
 							break;
+						case 'P':
+						case 'B':
+						case 'D':
+						case 'E':
+							//ignore
+							l_free(pkt.size, pkt.ptr);
+							continue;
+						case 'S':
 						default:
-							proxy_error("Not implemented yet\n");
-							assert(0);
+							proxy_error("Not implemented yet. Message type:'%c'\n", c);
+							client_myds->setDSS_STATE_QUERY_SENT_NET();
+							client_myds->myprot.generate_error_packet(true, true, "Feature not supported", PGSQL_ERROR_CODES::ERRCODE_FEATURE_NOT_SUPPORTED,
+								false, true);
+							l_free(pkt.size, pkt.ptr);
+							client_myds->DSS = STATE_SLEEP;
+							return handler_ret;
 						}
 					}
 					break;

test/tap/tests/pgsql-unsupported_feature_test-t.cpp

@@ -0,0 +1,135 @@
+/**
+ * @file pgsql-unsupported_feature_test-t.cpp
+ * @brief Ensures that ProxySQL does not crash and maintains the connection/session integrity when unsupported queries are executed.
+ * Currently validates:
+ * 1) Prepare Statement
+ * 2) COPY
+ */
+
+#include <string>
+#include <sstream>
+
+#include "libpq-fe.h"
+#include "command_line.h"
+#include "tap.h"
+#include "utils.h"
+
+CommandLine cl;
+
+PGconn* create_new_connection(bool with_ssl) {
+    std::stringstream ss;
+
+    ss << "host=" << cl.pgsql_host << " port=" << cl.pgsql_port;
+    ss << " user=" << cl.pgsql_username << " password=" << cl.pgsql_password;
+
+    if (with_ssl) {
+        ss << " sslmode=require";
+    } else {
+        ss << " sslmode=disable";
+    }
+
+    PGconn* conn = PQconnectdb(ss.str().c_str());
+    const bool res = (conn && PQstatus(conn) == CONNECTION_OK);
+    ok(res, "Connection created successfully. %s", PQerrorMessage(conn));
+
+    if (res) return conn;
+
+    PQfinish(conn);
+    return nullptr;
+}
+
+void check_transaction_state(PGconn* conn) {
+    PGresult* res;
+
+    // Check if the transaction is still active
+    res = PQexec(conn, "SELECT 1");
+    ok(PQresultStatus(res) == PGRES_TUPLES_OK && PQtransactionStatus(conn) == PQTRANS_INTRANS, 
+        "Transaction state was not affected by the error. %s", PQerrorMessage(conn));
+    PQclear(res);
+}
+
+void check_prepared_statement_binary(PGconn* conn) {
+    PGresult* res;
+    const char* paramValues[1] = { "1" };
+
+    // Start a transaction
+    res = PQexec(conn, "BEGIN");
+    if (PQresultStatus(res) != PGRES_COMMAND_OK) {
+        BAIL_OUT("Could not start transaction. %s", PQerrorMessage(conn));
+    }
+    PQclear(res);
+
+    // Test: Prepare a statement (using binary mode)
+    res = PQprepare(conn, "myplan", "SELECT $1::int", 1, NULL);
+    ok(PQresultStatus(res) != PGRES_COMMAND_OK, "Prepare statement failed. %s", PQerrorMessage(conn));
+    PQclear(res);
+
+    // Execute the prepared statement using binary protocol
+    res = PQexecPrepared(conn, "myplan", 1, paramValues, NULL, NULL, 1); // Binary result format (1)
+    ok(PQresultStatus(res) != PGRES_COMMAND_OK && PQresultStatus(res) != PGRES_TUPLES_OK, "Prepare statements are not supported for PostgreSQL: %s", PQerrorMessage(conn));
+    PQclear(res);
+
+    // Check if the transaction state is still active
+    check_transaction_state(conn);
+
+    // End the transaction
+    res = PQexec(conn, "ROLLBACK");
+    PQclear(res);
+}
+
+void check_copy_binary(PGconn* conn) {
+    PGresult* res;
+
+    // Start a transaction
+    res = PQexec(conn, "BEGIN");
+    if (PQresultStatus(res) != PGRES_COMMAND_OK) {
+        BAIL_OUT("Could not start transaction. %s", PQerrorMessage(conn));
+    }
+    PQclear(res);
+
+    // Test: COPY binary format
+    res = PQexec(conn, "COPY (SELECT 1) TO STDOUT (FORMAT BINARY)");
+    ok(PQresultStatus(res) != PGRES_COPY_OUT, "COPY binary command failed to start. %s", PQerrorMessage(conn));
+    PQclear(res);
+
+    // Attempt to fetch data in binary mode, expect it to fail
+    char buffer[256];
+    int ret = PQgetCopyData(conn, (char**)&buffer, 1); // Binary mode (1)
+    ok(ret == -2, "COPY in binary mode should have failed. %s", PQerrorMessage(conn));
+
+    // Check if the transaction state is still active
+    check_transaction_state(conn);
+
+    // End the transaction
+    res = PQexec(conn, "ROLLBACK");
+    PQclear(res);
+}
+
+void execute_tests(bool with_ssl) {
+    PGconn* conn = create_new_connection(with_ssl);
+
+    if (conn == nullptr)
+        return;
+
+    // Test 1: Prepared Statement in binary mode
+    check_prepared_statement_binary(conn);
+
+    // Test 2: COPY in binary mode
+    check_copy_binary(conn);
+
+    // Close the connection
+
+    PQfinish(conn);
+}
+
+int main(int argc, char** argv) {
+
+    plan(7); // Total number of tests planned
+
+    if (cl.getEnv())
+        return exit_status();
+
+    execute_tests(false); // without SSL
+
+    return exit_status();
+}