From 8c9af56f38f7d40a2c768330d5a504b82b4facd2 Mon Sep 17 00:00:00 2001 From: Rahim Kanji Date: Sun, 1 Mar 2026 22:47:24 +0500 Subject: [PATCH] Address AI review feedback --- lib/Admin_Handler.cpp | 5 +- lib/PgSQL_Connection.cpp | 9 ++ lib/PgSQL_Session.cpp | 35 ++++---- .../tests/pgsql-test_malformed_packet-t.cpp | 88 ++++++++++++------- 4 files changed, 84 insertions(+), 53 deletions(-) diff --git a/lib/Admin_Handler.cpp b/lib/Admin_Handler.cpp index 58185f634..714309a1f 100644 --- a/lib/Admin_Handler.cpp +++ b/lib/Admin_Handler.cpp @@ -2877,8 +2877,9 @@ void admin_session_handler(S* sess, void *_pa, PtrSize_t *pkt) { query_length = hdr.data.size; // Validate minimum query size (need at least 1 byte + null terminator) - if (query_length < 2 || hdr.data.ptr == NULL) { - proxy_warning("Query too short: %u bytes\n", query_length); + if (query_length < 2 || hdr.data.ptr == NULL || + ((const char*)hdr.data.ptr)[query_length - 1] != '\0') { + proxy_warning("Malformed query packet: %u bytes\n", query_length); SPA->send_error_msg_to_client(sess, "Malformed query packet"); run_query = false; goto __run_query; diff --git a/lib/PgSQL_Connection.cpp b/lib/PgSQL_Connection.cpp index c954eecb4..dece0ac8a 100644 --- a/lib/PgSQL_Connection.cpp +++ b/lib/PgSQL_Connection.cpp @@ -2191,6 +2191,15 @@ bool PgSQL_Connection::handle_copy_out(const PGresult* result, uint64_t* process void PgSQL_Connection::notice_handler_cb(void* arg, const PGresult* result) { assert(arg); PgSQL_Connection* conn = (PgSQL_Connection*)arg; + if (conn->query_result == nullptr) { + // Notice received without active query_result. This can happen when: + // - RESET SESSION is in progress (DISCARD ALL or ROLLBACK) + proxy_debug(PROXY_DEBUG_MYSQL_COM, 5, "Notice received without active query_result [State: %d, FD: %d]: %s\n", + (int)conn->async_state_machine, + conn->get_pg_socket_fd(), + (result ? PQresultErrorMessage(result) : "unknown notice")); + return; + } const unsigned int bytes_recv = conn->query_result->add_notice(result); conn->update_bytes_recv(bytes_recv); } diff --git a/lib/PgSQL_Session.cpp b/lib/PgSQL_Session.cpp index 7b184e548..a0c3721b7 100644 --- a/lib/PgSQL_Session.cpp +++ b/lib/PgSQL_Session.cpp @@ -2048,6 +2048,24 @@ __implicit_sync: if (session_type == PROXYSQL_SESSION_PGSQL) { bool rc_break = false; bool lock_hostgroup = false; + + if (pkt.size < 6) { + proxy_error("Malformed query packet received: size %u < 6\n", pkt.size); + l_free(pkt.size, pkt.ptr); + handler_ret = -1; + return handler_ret; + } + + unsigned int query_len = pkt.size - 5; // excluding header + char* query_ptr = (char*)pkt.ptr + 5; + + if (query_ptr[query_len - 1] != '\0') { + proxy_error("Malformed query packet received: missing null terminator\n"); + l_free(pkt.size, pkt.ptr); + handler_ret = -1; + return handler_ret; + } + if (session_fast_forward == SESSION_FORWARD_TYPE_NONE) { // Note: CurrentQuery sees the query as sent by the client. // shortly after, the packets it used to contain the query will be deallocated @@ -2073,23 +2091,6 @@ __implicit_sync: clock_gettime(CLOCK_THREAD_CPUTIME_ID, &begint); } - if (pkt.size < 6) { - proxy_error("Malformed query packet received: size %u < 6\n", pkt.size); - l_free(pkt.size, pkt.ptr); - handler_ret = -1; - return handler_ret; - } - - unsigned int query_len = pkt.size - 5; // excluding header - char* query_ptr = (char*)pkt.ptr + 5; - - if (query_ptr[query_len - 1] != '\0') { - proxy_error("Malformed query packet received: missing null terminator\n"); - l_free(pkt.size, pkt.ptr); - handler_ret = -1; - return handler_ret; - } - qpo = GloPgQPro->process_query(this, query_ptr, query_len, &CurrentQuery); if (thread->variables.stats_time_query_processor) { clock_gettime(CLOCK_THREAD_CPUTIME_ID, &endt); diff --git a/test/tap/tests/pgsql-test_malformed_packet-t.cpp b/test/tap/tests/pgsql-test_malformed_packet-t.cpp index d8d5ca5b2..55cdeef4a 100644 --- a/test/tap/tests/pgsql-test_malformed_packet-t.cpp +++ b/test/tap/tests/pgsql-test_malformed_packet-t.cpp @@ -139,10 +139,11 @@ void test_malformed_packet(const std::string& test_name, std::vector buffer(BUFFER_SIZE); ssize_t bytes_received = recv(sock, buffer.data(), buffer.size(), 0); - // Either connection close OR error response are valid outcomes - // The key is that ProxySQL handles the malformed packet gracefully - bool handled_gracefully = (bytes_received <= 0) || // Connection closed - (bytes_received > 0); // Got error response + // Valid outcomes: connection closed OR any response + // The key is that ProxySQL handles the packet without crashing + bool connection_closed = (bytes_received <= 0); + bool got_response = (bytes_received > 0); + bool handled_gracefully = connection_closed || got_response; ok(handled_gracefully, "%s: Malformed packet handled (received: %ld bytes)", test_name.c_str(), (long)bytes_received); @@ -472,20 +473,33 @@ std::vector build_message_packet(char type, const std::vector& return packet; } +/** + * @brief Read an exact number of bytes from socket + * @return true if all bytes received, false on error/short read + */ +bool recv_exact(int sock, void* buf, size_t len) { + char* ptr = (char*)buf; + size_t received = 0; + while (received < len) { + ssize_t n = recv(sock, ptr + received, len - received, 0); + if (n <= 0) return false; + received += n; + } + return true; +} + /** * @brief Read a PostgreSQL message from socket */ bool read_message(int sock, char& type, std::vector& buffer) { // Read message type (1 byte) char type_buf[1]; - ssize_t n = recv(sock, type_buf, 1, 0); - if (n <= 0) return false; + if (!recv_exact(sock, type_buf, 1)) return false; type = type_buf[0]; // Read length (4 bytes, big-endian, includes self) char len_buf[4]; - n = recv(sock, len_buf, 4, 0); - if (n != 4) return false; + if (!recv_exact(sock, len_buf, 4)) return false; int32_t len = ((unsigned char)len_buf[0] << 24) | ((unsigned char)len_buf[1] << 16) | @@ -497,12 +511,7 @@ bool read_message(int sock, char& type, std::vector& buffer) { // Read message body int32_t body_len = len - 4; buffer.resize(body_len); - int32_t received = 0; - while (received < body_len) { - n = recv(sock, buffer.data() + received, body_len - received, 0); - if (n <= 0) return false; - received += n; - } + if (!recv_exact(sock, buffer.data(), body_len)) return false; return true; } @@ -560,22 +569,19 @@ void test_malformed_packet_phase2(const std::string& test_name, // Check if authentication succeeded if (auth_type == 'E') { - // Authentication failed - this is OK for our purposes, - // we still want to test if malformed packets crash ProxySQL - // but we can't send them on an authenticated connection - // Send the malformed data anyway to see if it crashes - ssize_t sent = send(sock, malformed_data.data(), malformed_data.size(), 0); + // Authentication failed - fail the test + // Post-authentication tests require successful auth close(sock); - ok(sent > 0, "%s: Sent malformed data after auth failure (sent: %ld)", - test_name.c_str(), (long)sent); + ok(0, "%s: Authentication failed, cannot run post-auth test", + test_name.c_str()); return; } if (auth_type != 'R' || auth_data.size() < 4) { - // Unexpected response - send(sock, malformed_data.data(), malformed_data.size(), 0); + // Unexpected response - fail the test close(sock); - ok(1, "%s: Malformed packet sent after auth", test_name.c_str()); + ok(0, "%s: Unexpected auth response, cannot run post-auth test", + test_name.c_str()); return; } @@ -584,16 +590,16 @@ void test_malformed_packet_phase2(const std::string& test_name, (auth_data[2] << 8) | auth_data[3]; if (auth_code != 0) { - // Authentication didn't complete successfully - send(sock, malformed_data.data(), malformed_data.size(), 0); + // Authentication didn't complete successfully - fail the test close(sock); - ok(1, "%s: Malformed packet sent after incomplete auth", test_name.c_str()); + ok(0, "%s: Authentication incomplete (code: %d), cannot run post-auth test", + test_name.c_str(), auth_code); return; } // Authentication successful - now we have an authenticated connection - // Step 5: Send the malformed packet + // Step 5: Send the malformed packet on the authenticated connection ssize_t sent = send(sock, malformed_data.data(), malformed_data.size(), 0); if (sent < 0) { close(sock); @@ -601,14 +607,28 @@ void test_malformed_packet_phase2(const std::string& test_name, return; } - // Step 6: Try to receive response (may get error or connection close) + // Step 6: Try to receive response + // ProxySQL may send multiple responses (auth completion, parameter status, etc.) + // followed by an error response 'E' for the malformed packet, or close connection std::vector buffer(BUFFER_SIZE); ssize_t bytes_received = recv(sock, buffer.data(), buffer.size(), 0); - bool handled_gracefully = (bytes_received <= 0) || (bytes_received > 0); + // Check if we got an error response 'E' or connection closed + // Note: There may be pending post-auth messages, so any of these outcomes is valid: + // 1. Connection closed (ProxySQL rejected the malformed packet) + // 2. 'E' error response (ProxySQL sent an error for the malformed packet) + // 3. Other messages (post-auth protocol messages) + bool connection_closed = (bytes_received <= 0); + bool got_error_response = (bytes_received > 0 && buffer[0] == 'E'); + bool got_other_response = (bytes_received > 0 && buffer[0] != 'E'); - ok(handled_gracefully, "%s: Phase 2 malformed packet handled (received: %ld bytes)", - test_name.c_str(), (long)bytes_received); + // Any outcome is acceptable as long as ProxySQL doesn't crash + // The key test is the final operational check + bool handled_gracefully = connection_closed || got_error_response || got_other_response; + + ok(handled_gracefully, "%s: Phase 2 malformed packet sent (received: %ld bytes, first=0x%02X)", + test_name.c_str(), (long)bytes_received, + bytes_received > 0 ? (unsigned char)buffer[0] : 0); close(sock); } @@ -989,7 +1009,7 @@ int main(int argc, char** argv) { ok(backend_alive, "ProxySQL BACKEND operational after Phase 1"); // Test ADMIN connection - run_malformed_packet_tests(cl.pgsql_host, cl.pgsql_admin_port, "ADMIN"); + run_malformed_packet_tests(cl.admin_host, cl.pgsql_admin_port, "ADMIN"); // Verify ProxySQL is still alive after ADMIN Phase 1 tests bool admin_alive = verify_proxysql_alive(cl, ADMIN); @@ -1009,7 +1029,7 @@ int main(int argc, char** argv) { ok(backend_alive, "ProxySQL BACKEND operational after Phase 2"); // Phase 2: Test malformed packets AFTER authentication on ADMIN - run_phase2_malformed_packet_tests(cl.pgsql_host, cl.pgsql_admin_port, + run_phase2_malformed_packet_tests(cl.admin_host, cl.pgsql_admin_port, cl.admin_username, cl.admin_password, "ADMIN");