Handle short reads in concurrent File.ReadAt

client.go

@@ -1158,35 +1158,45 @@ func (f *File) Read(b []byte) (int, error) {
 // readChunkAt attempts to read the whole entire length of the buffer from the file starting at the offset.
 // It will continue progressively reading into the buffer until it fills the whole buffer, or an error occurs.
 func (f *File) readChunkAt(ch chan result, b []byte, off int64) (n int, err error) {
-	for err == nil && n < len(b) {
-		id := f.c.nextID()
-		typ, data, err := f.c.sendPacket(context.Background(), ch, &sshFxpReadPacket{
-			ID:     id,
-			Handle: f.handle,
-			Offset: uint64(off) + uint64(n),
-			Len:    uint32(len(b) - n),
-		})
-		if err != nil {
-			return n, err
+	for n < len(b) {
+		read, readErr := f.readOnceAt(ch, b[n:], off+int64(n))
+		n += read
+		if readErr != nil {
+			return n, readErr
 		}
+	}
 
-		switch typ {
-		case sshFxpStatus:
-			return n, normaliseError(unmarshalStatus(id, data))
+	return n, nil
+}
 
-		case sshFxpData:
-			data, err = unmarshalDataReply(id, data)
-			if err != nil {
-				return n, err
-			}
-			n += copy(b[n:], data)
+// readOnceAt requests data once at the specified offset. Unlike readChunkAt,
+// it does not retry when the server returns fewer bytes than requested.
+func (f *File) readOnceAt(ch chan result, b []byte, off int64) (n int, err error) {
+	id := f.c.nextID()
+	typ, data, err := f.c.sendPacket(context.Background(), ch, &sshFxpReadPacket{
+		ID:     id,
+		Handle: f.handle,
+		Offset: uint64(off),
+		Len:    uint32(len(b)),
+	})
+	if err != nil {
+		return 0, err
+	}
 
-		default:
-			return n, unimplementedPacketErr(typ)
+	switch typ {
+	case sshFxpStatus:
+		return 0, normaliseError(unmarshalStatus(id, data))
+
+	case sshFxpData:
+		data, err = unmarshalDataReply(id, data)
+		if err != nil {
+			return 0, err
 		}
-	}
+		return copy(b, data), nil
 
-	return
+	default:
+		return 0, unimplementedPacketErr(typ)
+	}
 }
 
 func (f *File) readAtSequential(b []byte, off int64) (read int, err error) {
@@ -1219,6 +1229,20 @@ func (f *File) ReadAt(b []byte, off int64) (int, error) {
 	return f.readAt(b, off)
 }
 
+// readFirstChunkAt reads once to find the chunk size for subsequent concurrent reads.
+// It returns the remaining buffer and its starting offset.
+func (f *File) readFirstChunkAt(b []byte, off int64, tryChunkSize int) (bNext []byte, offNext int64, chunkSize int, err error) {
+	first := b
+	if len(first) > tryChunkSize {
+		first = first[:tryChunkSize]
+	}
+	n, err := f.readOnceAt(nil, first, off)
+	if n == 0 && err == nil {
+		return b, off, 0, io.ErrNoProgress
+	}
+	return b[n:], off + int64(n), n, err
+}
+
 // readAt must be called while holding either the Read or Write mutex in File.
 // This code is concurrent safe with itself, but not with Close.
 func (f *File) readAt(b []byte, off int64) (int, error) {
@@ -1236,13 +1260,19 @@ func (f *File) readAt(b []byte, off int64) (int, error) {
 		return f.readAtSequential(b, off)
 	}
 
-	// Split the read into multiple maxPacket-sized concurrent reads bounded by maxConcurrentRequests.
-	// This allows writes with a suitably large buffer to transfer data at a much faster rate
+	b, off, chunkSize, err := f.readFirstChunkAt(b, off, f.c.maxPacket)
+	n := chunkSize
+	if err != nil {
+		return n, err
+	}
+
+	// Split the remaining read into concurrent chunks bounded by maxConcurrentRequests.
+	// This allows reads with a suitably large buffer to transfer data at a much faster rate
 	// by overlapping round trip times.
 
 	cancel := make(chan struct{})
 
-	concurrency := len(b)/f.c.maxPacket + 1
+	concurrency := len(b)/chunkSize + 1
 	if concurrency > f.c.maxConcurrentRequests || concurrency < 1 {
 		concurrency = f.c.maxConcurrentRequests
 	}
@@ -1264,7 +1294,6 @@ func (f *File) readAt(b []byte, off int64) (int, error) {
 
 		b := b
 		offset := off
-		chunkSize := f.c.maxPacket
 
 		for len(b) > 0 {
 			rb := b
@@ -1324,11 +1353,10 @@ func (f *File) readAt(b []byte, off int64) (int, error) {
 						if err == nil {
 							n = copy(packet.b, data)
 
-							// For normal disk files, it is guaranteed that this will read
-							// the specified number of bytes, or up to end of file.
-							// This implies, if we have a short read, that means EOF.
 							if n < len(packet.b) {
-								err = io.EOF
+								read, readErr := f.readAtSequential(packet.b[n:], packet.off+int64(n))
+								n += read
+								err = readErr
 							}
 						}
 
@@ -1371,11 +1399,11 @@ func (f *File) readAt(b []byte, off int64) (int, error) {
 
 	if firstErr.err != nil {
 		// firstErr.err != nil if and only if firstErr.off > our starting offset.
-		return int(firstErr.off - off), firstErr.err
+		return n + int(firstErr.off-off), firstErr.err
 	}
 
 	// As per spec for io.ReaderAt, we return nil error if and only if we read everything.
-	return len(b), nil
+	return n + len(b), nil
 }
 
 // writeToSequential implements WriteTo, but works sequentially with no parallelism.

client_test.go

@@ -3,6 +3,7 @@ package sftp
 import (
 	"bytes"
 	"errors"
+	"fmt"
 	"io"
 	"os"
 	"strings"
@@ -292,6 +293,99 @@ func TestClientReadDataReply(t *testing.T) {
 	}
 }
 
+func TestClientReadChunkAtShortReplies(t *testing.T) {
+	for _, reuseChannel := range []bool{false, true} {
+		for _, eof := range []bool{false, true} {
+			t.Run(fmt.Sprintf("reuseChannel=%t/eof=%t", reuseChannel, eof), func(t *testing.T) {
+				var requests int
+				c := startScriptedClient(t, func(_ fxp, id uint32) (fxp, []byte, bool) {
+					requests++
+					if eof && requests > 1 {
+						return sshFxpStatus, statusBody(marshalUint32(nil, id), sshFxEOF), true
+					}
+					return sshFxpData, marshalString(marshalUint32(nil, id), "AB"), true
+				})
+				defer c.Close()
+
+				var ch chan result
+				if reuseChannel {
+					ch = make(chan result, 1)
+				}
+				f := &File{c: c, handle: "somehandle"}
+				b := make([]byte, 6)
+				n, err := f.readChunkAt(ch, b, 11)
+				want, wantErr := "ABABAB", error(nil)
+				if eof {
+					want, wantErr = "AB", io.EOF
+				}
+				if !errors.Is(err, wantErr) || n != len(want) || string(b[:n]) != want {
+					t.Fatalf("readChunkAt = (%d, %v, %q), want (%d, %v, %q)", n, err, b[:n], len(want), wantErr, want)
+				}
+			})
+		}
+	}
+}
+
+func TestClientReadAtFirstChunk(t *testing.T) {
+	for _, tc := range []struct {
+		name    string
+		first   string
+		eof     bool
+		want    string
+		wantErr error
+	}{
+		{name: "short replies", first: "AB", want: "ABABAB"},
+		{name: "EOF after first chunk", first: "AB", eof: true, want: "AB", wantErr: io.EOF},
+		{name: "initial EOF", eof: true, wantErr: io.EOF},
+		{name: "no progress", wantErr: io.ErrNoProgress},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			var requests int
+			c := startScriptedClient(t, func(_ fxp, id uint32) (fxp, []byte, bool) {
+				requests++
+				if tc.eof && (tc.first == "" || requests > 1) {
+					return sshFxpStatus, statusBody(marshalUint32(nil, id), sshFxEOF), true
+				}
+				return sshFxpData, marshalString(marshalUint32(nil, id), tc.first), true
+			}, MaxPacket(4))
+			defer c.Close()
+
+			f := &File{c: c, handle: "somehandle"}
+			b := make([]byte, 6)
+			n, err := f.ReadAt(b, 11)
+			if !errors.Is(err, tc.wantErr) || n != len(tc.want) || string(b[:n]) != tc.want {
+				t.Fatalf("ReadAt = (%d, %v, %q), want (%d, %v, %q)", n, err, b[:n], len(tc.want), tc.wantErr, tc.want)
+			}
+		})
+	}
+}
+
+func TestClientReadAtConcurrentShortRepliesFallbackSequentially(t *testing.T) {
+	var requests int
+	c := startScriptedClient(t, func(_ fxp, id uint32) (fxp, []byte, bool) {
+		requests++
+		data := "AA"
+		if requests == 1 {
+			data = "AAAA"
+		}
+		return sshFxpData, marshalString(marshalUint32(nil, id), data), true
+	}, MaxPacket(4), MaxConcurrentRequestsPerFile(1))
+	defer c.Close()
+
+	f := &File{c: c, path: "somefile", handle: "somehandle"}
+	b := make([]byte, 16)
+	n, err := f.ReadAt(b, 0)
+	if err != nil {
+		t.Fatalf("ReadAt error = %v, want nil", err)
+	}
+	if n != len(b) || string(b) != strings.Repeat("A", len(b)) {
+		t.Fatalf("ReadAt = (%d, %q), want (%d, %q)", n, b, len(b), strings.Repeat("A", len(b)))
+	}
+	if requests != 7 {
+		t.Errorf("read requests = %d, want 7", requests)
+	}
+}
+
 func TestClientReadAtConcurrentDataReply(t *testing.T) {
 	c := startScriptedClient(t, func(_ fxp, id uint32) (fxp, []byte, bool) {
 		// Oversize declared length, no data.