From 3e1ac1c74b18e0a6f9fa7090bd359f026715a762 Mon Sep 17 00:00:00 2001
From: Sean Smith <scsmithr@gmail.com>
Date: Fri, 16 Sep 2022 09:58:29 -0500
Subject: [PATCH] feat: Implement parameter status message for pgsrv (#102)

* feat: Implement parameter status message for pgsrv

Sends parameter status messages when starting up a new connection.

Currently only sends `server_version`. This was enough to get pgJDBC to startup
properly. This was an attempt to try to get
benchbase (https://github.com/GlareDB/benchbase/tree/glaredb) running against a
glaredb server.

I have not gotten benchbase to successfully run a benchmark, see
https://github.com/GlareDB/glaredb/issues/101.

* Include both byte and char reps in error string
---
 crates/pgsrv/src/codec.rs    |  7 ++++++-
 crates/pgsrv/src/errors.rs   |  5 ++++-
 crates/pgsrv/src/handler.rs  | 21 ++++++++++++++++++++-
 crates/pgsrv/src/messages.rs |  1 +
 4 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/crates/pgsrv/src/codec.rs b/crates/pgsrv/src/codec.rs
index d7d406dd3..938458ab8 100644
--- a/crates/pgsrv/src/codec.rs
+++ b/crates/pgsrv/src/codec.rs
@@ -180,6 +180,7 @@ impl Encoder<BackendMessage> for PgCodec {
             BackendMessage::AuthenticationOk => b'R',
             BackendMessage::AuthenticationCleartextPassword => b'R',
             BackendMessage::EmptyQueryResponse => b'I',
+            BackendMessage::ParameterStatus { .. } => b'S',
             BackendMessage::ReadyForQuery(_) => b'Z',
             BackendMessage::CommandComplete { .. } => b'C',
             BackendMessage::RowDescription(_) => b'T',
@@ -197,6 +198,10 @@ impl Encoder<BackendMessage> for PgCodec {
             BackendMessage::AuthenticationOk => dst.put_i32(0),
             BackendMessage::AuthenticationCleartextPassword => dst.put_i32(3),
             BackendMessage::EmptyQueryResponse => (),
+            BackendMessage::ParameterStatus { key, val } => {
+                dst.put_cstring(&key);
+                dst.put_cstring(&val);
+            }
             BackendMessage::ReadyForQuery(status) => match status {
                 TransactionStatus::Idle => dst.put_u8(b'I'),
                 TransactionStatus::InBlock => dst.put_u8(b'T'),
@@ -279,7 +284,7 @@ impl Decoder for PgCodec {
         let msg_len = i32::from_be_bytes(src[1..5].try_into().unwrap()) as usize;
 
         // Not enough bytes to read the full message yet.
-        if src.len() < msg_len - 1 {
+        if src.len() < msg_len {
             src.reserve(msg_len - src.len());
             return Ok(None);
         }
diff --git a/crates/pgsrv/src/errors.rs b/crates/pgsrv/src/errors.rs
index dc70d9c30..1b02a88e2 100644
--- a/crates/pgsrv/src/errors.rs
+++ b/crates/pgsrv/src/errors.rs
@@ -21,7 +21,10 @@ pub enum PgSrvError {
     #[error("missing null byte")]
     MissingNullByte,
 
-    #[error("invalid message type: {0}")]
+    /// We've received an unexpected message identifier from the frontend.
+    /// Includes the char representation to allow for easy cross referencing
+    /// with the Postgres message format documentation.
+    #[error("invalid message type: byte={0}, char={}", *.0 as char)]
     InvalidMsgType(u8),
 
     #[error(transparent)]
diff --git a/crates/pgsrv/src/handler.rs b/crates/pgsrv/src/handler.rs
index 58f781d81..a59e22148 100644
--- a/crates/pgsrv/src/handler.rs
+++ b/crates/pgsrv/src/handler.rs
@@ -15,6 +15,15 @@ use std::collections::HashMap;
 use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt};
 use tracing::trace;
 
+/// Default parameters to send to the frontend on startup. Existing postgres
+/// drivers may expect these in the server response on startup.
+///
+/// See https://www.postgresql.org/docs/current/runtime-config-preset.html for
+/// other parameters we may want to provide.
+///
+/// Some parameters  will eventually be provided at runtime.
+const DEFAULT_READ_ONLY_PARAMS: &[(&str, &str)] = &[("server_version", "0.0.0")];
+
 /// A wrapper around a sqlengine that implements the Postgres frontend/backend
 /// protocol.
 pub struct Handler {
@@ -56,7 +65,7 @@ impl Handler {
                 }
             }
             StartupMessage::CancelRequest { .. } => {
-                trace!("recieved cancel request");
+                trace!("received cancel request");
                 // TODO: Properly handle requests to cancel sessions.
 
                 // Note that we should not respond to this request.
@@ -103,6 +112,16 @@ impl Handler {
             }
         };
 
+        // Send server parameters.
+        for (key, val) in DEFAULT_READ_ONLY_PARAMS {
+            framed
+                .send(BackendMessage::ParameterStatus {
+                    key: key.to_string(),
+                    val: val.to_string(),
+                })
+                .await?;
+        }
+
         let cs = ClientSession::new(sess, framed);
         cs.run().await
     }
diff --git a/crates/pgsrv/src/messages.rs b/crates/pgsrv/src/messages.rs
index 58e987d32..ad11d7e23 100644
--- a/crates/pgsrv/src/messages.rs
+++ b/crates/pgsrv/src/messages.rs
@@ -46,6 +46,7 @@ pub enum BackendMessage {
     NoticeResponse(NoticeResponse),
     AuthenticationOk,
     AuthenticationCleartextPassword,
+    ParameterStatus { key: String, val: String },
     EmptyQueryResponse,
     ReadyForQuery(TransactionStatus),
     CommandComplete { tag: String },