From f3367a1b20aa78c8e4ad228acad439f4d386f5b7 Mon Sep 17 00:00:00 2001
From: Sveinn <info@sveinn.is>
Date: Tue, 14 Nov 2023 17:31:00 +0000
Subject: [PATCH] Adding error handling for network errors in the SFTP layer
 (#18442)

---
 cmd/sftp-server-driver.go | 30 ++++++++++++++++++++++++------
 1 file changed, 24 insertions(+), 6 deletions(-)

diff --git a/cmd/sftp-server-driver.go b/cmd/sftp-server-driver.go
index 336ffbab7..365d55a65 100644
--- a/cmd/sftp-server-driver.go
+++ b/cmd/sftp-server-driver.go
@@ -201,23 +201,40 @@ func (f *sftpDriver) Fileread(r *sftp.Request) (ra io.ReaderAt, err error) {
 	return obj, nil
 }
 
+// TransferError will catch network errors during transfer.
+// When TransferError() is called Close() will also
+// be called, so we do not need to Wait() here.
+func (w *writerAt) TransferError(err error) {
+	_ = w.w.CloseWithError(err)
+	_ = w.r.CloseWithError(err)
+	w.err = err
+}
+
 func (w *writerAt) Close() (err error) {
-	if len(w.buffer) > 0 {
-		err = w.w.CloseWithError(errors.New("some file segments were not flushed from the queue"))
-		for i := range w.buffer {
-			delete(w.buffer, i)
-		}
-	} else {
+	switch {
+	case len(w.buffer) > 0:
+		err = errors.New("some file segments were not flushed from the queue")
+		_ = w.w.CloseWithError(err)
+	case w.err != nil:
+		// No need to close here since both pipes were
+		// closing inside TransferError()
+		err = w.err
+	default:
 		err = w.w.Close()
 	}
+	for i := range w.buffer {
+		delete(w.buffer, i)
+	}
 	w.wg.Wait()
 	return err
 }
 
 type writerAt struct {
 	w      *io.PipeWriter
+	r      *io.PipeReader
 	wg     *sync.WaitGroup
 	buffer map[int64][]byte
+	err    error
 
 	nextOffset int64
 	m          sync.Mutex
@@ -279,6 +296,7 @@ func (f *sftpDriver) Filewrite(r *sftp.Request) (w io.WriterAt, err error) {
 	wa := &writerAt{
 		buffer: make(map[int64][]byte),
 		w:      pw,
+		r:      pr,
 		wg:     &sync.WaitGroup{},
 	}
 	wa.wg.Add(1)