Skip to content

Commit

Permalink
Fix segfaults when closing the client and calling each
Browse files Browse the repository at this point in the history
  • Loading branch information
andyundso committed Dec 17, 2024
1 parent 25ab0bd commit 2dd9313
Show file tree
Hide file tree
Showing 3 changed files with 125 additions and 42 deletions.
4 changes: 4 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
@@ -1,3 +1,7 @@
## (unreleased)

* Fix segfaults when closing client before or during calling `each`. Fixes #435.

## 3.0.0

* Drop support for Ruby < 2.7
Expand Down
91 changes: 49 additions & 42 deletions ext/tiny_tds/result.c
Original file line number Diff line number Diff line change
Expand Up @@ -77,6 +77,12 @@ VALUE rb_tinytds_new_result_obj(tinytds_client_wrapper *cwrap) {
(rb_unblock_function_t*)dbcancel_ubf, _client ) \
)

void raise_if_client_closed(tinytds_client_wrapper *cwrap) {
if (cwrap->closed || cwrap->userdata->closed) {
rb_raise(cTinyTdsError, "closed connection");
}
}

static void dbcancel_ubf(DBPROCESS *client) {
GET_CLIENT_USERDATA(client);
dbcancel(client);
Expand Down Expand Up @@ -122,37 +128,40 @@ static void nogvl_cleanup(DBPROCESS *client) {
userdata->nonblocking_errors_size = 0;
}

static RETCODE nogvl_dbsqlok(DBPROCESS *client) {
static RETCODE nogvl_dbsqlok(tinytds_client_wrapper *cwrap) {
int retcode = FAIL;
GET_CLIENT_USERDATA(client);
nogvl_setup(client);
retcode = NOGVL_DBCALL(dbsqlok, client);
nogvl_cleanup(client);
userdata->dbsqlok_sent = 1;
raise_if_client_closed(cwrap);
nogvl_setup(cwrap->client);
retcode = NOGVL_DBCALL(dbsqlok, cwrap->client);
nogvl_cleanup(cwrap->client);
cwrap->userdata->dbsqlok_sent = 1;
return retcode;
}

static RETCODE nogvl_dbsqlexec(DBPROCESS *client) {
static RETCODE nogvl_dbsqlexec(tinytds_client_wrapper *cwrap) {
int retcode = FAIL;
nogvl_setup(client);
retcode = NOGVL_DBCALL(dbsqlexec, client);
nogvl_cleanup(client);
raise_if_client_closed(cwrap);
nogvl_setup(cwrap->client);
retcode = NOGVL_DBCALL(dbsqlexec, cwrap->client);
nogvl_cleanup(cwrap->client);
return retcode;
}

static RETCODE nogvl_dbresults(DBPROCESS *client) {
static RETCODE nogvl_dbresults(tinytds_client_wrapper *cwrap) {
int retcode = FAIL;
nogvl_setup(client);
retcode = NOGVL_DBCALL(dbresults, client);
nogvl_cleanup(client);
raise_if_client_closed(cwrap);
nogvl_setup(cwrap->client);
retcode = NOGVL_DBCALL(dbresults, cwrap->client);
nogvl_cleanup(cwrap->client);
return retcode;
}

static RETCODE nogvl_dbnextrow(DBPROCESS * client) {
static RETCODE nogvl_dbnextrow(tinytds_client_wrapper *cwrap) {
int retcode = FAIL;
nogvl_setup(client);
retcode = NOGVL_DBCALL(dbnextrow, client);
nogvl_cleanup(client);
raise_if_client_closed(cwrap);
nogvl_setup(cwrap->client);
retcode = NOGVL_DBCALL(dbnextrow, cwrap->client);
nogvl_cleanup(cwrap->client);
return retcode;
}

Expand All @@ -164,7 +173,7 @@ static RETCODE rb_tinytds_result_dbresults_retcode(VALUE self) {
GET_RESULT_WRAPPER(self);
ruby_rc = rb_ary_entry(rwrap->dbresults_retcodes, rwrap->number_of_results);
if (NIL_P(ruby_rc)) {
db_rc = nogvl_dbresults(rwrap->client);
db_rc = nogvl_dbresults(rwrap->cwrap);
ruby_rc = INT2FIX(db_rc);
rb_ary_store(rwrap->dbresults_retcodes, rwrap->number_of_results, ruby_rc);
} else {
Expand All @@ -173,35 +182,33 @@ static RETCODE rb_tinytds_result_dbresults_retcode(VALUE self) {
return db_rc;
}

static RETCODE rb_tinytds_result_ok_helper(DBPROCESS *client) {
GET_CLIENT_USERDATA(client);
if (userdata->dbsqlok_sent == 0) {
userdata->dbsqlok_retcode = nogvl_dbsqlok(client);
static RETCODE rb_tinytds_result_ok_helper(tinytds_client_wrapper *cwrap) {
if (cwrap->userdata->dbsqlok_sent == 0) {
cwrap->userdata->dbsqlok_retcode = nogvl_dbsqlok(cwrap);
}
return userdata->dbsqlok_retcode;
return cwrap->userdata->dbsqlok_retcode;
}

static void rb_tinytds_result_exec_helper(DBPROCESS *client) {
RETCODE dbsqlok_rc = rb_tinytds_result_ok_helper(client);
GET_CLIENT_USERDATA(client);
static void rb_tinytds_result_exec_helper(tinytds_client_wrapper *cwrap) {
RETCODE dbsqlok_rc = rb_tinytds_result_ok_helper(cwrap);
if (dbsqlok_rc == SUCCEED) {
/*
This is to just process each result set. Commands such as backup and
restore are not done when the first result set is returned, so we need to
exhaust the result sets before it is complete.
*/
while (nogvl_dbresults(client) == SUCCEED) {
while (nogvl_dbresults(cwrap) == SUCCEED) {
/*
If we don't loop through each row for calls to TinyTds::Result.do that
actually do return result sets, we will trigger error 20019 about trying
to execute a new command with pending results. Oh well.
*/
while (dbnextrow(client) != NO_MORE_ROWS);
while (dbnextrow(cwrap->client) != NO_MORE_ROWS);
}
}
dbcancel(client);
userdata->dbcancel_sent = 1;
userdata->dbsql_sent = 0;
dbcancel(cwrap->client);
cwrap->userdata->dbcancel_sent = 1;
cwrap->userdata->dbsql_sent = 0;
}

static VALUE rb_tinytds_result_fetch_row(VALUE self, ID timezone, int symbolize_keys, int as_array) {
Expand Down Expand Up @@ -368,7 +375,7 @@ static VALUE rb_tinytds_result_fields(VALUE self) {
RETCODE dbsqlok_rc, dbresults_rc;
VALUE fields_processed;
GET_RESULT_WRAPPER(self);
dbsqlok_rc = rb_tinytds_result_ok_helper(rwrap->client);
dbsqlok_rc = rb_tinytds_result_ok_helper(rwrap->cwrap);
dbresults_rc = rb_tinytds_result_dbresults_retcode(self);
fields_processed = rb_ary_entry(rwrap->fields_processed, rwrap->number_of_results);
if ((dbsqlok_rc == SUCCEED) && (dbresults_rc == SUCCEED) && (fields_processed == Qnil)) {
Expand Down Expand Up @@ -441,7 +448,7 @@ static VALUE rb_tinytds_result_each(int argc, VALUE * argv, VALUE self) {
if (NIL_P(rwrap->results)) {
RETCODE dbsqlok_rc, dbresults_rc;
rwrap->results = rb_ary_new();
dbsqlok_rc = rb_tinytds_result_ok_helper(rwrap->client);
dbsqlok_rc = rb_tinytds_result_ok_helper(rwrap->cwrap);
dbresults_rc = rb_tinytds_result_dbresults_retcode(self);
while ((dbsqlok_rc == SUCCEED) && (dbresults_rc == SUCCEED)) {
int has_rows = (DBROWS(rwrap->client) == SUCCEED) ? 1 : 0;
Expand All @@ -451,7 +458,7 @@ static VALUE rb_tinytds_result_each(int argc, VALUE * argv, VALUE self) {
/* Create rows for this result set. */
unsigned long rowi = 0;
VALUE result = rb_ary_new();
while (nogvl_dbnextrow(rwrap->client) != NO_MORE_ROWS) {
while (nogvl_dbnextrow(rwrap->cwrap) != NO_MORE_ROWS) {
VALUE row = rb_tinytds_result_fetch_row(self, timezone, symbolize_keys, as_array);
if (cache_rows)
rb_ary_store(result, rowi, row);
Expand Down Expand Up @@ -484,7 +491,7 @@ static VALUE rb_tinytds_result_each(int argc, VALUE * argv, VALUE self) {
} else {
// If we do not find results, side step the rb_tinytds_result_dbresults_retcode helper and
// manually populate its memoized array while nullifing any memoized fields too before loop.
dbresults_rc = nogvl_dbresults(rwrap->client);
dbresults_rc = nogvl_dbresults(rwrap->cwrap);
rb_ary_store(rwrap->dbresults_retcodes, rwrap->number_of_results, INT2FIX(dbresults_rc));
rb_ary_store(rwrap->fields_processed, rwrap->number_of_results, Qnil);
}
Expand All @@ -506,7 +513,7 @@ static VALUE rb_tinytds_result_cancel(VALUE self) {
GET_RESULT_WRAPPER(self);
userdata = (tinytds_client_userdata *)dbgetuserdata(rwrap->client);
if (rwrap->client && !userdata->dbcancel_sent) {
rb_tinytds_result_ok_helper(rwrap->client);
rb_tinytds_result_ok_helper(rwrap->cwrap);
dbcancel(rwrap->client);
userdata->dbcancel_sent = 1;
userdata->dbsql_sent = 0;
Expand All @@ -517,7 +524,7 @@ static VALUE rb_tinytds_result_cancel(VALUE self) {
static VALUE rb_tinytds_result_do(VALUE self) {
GET_RESULT_WRAPPER(self);
if (rwrap->client) {
rb_tinytds_result_exec_helper(rwrap->client);
rb_tinytds_result_exec_helper(rwrap->cwrap);
return LONG2NUM((long)dbcount(rwrap->client));
} else {
return Qnil;
Expand Down Expand Up @@ -547,12 +554,12 @@ static VALUE rb_tinytds_result_insert(VALUE self) {
GET_RESULT_WRAPPER(self);
if (rwrap->client) {
VALUE identity = Qnil;
rb_tinytds_result_exec_helper(rwrap->client);
rb_tinytds_result_exec_helper(rwrap->cwrap);
dbcmd(rwrap->client, rwrap->cwrap->identity_insert_sql);
if (nogvl_dbsqlexec(rwrap->client) != FAIL
&& nogvl_dbresults(rwrap->client) != FAIL
if (nogvl_dbsqlexec(rwrap->cwrap) != FAIL
&& nogvl_dbresults(rwrap->cwrap) != FAIL
&& DBROWS(rwrap->client) != FAIL) {
while (nogvl_dbnextrow(rwrap->client) != NO_MORE_ROWS) {
while (nogvl_dbnextrow(rwrap->cwrap) != NO_MORE_ROWS) {
int col = 1;
BYTE *data = dbdata(rwrap->client, col);
DBINT data_len = dbdatlen(rwrap->client, col);
Expand Down
72 changes: 72 additions & 0 deletions test/result_test.rb
Original file line number Diff line number Diff line change
Expand Up @@ -699,6 +699,78 @@ class ResultTest < TinyTds::TestCase
assert_equal 50000, e.db_error_number
end
end

it 'deals with closed client before calling each' do
client = new_connection
result = client.execute("SELECT [id] FROM [datatypes]")
client.close

assert_raise_tinytds_error(-> { result.to_a }) do |e|
assert_equal "closed connection", e.message
end

assert_new_connections_work
end

it 'deals with closed client mid-iteration' do
client = new_connection

action = lambda do
result = client.execute("SELECT [id] FROM [datatypes]")
result.each_with_index do |_row, index|
assert index <= 5

if index == 5
client.close
end
end
end

assert_raise_tinytds_error(action) do |e|
assert_equal "closed connection", e.message
end

assert_new_connections_work
end

it 'deals with dead connection before iterating over results' do
init_toxiproxy

client = new_connection port: 1234, host: ENV['TOXIPROXY_HOST']

Toxiproxy[:sqlserver_test].down do
result = client.execute("SELECT [id] FROM [datatypes]")

assert_raise_tinytds_error(-> { result.to_a }) do |e|
assert_equal "DBPROCESS is dead or not enabled", e.message
end
end
end

it 'deals with dead connection mid-iteration' do
init_toxiproxy

client = new_connection port: 1234, host: ENV['TOXIPROXY_HOST']
action = lambda do
result = client.execute("SELECT [id] FROM [datatypes]")
result.each_with_index do |_row, index|
assert index <= 5

if index == 5
puts "sleeping to wait for toxiproxy"
sleep 2
end
end
end

Toxiproxy[:sqlserver_test].toxic(:timeout, timeout: 1000).apply do
assert_raise_tinytds_error(action) do |e|
assert_equal "DBPROCESS is dead or not enabled", e.message
end
end

assert_new_connections_work
end
end
end

Expand Down

0 comments on commit 2dd9313

Please sign in to comment.