From 505f59b3549a14842a95e164848b8652c98f34ce Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jind=C5=99ich=20Moravec?= <xmorave2@fi.muni.cz>
Date: Sat, 23 Dec 2023 01:27:15 +0100
Subject: [PATCH] fix(proto): flush written startup message

---
 proto/src/handshake/client.rs   |  1 +
 proto/src/handshake/errors.rs   |  2 ++
 proto/src/handshake/request.rs  |  1 +
 proto/src/handshake/response.rs | 17 ++++++++++-------
 4 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/proto/src/handshake/client.rs b/proto/src/handshake/client.rs
index 9226b2a..aa6577b 100644
--- a/proto/src/handshake/client.rs
+++ b/proto/src/handshake/client.rs
@@ -14,6 +14,7 @@ pub async fn do_client_handshake(
 ) -> Result<HandshakeResponse, ClientHandshakeError> {
     let startup_message: StartupMessageData = request.into();
     writer.write_startup_message(startup_message).await?;
+    writer.flush().await?;
 
     let auth = reader.read_proto().await?;
     if !matches!(auth, BackendMessage::AuthenticationOk(AuthenticationOkData { status: 0 })) {
diff --git a/proto/src/handshake/errors.rs b/proto/src/handshake/errors.rs
index d2bc1bd..b5a31ed 100644
--- a/proto/src/handshake/errors.rs
+++ b/proto/src/handshake/errors.rs
@@ -11,6 +11,8 @@ pub enum ClientHandshakeError {
     UnexpectedResponse,
     #[error("unexpected auth response")]
     UnexpectedAuthResponse(BackendMessage),
+    #[error("socket communication failed")]
+    Io(#[from] io::Error),
     #[error("writing message to socket failed")]
     Write(#[from] ProtoWriteError),
     #[error("reading message from socket failed")]
diff --git a/proto/src/handshake/request.rs b/proto/src/handshake/request.rs
index ec3e123..1fc097f 100644
--- a/proto/src/handshake/request.rs
+++ b/proto/src/handshake/request.rs
@@ -1,6 +1,7 @@
 use crate::message::primitive::pgstring::PgString;
 use crate::message::special::StartupMessageData;
 
+#[derive(Debug)]
 pub struct HandshakeRequest {
     version: i32,
     parameters: Vec<(PgString, PgString)>,
diff --git a/proto/src/handshake/response.rs b/proto/src/handshake/response.rs
index 84a908f..38579b6 100644
--- a/proto/src/handshake/response.rs
+++ b/proto/src/handshake/response.rs
@@ -1,6 +1,7 @@
 use crate::handshake::errors::ClientHandshakeError;
 use crate::message::backend::{BackendKeyDataData, BackendMessage, ParameterStatusData};
 
+#[derive(Debug)]
 pub struct HandshakeResponse {
     pub version: String,
     pub process_id: i32,
@@ -41,20 +42,22 @@ impl TryFrom<&[BackendMessage]> for HandshakeResponse {
         }
 
         match (version, process_id, secret_key) {
-            (Some(version), Some(process_id), Some(secret_key)) => {
-                Ok(Self { version, process_id, secret_key })
-            }
+            (Some(version), Some(process_id), Some(secret_key)) => Ok(Self {
+                version,
+                process_id,
+                secret_key,
+            }),
             _ => Err(ClientHandshakeError::UnexpectedResponse),
         }
     }
 }
 
-impl From<&HandshakeResponse> for Vec<BackendMessage> {
-    fn from(response: &HandshakeResponse) -> Self {
+impl From<HandshakeResponse> for Vec<BackendMessage> {
+    fn from(response: HandshakeResponse) -> Self {
         vec![
             BackendMessage::ParameterStatus(ParameterStatusData {
                 name: "server_version".into(),
-                value: response.version.clone().into(),
+                value: response.version.into(),
             }),
             BackendMessage::BackendKeyData(BackendKeyDataData {
                 process: response.process_id,
@@ -62,4 +65,4 @@ impl From<&HandshakeResponse> for Vec<BackendMessage> {
             }),
         ]
     }
-}
\ No newline at end of file
+}