sftpgo/sftpd/transfer.go

package sftpd

import (
	"fmt"
	"io"
	"sync/atomic"

	"github.com/eikenb/pipeat"

	"github.com/drakkan/sftpgo/common"
	"github.com/drakkan/sftpgo/metrics"
	"github.com/drakkan/sftpgo/vfs"
)

type writerAtCloser interface {
	io.WriterAt
	io.Closer
}

type readerAtCloser interface {
	io.ReaderAt
	io.Closer
}

// transfer defines the transfer details.
// It implements the io.ReaderAt and io.WriterAt interfaces to handle SFTP downloads and uploads
type transfer struct {
	*common.BaseTransfer
	writerAt     writerAtCloser
	readerAt     readerAtCloser
	isFinished   bool
	maxWriteSize int64
}

func newTranfer(baseTransfer *common.BaseTransfer, pipeWriter *vfs.PipeWriter, pipeReader *pipeat.PipeReaderAt,
	maxWriteSize int64) *transfer {
	var writer writerAtCloser
	var reader readerAtCloser
	if baseTransfer.File != nil {
		writer = baseTransfer.File
		reader = baseTransfer.File
	} else if pipeWriter != nil {
		writer = pipeWriter
	} else if pipeReader != nil {
		reader = pipeReader
	}
	return &transfer{
		BaseTransfer: baseTransfer,
		writerAt:     writer,
		readerAt:     reader,
		isFinished:   false,
		maxWriteSize: maxWriteSize,
	}
}

// ReadAt reads len(p) bytes from the File to download starting at byte offset off and updates the bytes sent.
// It handles download bandwidth throttling too
func (t *transfer) ReadAt(p []byte, off int64) (n int, err error) {
	t.Connection.UpdateLastActivity()
	var readed int
	var e error

	readed, e = t.readerAt.ReadAt(p, off)
	atomic.AddInt64(&t.BytesSent, int64(readed))

	if e != nil && e != io.EOF {
		t.TransferError(e)
		return readed, e
	}
	t.HandleThrottle()
	return readed, e
}

// WriteAt writes len(p) bytes to the uploaded file starting at byte offset off and updates the bytes received.
// It handles upload bandwidth throttling too
func (t *transfer) WriteAt(p []byte, off int64) (n int, err error) {
	t.Connection.UpdateLastActivity()
	if off < t.MinWriteOffset {
		err := fmt.Errorf("Invalid write offset: %v minimum valid value: %v", off, t.MinWriteOffset)
		t.TransferError(err)
		return 0, err
	}
	var written int
	var e error

	written, e = t.writerAt.WriteAt(p, off)
	atomic.AddInt64(&t.BytesReceived, int64(written))

	if t.maxWriteSize > 0 && e == nil && atomic.LoadInt64(&t.BytesReceived) > t.maxWriteSize {
		e = common.ErrQuotaExceeded
	}
	if e != nil {
		t.TransferError(e)
		return written, e
	}
	t.HandleThrottle()
	return written, e
}

// Close it is called when the transfer is completed.
// It closes the underlying file, logs the transfer info, updates the user quota (for uploads)
// and executes any defined action.
// If there is an error no action will be executed and, in atomic mode, we try to delete
// the temporary file
func (t *transfer) Close() error {
	if err := t.setFinished(); err != nil {
		return err
	}
	err := t.closeIO()
	errBaseClose := t.BaseTransfer.Close()
	if errBaseClose != nil {
		err = errBaseClose
	}
	return t.Connection.GetFsError(err)
}

func (t *transfer) closeIO() error {
	var err error
	if t.File != nil {
		err = t.File.Close()
	} else if t.writerAt != nil {
		err = t.writerAt.Close()
		t.Lock()
		// we set ErrTransfer here so quota is not updated, in this case the uploads are atomic
		if err != nil && t.ErrTransfer == nil {
			t.ErrTransfer = err
		}
		t.Unlock()
	} else if t.readerAt != nil {
		err = t.readerAt.Close()
	}
	return err
}

func (t *transfer) setFinished() error {
	t.Lock()
	defer t.Unlock()
	if t.isFinished {
		return common.ErrTransferClosed
	}
	t.isFinished = true
	return nil
}

// used for ssh commands.
// It reads from src until EOF so it does not treat an EOF from Read as an error to be reported.
// EOF from Write is reported as error
func (t *transfer) copyFromReaderToWriter(dst io.Writer, src io.Reader) (int64, error) {
	defer t.Connection.RemoveTransfer(t)

	var written int64
	var err error

	if t.maxWriteSize < 0 {
		return 0, common.ErrQuotaExceeded
	}
	isDownload := t.GetType() == common.TransferDownload
	buf := make([]byte, 32768)
	for {
		t.Connection.UpdateLastActivity()
		nr, er := src.Read(buf)
		if nr > 0 {
			nw, ew := dst.Write(buf[0:nr])
			if nw > 0 {
				written += int64(nw)
				if isDownload {
					atomic.StoreInt64(&t.BytesSent, written)
				} else {
					atomic.StoreInt64(&t.BytesReceived, written)
				}
				if t.maxWriteSize > 0 && written > t.maxWriteSize {
					err = common.ErrQuotaExceeded
					break
				}
			}
			if ew != nil {
				err = ew
				break
			}
			if nr != nw {
				err = io.ErrShortWrite
				break
			}
		}
		if er != nil {
			if er != io.EOF {
				err = er
			}
			break
		}
		t.HandleThrottle()
	}
	t.ErrTransfer = err
	if written > 0 || err != nil {
		metrics.TransferCompleted(atomic.LoadInt64(&t.BytesSent), atomic.LoadInt64(&t.BytesReceived), t.GetType(), t.ErrTransfer)
	}
	return written, err
}