Skip to content

Commit

Permalink
Return an error if feature is not supported, instead of crashing (ass…
Browse files Browse the repository at this point in the history
…ert)
  • Loading branch information
rahim-kanji committed Sep 24, 2024
1 parent 0ed3347 commit 5f21844
Show file tree
Hide file tree
Showing 3 changed files with 182 additions and 8 deletions.
1 change: 1 addition & 0 deletions lib/PgSQL_Connection.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
54 changes: 46 additions & 8 deletions lib/PgSQL_Session.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -3307,17 +3325,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");
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 {
Expand Down Expand Up @@ -3515,9 +3540,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");
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;
Expand Down
135 changes: 135 additions & 0 deletions test/tap/tests/pgsql-unsupported_feature_test-t.cpp
Original file line number Diff line number Diff line change
@@ -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();
}

0 comments on commit 5f21844

Please sign in to comment.