From 50b72a9f68403af7215ea2105ddbeb4e23e8e500 Mon Sep 17 00:00:00 2001 From: Kushaal Shroff <51415286+KushaalShroff@users.noreply.github.com> Date: Thu, 26 Sep 2024 13:25:05 +0530 Subject: [PATCH] Reset db context at the time TDS resets the connection (#2971) T-SQL Behaviour suggests that if we connect to database db1 and if during the session we have changed the database context to db2 then at the time of reset connection, the server must reset the connection to db1. Earlier we were not resetting the database context to that of the database used to login, in the above example db1, this lead to clients being handed a stale connection. To Fix this we reset the database context to that from the loginInfo which was maintained at time of login. Changes were also made to avoid sending the environment change token for the implicit "USE DB" being run at time of reset. Issues Resolved BABEL-5256 Signed off by: Kushaal Shroff --- .../babelfishpg_tds/src/backend/tds/tds_srv.c | 1 + .../src/backend/tds/tdslogin.c | 198 +++++++++++------- .../src/backend/tds/tdsprotocol.c | 32 ++- contrib/babelfishpg_tds/src/include/tds_int.h | 1 + .../src/include/tds_protocol.h | 3 + contrib/babelfishpg_tsql/src/pl_exec-2.c | 21 +- contrib/babelfishpg_tsql/src/pltsql.h | 2 + contrib/babelfishpg_tsql/src/session.c | 1 - .../expected/Test-sp_reset_connection.out | 76 +++++++ ...ction.sql => Test-sp_reset_connection.mix} | 43 ++++ test/dotnet/ExpectedOutput/1_Setup.out | 5 + .../ExpectedOutput/2_Successful_reset.out | 4 + test/dotnet/input/ResetConnection/1_Setup.txt | 3 + .../ResetConnection/2_Successful_reset.txt | 2 + test/dotnet/src/ExecuteTests.cs | 2 + 15 files changed, 308 insertions(+), 86 deletions(-) rename test/JDBC/input/storedProcedures/{Test-sp_reset_connection.sql => Test-sp_reset_connection.mix} (67%) create mode 100644 test/dotnet/ExpectedOutput/1_Setup.out create mode 100644 test/dotnet/ExpectedOutput/2_Successful_reset.out create mode 100644 test/dotnet/input/ResetConnection/1_Setup.txt create mode 100644 test/dotnet/input/ResetConnection/2_Successful_reset.txt diff --git a/contrib/babelfishpg_tds/src/backend/tds/tds_srv.c b/contrib/babelfishpg_tds/src/backend/tds/tds_srv.c index 5783809bd9..46f0bd97dd 100644 --- a/contrib/babelfishpg_tds/src/backend/tds/tds_srv.c +++ b/contrib/babelfishpg_tds/src/backend/tds/tds_srv.c @@ -198,6 +198,7 @@ pe_tds_init(void) pltsql_plugin_handler_ptr->invalidate_stat_view = &invalidate_stat_table; pltsql_plugin_handler_ptr->get_host_name = &get_tds_host_name; pltsql_plugin_handler_ptr->set_reset_tds_connection_flag = &SetResetTDSConnectionFlag; + pltsql_plugin_handler_ptr->get_reset_tds_connection_flag = &GetResetTDSConnectionFlag; invalidate_stat_table_hook = invalidate_stat_table; guc_newval_hook = TdsSetGucStatVariable; diff --git a/contrib/babelfishpg_tds/src/backend/tds/tdslogin.c b/contrib/babelfishpg_tds/src/backend/tds/tdslogin.c index 37f7edd81e..2582a86fe5 100644 --- a/contrib/babelfishpg_tds/src/backend/tds/tdslogin.c +++ b/contrib/babelfishpg_tds/src/backend/tds/tdslogin.c @@ -60,6 +60,7 @@ #include "src/include/tds_debug.h" #include "src/include/tds_int.h" +#include "src/include/tds_protocol.h" #include "src/include/tds_request.h" #include "src/include/tds_response.h" #include "src/include/guc.h" @@ -2008,6 +2009,128 @@ TdsProcessLogin(Port *port, bool loadedSsl) return rc; } +/* + * TdsSetDbContext: + * Used to Set the Database Context during login + * and during reset connection. + * Note: We should not optimize the scenario during + * reset connection to reset to the same database + * which might be in use since the USE db command + * will reset other configurations which might + * have changed. + */ +void +TdsSetDbContext() +{ + char *dbname = NULL; + char *useDbCommand = NULL; + char *user = NULL; + MemoryContext oldContext = CurrentMemoryContext; + + PG_TRY(); + { + if (loginInfo->database != NULL && loginInfo->database[0] != '\0') + { + Oid db_id; + + /* + * Before preparing the query, first check whether we got a valid + * database name and it exists. Otherwise, there'll be risk of + * SQL injection. + */ + StartTransactionCommand(); + db_id = pltsql_plugin_handler_ptr->pltsql_get_database_oid(loginInfo->database); + CommitTransactionCommand(); + MemoryContextSwitchTo(oldContext); + + if (!OidIsValid(db_id)) + ereport(ERROR, + (errcode(ERRCODE_UNDEFINED_DATABASE), + errmsg("database \"%s\" does not exist", loginInfo->database))); + + /* + * Any delimitated/quoted db name identifier requested in login + * must be already handled before this point. + */ + useDbCommand = psprintf("USE [%s]", loginInfo->database); + dbname = pstrdup(loginInfo->database); + } + else + { + char *temp = NULL; + + StartTransactionCommand(); + temp = pltsql_plugin_handler_ptr->pltsql_get_login_default_db(loginInfo->username); + MemoryContextSwitchTo(oldContext); + + if (temp == NULL) + ereport(ERROR, + (errcode(ERRCODE_UNDEFINED_DATABASE), + errmsg("could not find default database for user \"%s\"", loginInfo->username))); + + useDbCommand = psprintf("USE [%s]", temp); + dbname = pstrdup(temp); + CommitTransactionCommand(); + MemoryContextSwitchTo(oldContext); + } + + StartTransactionCommand(); + /* + * Check if user has privileges to access current database. + */ + user = pltsql_plugin_handler_ptr->pltsql_get_user_for_database(dbname); + if (!user) + ereport(ERROR, + (errcode(ERRCODE_UNDEFINED_DATABASE), + errmsg("Cannot open database \"%s\" requested by the login. The login failed", dbname))); + + /* + * loginInfo has a database name provided, so we execute a "USE + * []" through pltsql inline handler. + */ + ExecuteSQLBatch(useDbCommand); + CommitTransactionCommand(); + } + PG_CATCH(); + { + /* + * If this is during reset phase and we encounter an error + * with mapped user or db not found then we should terminate + * the connection. + */ + if (resetTdsConnectionFlag) + { + /* Before terminating the connection, send the response to the client. */ + EmitErrorReport(); + FlushErrorState(); + + /* + * Client driver terminates the connection with a + * dual error token and with error 596. Otherwise + * it sends the next requests before realising the + * session was terminated. + */ + TdsSendError(596, 1, ERROR, + "Cannot continue the execution because the session is in the kill state.", 1); + + TdsSendDone(TDS_TOKEN_DONE, TDS_DONE_ERROR, 0, 0); + TdsFlush(); + + /* Terminate the connection. */ + ereport(FATAL, + (errcode(ERRCODE_INVALID_AUTHORIZATION_SPECIFICATION), + errmsg("Reset Connection Failed"))); + } + /* Else rethrow the error. */ + PG_RE_THROW(); + } + PG_END_TRY(); + if (useDbCommand) + pfree(useDbCommand); + if (dbname) + pfree(dbname); +} + /* * TdsSendLoginAck - Send a login acknowledgement to the client * @@ -2017,16 +2140,13 @@ void TdsSendLoginAck(Port *port) { uint16_t temp16; - char *dbname = NULL; int prognameLen = pg_mbstrlen(default_server_name); LoginRequest request; StringInfoData buf; uint8 temp8; uint32_t collationInfo; char collationBytesNew[5]; - char *useDbCommand = NULL; - char *user = NULL; - MemoryContext oldContext; + Oid roleid = InvalidOid; uint32_t tdsVersion = pg_hton32(loginInfo->tdsVersion); char srvVersionBytes[4]; @@ -2138,75 +2258,7 @@ TdsSendLoginAck(Port *port) errmsg("\"%s\" is not a Babelfish user", port->user_name))); } - oldContext = CurrentMemoryContext; - - if (request->database != NULL && request->database[0] != '\0') - { - Oid db_id; - - /* - * Before preparing the query, first check whether we got a valid - * database name and it exists. Otherwise, there'll be risk of - * SQL injection. - */ - StartTransactionCommand(); - db_id = pltsql_plugin_handler_ptr->pltsql_get_database_oid(request->database); - CommitTransactionCommand(); - MemoryContextSwitchTo(oldContext); - - if (!OidIsValid(db_id)) - ereport(ERROR, - (errcode(ERRCODE_UNDEFINED_DATABASE), - errmsg("database \"%s\" does not exist", request->database))); - - /* - * Any delimitated/quoted db name identifier requested in login - * must be already handled before this point. - */ - useDbCommand = psprintf("USE [%s]", request->database); - dbname = pstrdup(request->database); - } - else - { - char *temp = NULL; - - StartTransactionCommand(); - temp = pltsql_plugin_handler_ptr->pltsql_get_login_default_db(port->user_name); - MemoryContextSwitchTo(oldContext); - - if (temp == NULL) - ereport(ERROR, - (errcode(ERRCODE_UNDEFINED_DATABASE), - errmsg("could not find default database for user \"%s\"", port->user_name))); - - useDbCommand = psprintf("USE [%s]", temp); - dbname = pstrdup(temp); - CommitTransactionCommand(); - MemoryContextSwitchTo(oldContext); - } - - /* - * Check if user has privileges to access current database - */ - StartTransactionCommand(); - user = pltsql_plugin_handler_ptr->pltsql_get_user_for_database(dbname); - if (!user) - ereport(ERROR, - (errcode(ERRCODE_UNDEFINED_DATABASE), - errmsg("Cannot open database \"%s\" requested by the login. The login failed", dbname))); - CommitTransactionCommand(); - if (dbname) - pfree(dbname); - - /* - * Request has a database name provided, so we execute a "USE - * []" through pgtsql inline handler - */ - StartTransactionCommand(); - ExecuteSQLBatch(useDbCommand); - CommitTransactionCommand(); - if (useDbCommand) - pfree(useDbCommand); + TdsSetDbContext(); /* * Set the GUC for language, it will take care of changing the GUC, diff --git a/contrib/babelfishpg_tds/src/backend/tds/tdsprotocol.c b/contrib/babelfishpg_tds/src/backend/tds/tdsprotocol.c index 22b6c17c68..b7b0f5daca 100644 --- a/contrib/babelfishpg_tds/src/backend/tds/tdsprotocol.c +++ b/contrib/babelfishpg_tds/src/backend/tds/tdsprotocol.c @@ -73,7 +73,7 @@ typedef ResetConnectionData *ResetConnection; TdsRequestCtrlData *TdsRequestCtrl = NULL; ResetConnection resetCon = NULL; -static bool resetTdsConnectionFlag = false; +bool resetTdsConnectionFlag = false; /* Local functions */ static void ResetTDSConnection(void); @@ -155,16 +155,15 @@ ResetTDSConnection(void) TdsResetCache(); TdsResponseReset(); TdsResetBcpOffset(); - /* Retore previous isolation level when not called by sys.sp_reset_connection */ + /* Retore previous isolation level when not called by sys.sp_reset_connection. */ if (!resetTdsConnectionFlag) { SetConfigOption("default_transaction_isolation", isolationOld, PGC_BACKEND, PGC_S_CLIENT); } - tvp_lookup_list = NIL; - /* send an environement change token is its not called via sys.sp_reset_connection procedure */ + /* Send an environement change token is its not called via sys.sp_reset_connection procedure. */ if (!resetTdsConnectionFlag) { TdsSendEnvChange(TDS_ENVID_RESETCON, NULL, NULL); @@ -179,6 +178,11 @@ void SetResetTDSConnectionFlag() resetTdsConnectionFlag = true; } +bool GetResetTDSConnectionFlag() +{ + return resetTdsConnectionFlag; +} + /* * GetTDSRequest - Fetch and parse a TDS packet and generate a TDS request that * can be processed later. @@ -290,7 +294,16 @@ GetTDSRequest(bool *resetProtocol) resetCon->messageType = messageType; resetCon->status = (status & ~TDS_PACKET_HEADER_STATUS_RESETCON); + /* + * Set resetTdsConnectionFlag to true so that we avoid + * sending any env change token for the USE DB command + * which will get executed. + */ + resetTdsConnectionFlag = true; + TdsSetDbContext(); + resetTdsConnectionFlag = false; ResetTDSConnection(); + TdsErrorContext->err_text = "Fetching TDS Request"; *resetProtocol = true; return NULL; @@ -678,6 +691,17 @@ TdsSocketBackend(void) case TDS_REQUEST_PHASE_FLUSH: { TdsErrorContext->phase = "TDS_REQUEST_PHASE_FLUSH"; + + if (resetTdsConnectionFlag) + { + /* + * We must set the Db Context before resetting TDS state, + * becasue we need the existing TDS state to flush any errors + * along with the reset. + */ + TdsSetDbContext(); + } + /* Send the response now */ TdsFlush(); diff --git a/contrib/babelfishpg_tds/src/include/tds_int.h b/contrib/babelfishpg_tds/src/include/tds_int.h index 68e7dd1a8c..e09a4840a1 100644 --- a/contrib/babelfishpg_tds/src/include/tds_int.h +++ b/contrib/babelfishpg_tds/src/include/tds_int.h @@ -295,6 +295,7 @@ extern int TdsProcessLogin(Port *port, bool LoadSsl); extern void TdsSendLoginAck(Port *port); extern uint32_t GetClientTDSVersion(void); extern char *get_tds_login_domainname(void); +extern void TdsSetDbContext(void); /* Functions in backend/tds/tdsprotocol.c */ extern int TdsSocketBackend(void); diff --git a/contrib/babelfishpg_tds/src/include/tds_protocol.h b/contrib/babelfishpg_tds/src/include/tds_protocol.h index 4018a65382..7a80af2cea 100644 --- a/contrib/babelfishpg_tds/src/include/tds_protocol.h +++ b/contrib/babelfishpg_tds/src/include/tds_protocol.h @@ -78,5 +78,8 @@ typedef struct extern TdsRequestCtrlData *TdsRequestCtrl; extern void SetResetTDSConnectionFlag(void); +extern bool GetResetTDSConnectionFlag(void); + +extern bool resetTdsConnectionFlag; #endif /* TDS_PROTOCOL_H */ diff --git a/contrib/babelfishpg_tsql/src/pl_exec-2.c b/contrib/babelfishpg_tsql/src/pl_exec-2.c index a11280058c..7e3583dc9b 100644 --- a/contrib/babelfishpg_tsql/src/pl_exec-2.c +++ b/contrib/babelfishpg_tsql/src/pl_exec-2.c @@ -2725,14 +2725,19 @@ exec_stmt_usedb(PLtsql_execstate *estate, PLtsql_stmt_usedb *stmt) top_es_entry = top_es_entry->next; } - snprintf(message, sizeof(message), "Changed database context to '%s'.", stmt->db_name); - /* send env change token to user */ - if (*pltsql_protocol_plugin_ptr && (*pltsql_protocol_plugin_ptr)->send_env_change) - ((*pltsql_protocol_plugin_ptr)->send_env_change) (1, stmt->db_name, old_db_name); - /* send message to user */ - if (*pltsql_protocol_plugin_ptr && (*pltsql_protocol_plugin_ptr)->send_info) - ((*pltsql_protocol_plugin_ptr)->send_info) (0, 1, 0, message, 0); - + /* + * In case of reset-connection we do not need to send the environment change token. + */ + if (!((*pltsql_protocol_plugin_ptr) && (*pltsql_protocol_plugin_ptr)->get_reset_tds_connection_flag())) + { + snprintf(message, sizeof(message), "Changed database context to '%s'.", stmt->db_name); + /* send env change token to user */ + if (*pltsql_protocol_plugin_ptr && (*pltsql_protocol_plugin_ptr)->send_env_change) + ((*pltsql_protocol_plugin_ptr)->send_env_change) (1, stmt->db_name, old_db_name); + /* send message to user */ + if (*pltsql_protocol_plugin_ptr && (*pltsql_protocol_plugin_ptr)->send_info) + ((*pltsql_protocol_plugin_ptr)->send_info) (0, 1, 0, message, 0); + } return PLTSQL_RC_OK; } diff --git a/contrib/babelfishpg_tsql/src/pltsql.h b/contrib/babelfishpg_tsql/src/pltsql.h index ad2a1b41e3..05edff2795 100644 --- a/contrib/babelfishpg_tsql/src/pltsql.h +++ b/contrib/babelfishpg_tsql/src/pltsql.h @@ -1666,6 +1666,8 @@ typedef struct PLtsql_protocol_plugin void (*set_reset_tds_connection_flag) (); + bool (*get_reset_tds_connection_flag) (); + /* Session level GUCs */ bool quoted_identifier; bool arithabort; diff --git a/contrib/babelfishpg_tsql/src/session.c b/contrib/babelfishpg_tsql/src/session.c index d522e82cdd..a746017e2c 100644 --- a/contrib/babelfishpg_tsql/src/session.c +++ b/contrib/babelfishpg_tsql/src/session.c @@ -201,7 +201,6 @@ void reset_session_properties(void) { reset_cached_batch(); - set_session_properties(get_cur_db_name()); } void diff --git a/test/JDBC/expected/Test-sp_reset_connection.out b/test/JDBC/expected/Test-sp_reset_connection.out index 6398ab3ece..7c4f6e2be0 100644 --- a/test/JDBC/expected/Test-sp_reset_connection.out +++ b/test/JDBC/expected/Test-sp_reset_connection.out @@ -1,3 +1,4 @@ +-- tsql -- 1. Test resets GUC variables SET lock_timeout 0; GO @@ -107,3 +108,78 @@ smallint 2 ~~END~~ + +-- 6. Test Database Context being reset +-- Tests include negative cases where db is dropped or renamed +Create database reset_con_db1; +GO +Create database reset_con_db2; +GO + +-- tsql database=reset_con_db1 +select db_name(); +GO +~~START~~ +nvarchar +reset_con_db1 +~~END~~ + +exec sys.sp_reset_connection +GO +use master +GO +select db_name(); +GO +~~START~~ +nvarchar +master +~~END~~ + +exec sys.sp_reset_connection +GO +select db_name(); +GO +~~START~~ +nvarchar +reset_con_db1 +~~END~~ + +-- test db being dropped before resetting to same db +use master; +drop database reset_con_db1; +GO +exec sys.sp_reset_connection +GO +~~ERROR (Code: 911)~~ + +~~ERROR (Message: database "reset_con_db1" does not exist)~~ + +-- tsql database=reset_con_db2 +select db_name(); +GO +~~START~~ +nvarchar +reset_con_db2 +~~END~~ + +use master +GO +select db_name(); +GO +~~START~~ +nvarchar +master +~~END~~ + +ALTER DATABASE reset_con_db2 MODIFY NAME=reset_con_db3 +GO +exec sys.sp_reset_connection +GO +~~ERROR (Code: 911)~~ + +~~ERROR (Message: database "reset_con_db2" does not exist)~~ + + +-- tsql +DROP DATABASE reset_con_db3 +GO diff --git a/test/JDBC/input/storedProcedures/Test-sp_reset_connection.sql b/test/JDBC/input/storedProcedures/Test-sp_reset_connection.mix similarity index 67% rename from test/JDBC/input/storedProcedures/Test-sp_reset_connection.sql rename to test/JDBC/input/storedProcedures/Test-sp_reset_connection.mix index f29123db9b..30cf8469c6 100644 --- a/test/JDBC/input/storedProcedures/Test-sp_reset_connection.sql +++ b/test/JDBC/input/storedProcedures/Test-sp_reset_connection.mix @@ -1,3 +1,4 @@ +-- tsql -- 1. Test resets GUC variables SET lock_timeout 0; GO @@ -54,3 +55,45 @@ GO GO select transaction_isolation_level from sys.dm_exec_sessions where session_id = @@SPID GO + +-- 6. Test Database Context being reset +-- Tests include negative cases where db is dropped or renamed +Create database reset_con_db1; +GO +Create database reset_con_db2; +GO + +-- tsql database=reset_con_db1 +select db_name(); +GO +exec sys.sp_reset_connection +GO +use master +GO +select db_name(); +GO +exec sys.sp_reset_connection +GO +select db_name(); +GO +-- test db being dropped before resetting to same db +use master; +drop database reset_con_db1; +GO +exec sys.sp_reset_connection +GO +-- tsql database=reset_con_db2 +select db_name(); +GO +use master +GO +select db_name(); +GO +ALTER DATABASE reset_con_db2 MODIFY NAME=reset_con_db3 +GO +exec sys.sp_reset_connection +GO + +-- tsql +DROP DATABASE reset_con_db3 +GO \ No newline at end of file diff --git a/test/dotnet/ExpectedOutput/1_Setup.out b/test/dotnet/ExpectedOutput/1_Setup.out new file mode 100644 index 0000000000..55c7fa7e2b --- /dev/null +++ b/test/dotnet/ExpectedOutput/1_Setup.out @@ -0,0 +1,5 @@ +#Q#create database db1; +#Q#use db1; +#Q#Select db_name() +#D#nvarchar +db1 diff --git a/test/dotnet/ExpectedOutput/2_Successful_reset.out b/test/dotnet/ExpectedOutput/2_Successful_reset.out new file mode 100644 index 0000000000..394dc764ee --- /dev/null +++ b/test/dotnet/ExpectedOutput/2_Successful_reset.out @@ -0,0 +1,4 @@ +#Q#select db_name(); +#D#nvarchar +master +#Q#drop database db1; diff --git a/test/dotnet/input/ResetConnection/1_Setup.txt b/test/dotnet/input/ResetConnection/1_Setup.txt new file mode 100644 index 0000000000..2b363689d0 --- /dev/null +++ b/test/dotnet/input/ResetConnection/1_Setup.txt @@ -0,0 +1,3 @@ +create database db1; +use db1; +Select db_name() \ No newline at end of file diff --git a/test/dotnet/input/ResetConnection/2_Successful_reset.txt b/test/dotnet/input/ResetConnection/2_Successful_reset.txt new file mode 100644 index 0000000000..123124922d --- /dev/null +++ b/test/dotnet/input/ResetConnection/2_Successful_reset.txt @@ -0,0 +1,2 @@ +select db_name(); +drop database db1; diff --git a/test/dotnet/src/ExecuteTests.cs b/test/dotnet/src/ExecuteTests.cs index 40b212e85a..3c326d3642 100644 --- a/test/dotnet/src/ExecuteTests.cs +++ b/test/dotnet/src/ExecuteTests.cs @@ -52,6 +52,8 @@ public void Test() } allFiles = tempList; } + allFiles = allFiles.OrderBy(file => file.DirectoryName) + .ThenBy(file => file.Name); Task[] tasksInParallel = new Task[allFiles.Count()]; bool [] result = new bool[allFiles.Count()]; int i = 0;