Fix #487 (#501)

Author: levkk

pgdog/src/net/error.rs

@@ -7,9 +7,12 @@ use tokio_rustls::rustls;
 
 #[derive(Debug, Error)]
 pub enum Error {
-    #[error("{0}")]
+    #[error("io: {0}")]
     Io(#[from] std::io::Error),
 
+    #[error("connection closed by peer")]
+    UnexpectedEof,
+
     #[error("unsupported startup request: {0}")]
     UnsupportedStartup(i32),
 
@@ -58,9 +61,6 @@ pub enum Error {
     #[error("unknown transaction state identifier: {0}")]
     UnknownTransactionStateIdentifier(char),
 
-    #[error("eof")]
-    Eof,
-
     #[error("not text encoding")]
     NotTextEncoding,
 

pgdog/src/net/stream.rs

@@ -6,7 +6,7 @@ use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufStream, R
 use tokio::net::TcpStream;
 use tracing::{debug, enabled, trace, Level};
 
-use std::io::Error;
+use std::io::{Error, ErrorKind};
 use std::net::SocketAddr;
 use std::ops::Deref;
 use std::pin::Pin;
@@ -108,8 +108,8 @@ impl Stream {
     pub async fn check(&mut self) -> Result<(), crate::net::Error> {
         let mut buf = [0u8; 1];
         match self {
-            Self::Plain(plain) => plain.get_mut().peek(&mut buf).await?,
-            Self::Tls(tls) => tls.get_mut().get_mut().0.peek(&mut buf).await?,
+            Self::Plain(plain) => eof(plain.get_mut().peek(&mut buf).await)?,
+            Self::Tls(tls) => eof(tls.get_mut().get_mut().0.peek(&mut buf).await)?,
             Self::DevNull => 0,
         };
 
@@ -126,8 +126,8 @@ impl Stream {
         let bytes = message.to_bytes()?;
 
         match self {
-            Stream::Plain(ref mut stream) => stream.write_all(&bytes).await?,
-            Stream::Tls(ref mut stream) => stream.write_all(&bytes).await?,
+            Stream::Plain(ref mut stream) => eof(stream.write_all(&bytes).await)?,
+            Stream::Tls(ref mut stream) => eof(stream.write_all(&bytes).await)?,
             Self::DevNull => (),
         }
 
@@ -165,7 +165,7 @@ impl Stream {
         message: &impl Protocol,
     ) -> Result<usize, crate::net::Error> {
         let sent = self.send(message).await?;
-        self.flush().await?;
+        eof(self.flush().await)?;
         trace!("😳");
 
         Ok(sent)
@@ -180,7 +180,7 @@ impl Stream {
         for message in messages {
             sent += self.send(message).await?;
         }
-        self.flush().await?;
+        eof(self.flush().await)?;
         trace!("😳");
         Ok(sent)
     }
@@ -199,15 +199,15 @@ impl Stream {
 
     /// Read data into a buffer, avoiding unnecessary allocations.
     pub async fn read_buf(&mut self, bytes: &mut BytesMut) -> Result<Message, crate::net::Error> {
-        let code = self.read_u8().await?;
-        let len = self.read_i32().await?;
+        let code = eof(self.read_u8().await)?;
+        let len = eof(self.read_i32().await)?;
 
         bytes.put_u8(code);
         bytes.put_i32(len);
 
         // Length must be at least 4 bytes.
         if len < 4 {
-            return Err(crate::net::Error::Eof);
+            return Err(crate::net::Error::UnexpectedEof);
         }
 
         let capacity = len as usize + 1;
@@ -218,7 +218,7 @@ impl Stream {
             bytes.set_len(capacity);
         }
 
-        self.read_exact(&mut bytes[5..capacity]).await?;
+        eof(self.read_exact(&mut bytes[5..capacity]).await)?;
 
         let message = Message::new(bytes.split().freeze());
 
@@ -261,6 +261,19 @@ impl Stream {
     }
 }
 
+fn eof<T>(result: std::io::Result<T>) -> Result<T, crate::net::Error> {
+    match result {
+        Ok(val) => Ok(val),
+        Err(err) => {
+            if err.kind() == ErrorKind::UnexpectedEof {
+                Err(crate::net::Error::UnexpectedEof)
+            } else {
+                Err(crate::net::Error::Io(err))
+            }
+        }
+    }
+}
+
 /// Wrapper around SocketAddr
 /// to make it easier to debug.
 pub struct PeerAddr {