[proxy] Pass PostgreSQL version to client

Fixes #779

zenith_utils/src/postgres_backend.rs

@@ -3,7 +3,9 @@
 //! implementation determining how to process the queries. Currently its API
 //! is rather narrow, but we can extend it once required.
 
-use crate::pq_proto::{BeMessage, FeMessage, FeStartupMessage, StartupRequestCode};
+use crate::pq_proto::{
+    BeMessage, BeParameterStatusMessage, FeMessage, FeStartupMessage, StartupRequestCode,
+};
 use crate::sock_split::{BidiStream, ReadStream, WriteStream};
 use anyhow::{anyhow, bail, ensure, Result};
 use bytes::{Bytes, BytesMut};
@@ -355,11 +357,9 @@ impl PostgresBackend {
 
                         match self.auth_type {
                             AuthType::Trust => {
-                                self.write_message_noflush(&BeMessage::AuthenticationOk)?;
-                                // psycopg2 will not connect if client_encoding is not
-                                // specified by the server
-                                self.write_message_noflush(&BeMessage::ParameterStatus)?;
-                                self.write_message(&BeMessage::ReadyForQuery)?;
+                                self.write_message_noflush(&BeMessage::AuthenticationOk)?
+                                    .write_message_noflush(&BeParameterStatusMessage::encoding())?
+                                    .write_message(&BeMessage::ReadyForQuery)?;
                                 self.state = ProtoState::Established;
                             }
                             AuthType::MD5 => {
@@ -410,11 +410,9 @@ impl PostgresBackend {
                         }
                     }
                 }
-                self.write_message_noflush(&BeMessage::AuthenticationOk)?;
-                // psycopg2 will not connect if client_encoding is not
-                // specified by the server
-                self.write_message_noflush(&BeMessage::ParameterStatus)?;
-                self.write_message(&BeMessage::ReadyForQuery)?;
+                self.write_message_noflush(&BeMessage::AuthenticationOk)?
+                    .write_message_noflush(&BeParameterStatusMessage::encoding())?
+                    .write_message(&BeMessage::ReadyForQuery)?;
                 self.state = ProtoState::Established;
             }
 

zenith_utils/src/pq_proto.rs

@@ -352,7 +352,7 @@ pub enum BeMessage<'a> {
     EncryptionResponse(bool),
     NoData,
     ParameterDescription,
-    ParameterStatus,
+    ParameterStatus(BeParameterStatusMessage<'a>),
     ParseComplete,
     ReadyForQuery,
     RowDescription(&'a [RowDescriptor<'a>]),
@@ -361,6 +361,18 @@ pub enum BeMessage<'a> {
     KeepAlive(WalSndKeepAlive),
 }
 
+#[derive(Debug)]
+pub enum BeParameterStatusMessage<'a> {
+    Encoding(&'a str),
+    ServerVersion(&'a str),
+}
+
+impl BeParameterStatusMessage<'static> {
+    pub fn encoding() -> BeMessage<'static> {
+        BeMessage::ParameterStatus(Self::Encoding("UTF8"))
+    }
+}
+
 // One row desciption in RowDescription packet.
 #[derive(Debug)]
 pub struct RowDescriptor<'a> {
@@ -665,12 +677,23 @@ impl<'a> BeMessage<'a> {
                 buf.put_u8(response);
             }
 
-            BeMessage::ParameterStatus => {
+            BeMessage::ParameterStatus(param) => {
+                use std::io::{IoSlice, Write};
+                use BeParameterStatusMessage::*;
+
+                let [name, value] = match param {
+                    Encoding(name) => [b"client_encoding", name.as_bytes()],
+                    ServerVersion(version) => [b"server_version", version.as_bytes()],
+                };
+
+                // Parameter names and values are passed as null-terminated strings
+                let iov = &mut [name, b"\0", value, b"\0"].map(IoSlice::new);
+                let mut buffer = [0u8; 64]; // this should be enough
+                let cnt = buffer.as_mut().write_vectored(iov).unwrap();
+
                 buf.put_u8(b'S');
-                // parameter names and values are specified by null terminated strings
-                const PARAM_NAME_VALUE: &[u8] = b"client_encoding\0UTF8\0";
                 write_body(buf, |buf| {
-                    buf.put_slice(PARAM_NAME_VALUE);
+                    buf.put_slice(&buffer[..cnt]);
                     Ok::<_, io::Error>(())
                 })
                 .unwrap();