1package sftp
  2
  3import (
  4	"bytes"
  5	"fmt"
  6	"sync"
  7	"time"
  8
  9	"github.com/picosh/pico/pkg/send/utils"
 10)
 11
 12type buffer struct {
 13	buf []byte
 14	m   sync.Mutex
 15}
 16
 17func (b *buffer) WriteAt(p []byte, pos int64) (n int, err error) {
 18	pLen := len(p)
 19	expLen := pos + int64(pLen)
 20	b.m.Lock()
 21	defer b.m.Unlock()
 22	if int64(len(b.buf)) < expLen {
 23		oldLen := len(b.buf)
 24		if int64(cap(b.buf)) < expLen {
 25			newBuf := make([]byte, expLen)
 26			copy(newBuf, b.buf)
 27			b.buf = newBuf
 28		}
 29		// A truncated buffer keeps its old bytes past its length.
 30		b.buf = b.buf[:expLen]
 31		clear(b.buf[oldLen:])
 32	}
 33	copy(b.buf[pos:], p)
 34	return pLen, nil
 35}
 36
 37func (b *buffer) truncate(size int64) {
 38	b.m.Lock()
 39	defer b.m.Unlock()
 40	if size < int64(len(b.buf)) {
 41		b.buf = b.buf[:size]
 42		return
 43	}
 44	b.buf = append(b.buf, make([]byte, size-int64(len(b.buf)))...)
 45}
 46
 47func (b *buffer) bytes() []byte {
 48	b.m.Lock()
 49	defer b.m.Unlock()
 50	return bytes.Clone(b.buf)
 51}
 52
 53// pendingFile is a file open for writing. Editors over sshfs stat a file
 54// right after creating it, so it is listed from here until the last handle
 55// closes and its contents are written.
 56type pendingFile struct {
 57	path    string
 58	entry   *utils.FileEntry
 59	buf     *buffer
 60	modTime time.Time
 61	handles int
 62	// stored is set when storage had the file at open. Unless a handle
 63	// truncated or wrote to it, it's left as it is.
 64	stored bool
 65	dirty  bool
 66	// replaces holds stored files this one was renamed from while open.
 67	// They are deleted once it is stored at its new path.
 68	replaces []renamedFile
 69}
 70
 71// renamedFile is a stored file renamed away, with the checksum it had then.
 72// Another session may replace it before it's deleted.
 73type renamedFile struct {
 74	path string
 75	sum  []byte
 76}
 77
 78func (p *pendingFile) info(name string) *utils.VirtualFile {
 79	return &utils.VirtualFile{
 80		FName:    name,
 81		FSize:    int64(len(p.buf.bytes())),
 82		FModTime: p.modTime,
 83	}
 84}
 85
 86type pendingWriter struct {
 87	handler *handler
 88	file    *pendingFile
 89}
 90
 91func (w *pendingWriter) WriteAt(p []byte, off int64) (int, error) {
 92	w.handler.mu.Lock()
 93	w.file.dirty = true
 94	w.handler.mu.Unlock()
 95	return w.file.buf.WriteAt(p, off)
 96}
 97
 98func (w *pendingWriter) Close() error {
 99	h, f := w.handler, w.file
100	h.mu.Lock()
101	f.handles--
102	// A file removed while open is dropped.
103	commit := f.handles == 0 && h.pending[f.path] == f
104	if commit {
105		delete(h.pending, f.path)
106	}
107	name := f.path
108	h.mu.Unlock()
109	if !commit || (f.stored && !f.dirty) {
110		return nil
111	}
112
113	data := f.buf.bytes()
114	entry := *f.entry
115	entry.Filepath = name
116	entry.Reader = bytes.NewReader(data)
117	entry.Size = int64(len(data))
118	msg, err := h.writeHandler.Write(h.session, &entry)
119	h.mu.Lock()
120	replaced := h.release(f)
121	h.mu.Unlock()
122	if err != nil {
123		_, _ = fmt.Fprintf(h.session.Stderr(), "%s\r\n", err)
124		return err
125	}
126	if msg != "" {
127		_, _ = fmt.Fprintf(h.session.Stderr(), "%s\r\n", msg)
128	}
129	for _, old := range replaced {
130		if err := h.deleteIfUnchanged(old.path, old.sum); err != nil {
131			_, _ = fmt.Fprintf(h.session.Stderr(), "%s\r\n", err)
132			return err
133		}
134	}
135	return nil
136}
137
138// ignoredWriter accepts a file the handler ignores by policy, so a
139// transfer moves on to the next one.
140type ignoredWriter struct {
141	handler *handler
142	path    string
143}
144
145func (ignoredWriter) WriteAt(p []byte, _ int64) (int, error) { return len(p), nil }
146
147func (w ignoredWriter) Close() error {
148	w.handler.mu.Lock()
149	defer w.handler.mu.Unlock()
150	w.handler.ignored[w.path]--
151	if w.handler.ignored[w.path] <= 0 {
152		delete(w.handler.ignored, w.path)
153	}
154	return nil
155}