From c115fa1a42db31d645bc0c561c16ada3b52a806e Mon Sep 17 00:00:00 2001
From: Sayan Nandan <nandansayan@outlook.com>
Date: Mon, 11 Mar 2024 18:43:39 +0530
Subject: [PATCH] Version handshake

---
 src/io/aio.rs             | 4 ++--
 src/io/sync.rs            | 4 ++--
 src/protocol/handshake.rs | 8 ++++++--
 src/protocol/mod.rs       | 1 +
 4 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/src/io/aio.rs b/src/io/aio.rs
index 34b4005..0ea9f54 100644
--- a/src/io/aio.rs
+++ b/src/io/aio.rs
@@ -86,7 +86,7 @@ impl Config {
     /// Establish an async connection to the database using the current configuration
     pub async fn connect_async(&self) -> ClientResult<ConnectionAsync> {
         let mut tcpstream = TcpStream::connect((self.host(), self.port())).await?;
-        let handshake = ClientHandshake::new(self);
+        let handshake = ClientHandshake::new_v1(self);
         tcpstream.write_all(handshake.inner()).await?;
         let mut resp = [0u8; 4];
         tcpstream.read_exact(&mut resp).await?;
@@ -121,7 +121,7 @@ impl Config {
             .await
             .map_err(|e| ConnectionSetupError::Other(format!("TLS handshake failed: {e}")))?;
         // handshake
-        let handshake = ClientHandshake::new(self);
+        let handshake = ClientHandshake::new_v1(self);
         stream.write_all(handshake.inner()).await?;
         let mut resp = [0u8; 4];
         stream.read_exact(&mut resp).await?;
diff --git a/src/io/sync.rs b/src/io/sync.rs
index aae4647..6329a5d 100644
--- a/src/io/sync.rs
+++ b/src/io/sync.rs
@@ -87,7 +87,7 @@ impl Config {
     /// Establish a connection to the database using the current configuration
     pub fn connect(&self) -> ClientResult<Connection> {
         let mut tcpstream = TcpStream::connect((self.host(), self.port()))?;
-        let handshake = ClientHandshake::new(self);
+        let handshake = ClientHandshake::new_v1(self);
         tcpstream.write_all(handshake.inner())?;
         let mut resp = [0u8; 4];
         tcpstream.read_exact(&mut resp)?;
@@ -113,7 +113,7 @@ impl Config {
             })?
             .connect(self.host(), stream)
             .map_err(|e| ConnectionSetupError::Other(format!("TLS handshake failed: {e}")))?;
-        let handshake = ClientHandshake::new(self);
+        let handshake = ClientHandshake::new_v1(self);
         stream.write_all(handshake.inner())?;
         let mut resp = [0u8; 4];
         stream.read_exact(&mut resp)?;
diff --git a/src/protocol/handshake.rs b/src/protocol/handshake.rs
index d885aed..51b8fcc 100644
--- a/src/protocol/handshake.rs
+++ b/src/protocol/handshake.rs
@@ -21,9 +21,13 @@ use crate::{
 
 pub struct ClientHandshake(Box<[u8]>);
 impl ClientHandshake {
-    pub(crate) fn new(cfg: &Config) -> Self {
+    const HANDSHAKE_PROTO_V1: [u8; 6] = [b'H', 0, 0, 0, 0, 0];
+    pub(crate) fn new_v1(cfg: &Config) -> Self {
+        Self::_new(Self::HANDSHAKE_PROTO_V1, cfg)
+    }
+    fn _new(hs: [u8; 6], cfg: &Config) -> Self {
         let mut v = Vec::with_capacity(6 + cfg.username().len() + cfg.password().len() + 5);
-        v.extend(b"H\x00\x00\x00\x00\x00");
+        v.extend(hs);
         pushlen!(v, cfg.username().len());
         pushlen!(v, cfg.password().len());
         v.extend(cfg.username().as_bytes());
diff --git a/src/protocol/mod.rs b/src/protocol/mod.rs
index e102b4b..8756c33 100644
--- a/src/protocol/mod.rs
+++ b/src/protocol/mod.rs
@@ -203,6 +203,7 @@ impl<'a> Decoder<'a> {
         }
         let lf = self._creq(b'\n');
         self._cursor_incr_if(lf);
+        // FIXME(@ohsayan): the below is not exactly necessary and we can actually remove this if it complicates state management
         okay &= !(lf & (self._cursor() == meta.start));
         if okay & lf {
             let start = meta.start;