fix(postgres): return UnexpectedEof when server closes connection at SSLRequest

sqlx-postgres/src/connection/tls.rs

@@ -1,3 +1,5 @@
+use std::io;
+
 use crate::error::Error;
 use crate::net::tls::{self, TlsConfig};
 use crate::net::{Socket, SocketIntoBox, WithSocket};
@@ -88,7 +90,10 @@ async fn request_upgrade(
 
     let mut response = [0u8];
 
-    socket.read(&mut &mut response[..]).await?;
+    let n = socket.read(&mut &mut response[..]).await?;
+    if n == 0 {
+        return Err(io::Error::from(io::ErrorKind::UnexpectedEof).into());
+    }
 
     match response[0] {
         b'S' => {
@@ -107,3 +112,133 @@ async fn request_upgrade(
         )),
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use sqlx_core::io::ReadBuf;
+    use std::io;
+    use std::pin::pin;
+    use std::task::{Context, Poll, Waker};
+
+    struct MockSocket {
+        read_byte: Option<u8>,
+    }
+
+    impl Socket for MockSocket {
+        fn try_read(&mut self, buf: &mut dyn ReadBuf) -> io::Result<usize> {
+            match self.read_byte {
+                Some(b) => {
+                    buf.put_slice(&[b]);
+                    Ok(1)
+                }
+                None => Ok(0), // EOF
+            }
+        }
+
+        fn try_write(&mut self, buf: &[u8]) -> io::Result<usize> {
+            Ok(buf.len())
+        }
+
+        fn poll_read_ready(&mut self, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+            Poll::Ready(Ok(()))
+        }
+
+        fn poll_write_ready(&mut self, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+            Poll::Ready(Ok(()))
+        }
+
+        fn poll_shutdown(&mut self, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+            Poll::Ready(Ok(()))
+        }
+    }
+
+    /// Polls a future that is guaranteed to complete immediately on its first poll.
+    ///
+    /// This safe test runner uses `Waker::noop()` and `std::pin::pin!` without needing
+    /// a full runtime or unsafe waker implementations.
+    fn poll_immediate<F: std::future::Future>(fut: F) -> F::Output {
+        let mut fut = pin!(fut);
+        let mut cx = Context::from_waker(Waker::noop());
+        match fut.as_mut().poll(&mut cx) {
+            Poll::Ready(val) => val,
+            Poll::Pending => panic!("test future was pending unexpectedly"),
+        }
+    }
+
+    #[test]
+    fn test_request_upgrade_eof() {
+        poll_immediate(async {
+            let mut socket = MockSocket { read_byte: None };
+            let options = PgConnectOptions::new();
+            let err = request_upgrade(&mut socket, &options).await.unwrap_err();
+
+            match err {
+                Error::Io(io_err) => {
+                    assert_eq!(io_err.kind(), io::ErrorKind::UnexpectedEof);
+                }
+                other => panic!("expected Error::Io(UnexpectedEof), but got: {:?}", other),
+            }
+        });
+    }
+
+    #[test]
+    fn test_request_upgrade_wire_null_byte() {
+        poll_immediate(async {
+            let mut socket = MockSocket {
+                read_byte: Some(0x00),
+            };
+            let options = PgConnectOptions::new();
+            let err = request_upgrade(&mut socket, &options).await.unwrap_err();
+
+            match err {
+                Error::Protocol(msg) => {
+                    assert!(msg.contains("unexpected response from SSLRequest: 0x00"));
+                }
+                other => panic!("expected Error::Protocol, but got: {:?}", other),
+            }
+        });
+    }
+
+    #[test]
+    fn test_request_upgrade_supported() {
+        poll_immediate(async {
+            let mut socket = MockSocket {
+                read_byte: Some(b'S'),
+            };
+            let options = PgConnectOptions::new();
+            let res = request_upgrade(&mut socket, &options).await.unwrap();
+            assert!(res);
+        });
+    }
+
+    #[test]
+    fn test_request_upgrade_unsupported() {
+        poll_immediate(async {
+            let mut socket = MockSocket {
+                read_byte: Some(b'N'),
+            };
+            let options = PgConnectOptions::new();
+            let res = request_upgrade(&mut socket, &options).await.unwrap();
+            assert!(!res);
+        });
+    }
+
+    #[test]
+    fn test_request_upgrade_unexpected_byte() {
+        poll_immediate(async {
+            let mut socket = MockSocket {
+                read_byte: Some(0x42),
+            };
+            let options = PgConnectOptions::new();
+            let err = request_upgrade(&mut socket, &options).await.unwrap_err();
+
+            match err {
+                Error::Protocol(msg) => {
+                    assert!(msg.contains("unexpected response from SSLRequest: 0x42"));
+                }
+                other => panic!("expected Error::Protocol, but got: {:?}", other),
+            }
+        });
+    }
+}