From e9377c8373b34dca1f9655b2444adf52bcd32cb6 Mon Sep 17 00:00:00 2001
From: Joe Tsai <joetsai@digital-static.net>
Date: Fri, 21 Jul 2023 01:17:22 -0700
Subject: [PATCH] Properly handle io.EOF error conditions when reading

Previously, the Server.Serve method would never return nil,
because the infinite for-loop handling request packets would
only break if reading a packet reported an error.
A common termination condition is when the underlying connection
is closed and recvPacket returns io.EOF.
In which case Serve should ignore io.EOF and
treat it as a normal shutdown.

However, this means that recvPacket must correctly handle io.EOF
such that it never reports io.EOF if a packet is partially read.
There are two calls to io.ReadFull in recvPacket.
The first call correctly forwards an io.EOF error
if no additional bytes of the next packet are read.
However, the second call incorrectly forwards io.EOF
when no bytes of the payload could be read.
This is incorrect since we already read the length and
should convert the io.EOF into an io.ErrUnexpectedEOF.
---
 conn.go   | 4 +++-
 packet.go | 5 +++++
 server.go | 6 +++++-
 3 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/conn.go b/conn.go
index 1ce9725f..3bb2ba15 100644
--- a/conn.go
+++ b/conn.go
@@ -18,7 +18,9 @@ type conn struct {
 }
 
 // the orderID is used in server mode if the allocator is enabled.
-// For the client mode just pass 0
+// For the client mode just pass 0.
+// It returns io.EOF if the connection is closed and
+// there are no more packets to read.
 func (c *conn) recvPacket(orderID uint32) (uint8, []byte, error) {
 	return recvPacket(c, c.alloc, orderID)
 }
diff --git a/packet.go b/packet.go
index b06d0b95..1232ff1e 100644
--- a/packet.go
+++ b/packet.go
@@ -290,6 +290,11 @@ func recvPacket(r io.Reader, alloc *allocator, orderID uint32) (uint8, []byte, e
 		b = make([]byte, length)
 	}
 	if _, err := io.ReadFull(r, b[:length]); err != nil {
+		// ReadFull only returns EOF if it has read no bytes.
+		// In this case, that means a partial packet, and thus unexpected.
+		if err == io.EOF {
+			err = io.ErrUnexpectedEOF
+		}
 		debug("recv packet %d bytes: err %v", length, err)
 		return 0, nil, err
 	}
diff --git a/server.go b/server.go
index 503454e6..2e419f59 100644
--- a/server.go
+++ b/server.go
@@ -327,7 +327,7 @@ func handlePacket(s *Server, p orderedRequest) error {
 }
 
 // Serve serves SFTP connections until the streams stop or the SFTP subsystem
-// is stopped.
+// is stopped. It returns nil if the server exits cleanly.
 func (svr *Server) Serve() error {
 	defer func() {
 		if svr.pktMgr.alloc != nil {
@@ -353,6 +353,10 @@ func (svr *Server) Serve() error {
 	for {
 		pktType, pktBytes, err = svr.serverConn.recvPacket(svr.pktMgr.getNextOrderID())
 		if err != nil {
+			// Check whether the connection terminated cleanly in-between packets.
+			if err == io.EOF {
+				err = nil
+			}
 			// we don't care about releasing allocated pages here, the server will quit and the allocator freed
 			break
 		}