From fceb428878087dbf1867dcd491fa6ed555c15470 Mon Sep 17 00:00:00 2001
From: Bojan Serafimov <bojan.serafimov7@gmail.com>
Date: Tue, 18 Jan 2022 15:30:33 -0500
Subject: [PATCH] Fix logic

---
 proxy/src/proxy.rs | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/proxy/src/proxy.rs b/proxy/src/proxy.rs
index ebdee14505..b74218c603 100644
--- a/proxy/src/proxy.rs
+++ b/proxy/src/proxy.rs
@@ -164,14 +164,22 @@ impl ProxyConnection {
     fn handle_startup(&mut self) -> anyhow::Result<Option<(String, String)>> {
         let have_tls = self.pgb.tls_config.is_some();
         let mut encrypted = false;
+        let mut received_something = false;
 
         loop {
             let msg = match self.pgb.read_message()? {
                 Some(Fe::StartupPacket(msg)) => msg,
-                None => return Ok(None), // Probably load balancer health check
+                None => {
+                    if received_something {
+                        bail!("connection is lost");
+                    } else {
+                        return Ok(None); // Probably load balancer health check
+                    }
+                }
                 bad => bail!("unexpected message type: {:?}", bad),
             };
             println!("got message: {:?}", msg);
+            received_something = true;
 
             match msg {
                 FeStartupPacket::GssEncRequest => {