From ccbca66e7ab0a9a0a90f1e734d193e0f5037669d Mon Sep 17 00:00:00 2001
From: Lev Kokotov <levkk@users.noreply.github.com>
Date: Mon, 9 May 2022 11:09:22 -0500
Subject: [PATCH] Poorly behaved client fix (#65)

* Poorly behaved client fix

* yes officer

* fix tests

* no useless rescue

* Looks ok
---
 src/client.rs           |  6 ++++--
 tests/ruby/Gemfile      |  1 +
 tests/ruby/Gemfile.lock | 21 +++++++++++++++++++++
 tests/ruby/tests.rb     | 17 +++++++++++++++++
 4 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/src/client.rs b/src/client.rs
index 87f6215..976121b 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -359,7 +359,8 @@ impl Client {
                             // Clean up the server and re-use it.
                             // This prevents connection thrashing by bad clients.
                             if server.in_transaction() {
-                                server.query("ROLLBACK; DISCARD ALL;").await?;
+                                server.query("ROLLBACK").await?;
+                                server.query("DISCARD ALL").await?;
                             }
 
                             return Err(err);
@@ -429,7 +430,8 @@ impl Client {
                         // Pgbouncer closes the connection which leads to
                         // connection thrashing when clients misbehave.
                         if server.in_transaction() {
-                            server.query("ROLLBACK; DISCARD ALL;").await?;
+                            server.query("ROLLBACK").await?;
+                            server.query("DISCARD ALL").await?;
                         }
 
                         return Ok(());
diff --git a/tests/ruby/Gemfile b/tests/ruby/Gemfile
index 94fb8c3..05684c9 100644
--- a/tests/ruby/Gemfile
+++ b/tests/ruby/Gemfile
@@ -2,3 +2,4 @@ source "https://rubygems.org"
 
 gem "pg"
 gem "activerecord"
+gem "rubocop"
diff --git a/tests/ruby/Gemfile.lock b/tests/ruby/Gemfile.lock
index 7b7d448..045e2f6 100644
--- a/tests/ruby/Gemfile.lock
+++ b/tests/ruby/Gemfile.lock
@@ -11,13 +11,33 @@ GEM
       i18n (>= 1.6, < 2)
       minitest (>= 5.1)
       tzinfo (~> 2.0)
+    ast (2.4.2)
     concurrent-ruby (1.1.9)
     i18n (1.10.0)
       concurrent-ruby (~> 1.0)
     minitest (5.15.0)
+    parallel (1.22.1)
+    parser (3.1.2.0)
+      ast (~> 2.4.1)
     pg (1.3.2)
+    rainbow (3.1.1)
+    regexp_parser (2.3.1)
+    rexml (3.2.5)
+    rubocop (1.29.0)
+      parallel (~> 1.10)
+      parser (>= 3.1.0.0)
+      rainbow (>= 2.2.2, < 4.0)
+      regexp_parser (>= 1.8, < 3.0)
+      rexml (>= 3.2.5, < 4.0)
+      rubocop-ast (>= 1.17.0, < 2.0)
+      ruby-progressbar (~> 1.7)
+      unicode-display_width (>= 1.4.0, < 3.0)
+    rubocop-ast (1.17.0)
+      parser (>= 3.1.1.0)
+    ruby-progressbar (1.11.0)
     tzinfo (2.0.4)
       concurrent-ruby (~> 1.0)
+    unicode-display_width (2.1.0)
 
 PLATFORMS
   x86_64-linux
@@ -25,6 +45,7 @@ PLATFORMS
 DEPENDENCIES
   activerecord
   pg
+  rubocop
 
 BUNDLED WITH
    2.3.7
diff --git a/tests/ruby/tests.rb b/tests/ruby/tests.rb
index 622ac01..5b032a9 100644
--- a/tests/ruby/tests.rb
+++ b/tests/ruby/tests.rb
@@ -1,6 +1,9 @@
 # frozen_string_literal: true
 
 require 'active_record'
+require 'pg'
+
+$stdout.sync = true
 
 # Uncomment these two to see all queries.
 # ActiveRecord.verbose_query_logs = true
@@ -110,3 +113,17 @@ begin
 rescue ActiveRecord::StatementInvalid
   puts 'OK'
 end
+
+# Test evil clients
+def poorly_behaved_client
+  conn = PG::connect("postgres://sharding_user:sharding_user@127.0.0.1:6432/rails_dev")
+  conn.async_exec 'BEGIN'
+  conn.async_exec 'SELECT 1'
+
+  conn.close
+  puts 'Bad client ok'
+end
+
+25.times do
+  poorly_behaved_client
+end