From 15442d6625b33e23d5220d7c3bba573740606b78 Mon Sep 17 00:00:00 2001
From: Jeremy Howard <github@jhoward.fastmail.fm>
Date: Mon, 31 Aug 2026 09:39:56 +1000
Subject: [PATCH] Distinguish clean peer closure from protocol errors

---
 README.md     |  2 +-
 src/socket.rs | 21 +++++++++++++++++----
 src/wire.rs   |  8 +++++++-
 3 files changed, 25 insertions(+), 6 deletions(-)

diff --git a/README.md b/README.md
index aefab32..d5d6963 100644
--- a/README.md
+++ b/README.md
@@ -14,7 +14,7 @@ zmtpmini implements the corner of ZeroMQ that a Jupyter-style client needs. It c
 | TCP (with `TCP_NODELAY`) and generic `AsyncRead + AsyncWrite` streams | automatic reconnection |
 | a configurable frame-size cap (default 64 MiB) | unbounded internal queues |
 
-A dropped connection or a protocol violation returns an error. The crate never redials or buffers on its own.
+A dropped connection returns `Error::Closed`, while I/O failures and protocol violations retain their own error variants. The crate never redials or buffers on its own.
 
 ## Use
 
diff --git a/src/socket.rs b/src/socket.rs
index 5c77180..66d7e07 100644
--- a/src/socket.rs
+++ b/src/socket.rs
@@ -62,7 +62,10 @@ impl<S: AsyncRead + AsyncWrite + Unpin> ZmtpStream<S> {
     async fn read_frame(&mut self) -> Result<RawFrame> {
         loop {
             if let Some(f) = decode_frame(&mut self.rbuf, self.max_frame)? { return Ok(f); }
-            if self.s.read_buf(&mut self.rbuf).await? == 0 { return perr("connection closed by peer"); }
+            if self.s.read_buf(&mut self.rbuf).await? == 0 {
+                if self.rbuf.is_empty() && self.parts.is_empty() { return Err(Error::Closed); }
+                return perr("connection closed during a multipart message");
+            }
         }
     }
 
@@ -196,7 +199,10 @@ impl<R: AsyncRead + Unpin> PeerReader<R> {
     async fn read_frame(&mut self) -> Result<RawFrame> {
         loop {
             if let Some(frame) = decode_frame(&mut self.rbuf, self.max_frame)? { return Ok(frame); }
-            if self.s.read_buf(&mut self.rbuf).await? == 0 { return perr("connection closed by peer"); }
+            if self.s.read_buf(&mut self.rbuf).await? == 0 {
+                if self.rbuf.is_empty() && self.parts.is_empty() { return Err(Error::Closed); }
+                return perr("connection closed during a multipart message");
+            }
         }
     }
 
@@ -368,10 +374,15 @@ mod tests {
         let (f1, f2) = reader.await.unwrap();
         assert!(f1.body.len() == 1 << 20 && &f2.body[..] == b"after");
 
-        // peer disappearing surfaces as an error
+        // ordinary EOF is distinct from a protocol violation
         let (mut d, server) = scripted_peer("ROUTER").await;
         drop(server);
-        assert!(d.recv().await.is_err());
+        assert!(matches!(d.recv().await, Err(Error::Closed)));
+
+        let (mut d, mut server) = scripted_peer("ROUTER").await;
+        server.write_all(b"\0\x05x").await.unwrap();
+        drop(server);
+        assert!(matches!(d.recv().await, Err(Error::Protocol(_))));
     }
 
     #[tokio::test]
@@ -398,6 +409,8 @@ mod tests {
         assert_eq!(reader.recv().await.unwrap(), Incoming::Message(vec![Bytes::from_static(b"one"), Bytes::from_static(b"two")]));
         writer.send([b"reply".as_ref()]).await.unwrap();
         assert_eq!(dealer.recv().await.unwrap(), vec![Bytes::from_static(b"reply")]);
+        drop(dealer);
+        assert!(matches!(reader.recv().await, Err(Error::Closed)));
 
         let (client, server) = tokio::io::duplex(1 << 16);
         let (sub, peer) = tokio::join!(Sub::from_stream(client), Peer::xpublisher(server));
diff --git a/src/wire.rs b/src/wire.rs
index f3618b2..9a5d5eb 100644
--- a/src/wire.rs
+++ b/src/wire.rs
@@ -7,6 +7,8 @@ use bytes::{Buf, BufMut, Bytes, BytesMut};
 pub enum Error {
     /// Underlying I/O failure.
     Io(std::io::Error),
+    /// The peer closed an established connection.
+    Closed,
     /// The peer violated ZMTP or spoke an unsupported variant.
     Protocol(String),
 }
@@ -15,7 +17,11 @@ impl From<std::io::Error> for Error { fn from(e: std::io::Error) -> Self { Error
 
 impl std::fmt::Display for Error {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        match self { Error::Io(e) => write!(f, "io error: {e}"), Error::Protocol(m) => write!(f, "protocol error: {m}") }
+        match self {
+            Error::Io(e) => write!(f, "io error: {e}"),
+            Error::Closed => write!(f, "connection closed by peer"),
+            Error::Protocol(m) => write!(f, "protocol error: {m}"),
+        }
     }
 }