From 560f088f05e3132201d51e5625e16ef13a0fe16b Mon Sep 17 00:00:00 2001
From: Stas Kelvich <stas.kelvich@gmail.com>
Date: Wed, 30 Jun 2021 12:35:07 +0300
Subject: [PATCH] fix query_string parsing for mgmt callback

---
 proxy/src/mgmt.rs | 17 ++++++++---------
 1 file changed, 8 insertions(+), 9 deletions(-)

diff --git a/proxy/src/mgmt.rs b/proxy/src/mgmt.rs
index 6fab902c8d..c91605951f 100644
--- a/proxy/src/mgmt.rs
+++ b/proxy/src/mgmt.rs
@@ -78,15 +78,14 @@ impl postgres_backend::Handler for MgmtHandler {
         pgb: &mut PostgresBackend,
         query_string: Bytes,
     ) -> anyhow::Result<()> {
-        let (_, query_string) = query_string
-            .split_last()
-            .ok_or_else(|| anyhow::Error::msg("protocol violation"))?;
+        let mut query_string = query_string.to_vec();
+        if let Some(ch) = query_string.last() {
+            if *ch == 0 {
+                query_string.pop();
+            }
+        }
 
-        let (_, query_string) = query_string
-            .split_last()
-            .ok_or_else(|| anyhow::Error::msg("protocol violation"))?;
-
-        println!("Got mgmt query: {:?}", query_string);
+        println!("Got mgmt query: '{}'", std::str::from_utf8(&query_string)?);
 
         let resp: PsqlSessionResponse = serde_json::from_slice(&query_string)?;
 
@@ -94,7 +93,7 @@ impl postgres_backend::Handler for MgmtHandler {
 
         let sender = waiters
             .get(&resp.session_id)
-            .ok_or_else(|| anyhow::Error::msg("psql session_id is not found"))?;
+            .ok_or_else(|| anyhow::Error::msg("psql_session_id is not found"))?;
 
         match resp.result {
             PsqlSessionResult::Success(db_info) => {