VoiletSshTools/internal/scp/upload.go

// Package scp implements the client half of the SCP wire protocol on top of an
// SSH session.
//
// It exists so the sshtools binaries stay dependency free: the protocol is
// small, and every normal OpenSSH server still ships the scp binary that the
// remote side of this protocol talks to.
package scp

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"io/fs"
	"os"
	"path/filepath"
	"strings"

	"golang.org/x/crypto/ssh"
)

// Options describes one upload batch.
type Options struct {
	// Dest is the remote directory that receives the paths. Empty means the
	// remote user's home directory.
	Dest string
	// Paths are local files and directories, uploaded in order. Each one is
	// handled according to what it is: a file is sent as one entry, a directory
	// is walked recursively. Callers do not have to distinguish them.
	Paths []string
	// Logf reports each entry as it is sent. nil disables reporting.
	Logf func(format string, args ...any)
}

// Stats summarises an upload.
type Stats struct {
	Files int
	Dirs  int
	Bytes int64
}

// Upload copies opts.Paths into opts.Dest on the remote host. The session must
// not have been started yet.
func Upload(sess *ssh.Session, opts Options) (Stats, error) {
	var stats Stats
	dest := opts.Dest
	if dest == "" {
		dest = "."
	}

	stdin, err := sess.StdinPipe()
	if err != nil {
		return stats, err
	}
	stdout, err := sess.StdoutPipe()
	if err != nil {
		return stats, err
	}
	sess.Stderr = os.Stderr

	// -d tells scp the target is a directory, -r that directories may arrive;
	// mkdir -p makes sure the target exists in the first place.
	cmd := "mkdir -p " + shellQuote(dest) + " && scp -t -r -d " + shellQuote(dest)
	if err := sess.Start(cmd); err != nil {
		return stats, err
	}

	writer := bufio.NewWriter(stdin)
	reader := bufio.NewReader(stdout)
	if err := readAck(reader); err != nil {
		stdin.Close()
		return stats, err
	}

	for _, path := range opts.Paths {
		if err := sendPath(writer, reader, path, opts, &stats); err != nil {
			stdin.Close()
			return stats, err
		}
	}
	if err := writer.Flush(); err != nil {
		stdin.Close()
		return stats, err
	}
	if err := stdin.Close(); err != nil {
		return stats, err
	}
	return stats, sess.Wait()
}

func sendPath(w *bufio.Writer, r *bufio.Reader, local string, opts Options, stats *Stats) error {
	// Stat, not Lstat: a symlink named on the command line is uploaded as the
	// file or directory it points at.
	info, err := os.Stat(local)
	if err != nil {
		return err
	}
	switch {
	case info.IsDir():
		return sendDir(w, r, local, opts, stats)
	case info.Mode().IsRegular():
		return sendFile(w, r, local, info.Size(), opts, stats)
	default:
		return fmt.Errorf("%s: not a regular file or directory", local)
	}
}

func sendFile(w *bufio.Writer, r *bufio.Reader, local string, size int64, opts Options, stats *Stats) error {
	name := filepath.Base(local)
	if err := checkName(name); err != nil {
		return err
	}
	mode := fs.FileMode(0644)
	if isExecutable(local) {
		mode = 0755
	}
	f, err := os.Open(local)
	if err != nil {
		return err
	}
	defer f.Close()

	if _, err := fmt.Fprintf(w, "C%04o %d %s\n", mode, size, name); err != nil {
		return err
	}
	if err := w.Flush(); err != nil {
		return err
	}
	if err := readAck(r); err != nil {
		return fmt.Errorf("%s: %w", local, err)
	}
	if _, err := io.CopyN(w, f, size); err != nil {
		return fmt.Errorf("%s: %w", local, err)
	}
	if err := w.WriteByte(0); err != nil {
		return err
	}
	if err := w.Flush(); err != nil {
		return err
	}
	if err := readAck(r); err != nil {
		return fmt.Errorf("%s: %w", local, err)
	}
	stats.Files++
	stats.Bytes += size
	if opts.Logf != nil {
		opts.Logf("%-40s %s", name, humanBytes(size))
	}
	return nil
}

func sendDir(w *bufio.Writer, r *bufio.Reader, local string, opts Options, stats *Stats) error {
	name := filepath.Base(local)
	if err := checkName(name); err != nil {
		return err
	}
	if _, err := fmt.Fprintf(w, "D%04o 0 %s\n", 0755, name); err != nil {
		return err
	}
	if err := w.Flush(); err != nil {
		return err
	}
	if err := readAck(r); err != nil {
		return fmt.Errorf("%s: %w", local, err)
	}
	stats.Dirs++
	if opts.Logf != nil {
		opts.Logf("%-40s dir", name+"/")
	}

	entries, err := os.ReadDir(local)
	if err != nil {
		return err
	}
	for _, entry := range entries {
		child := filepath.Join(local, entry.Name())
		info, err := entry.Info()
		if err != nil {
			return err
		}
		switch {
		case info.IsDir():
			if err := sendDir(w, r, child, opts, stats); err != nil {
				return err
			}
		case info.Mode().IsRegular():
			if err := sendFile(w, r, child, info.Size(), opts, stats); err != nil {
				return err
			}
		default:
			fmt.Fprintf(os.Stderr, "goscp: skipping %s (not a regular file)\n", child)
		}
	}

	if _, err := w.WriteString("E\n"); err != nil {
		return err
	}
	if err := w.Flush(); err != nil {
		return err
	}
	if err := readAck(r); err != nil {
		return fmt.Errorf("%s: %w", local, err)
	}
	return nil
}

// readAck consumes the single-byte reply that follows every SCP message.
func readAck(r *bufio.Reader) error {
	b, err := r.ReadByte()
	if err != nil {
		return fmt.Errorf("scp: %w", err)
	}
	switch b {
	case 0:
		return nil
	case 1, 2:
		msg, _ := r.ReadString('\n')
		msg = strings.TrimSpace(msg)
		if msg == "" {
			msg = "unspecified error"
		}
		return errors.New("scp: " + msg)
	default:
		return fmt.Errorf("scp: unexpected reply %q", b)
	}
}

func checkName(name string) error {
	if name == "" || name == "." || name == ".." {
		return fmt.Errorf("invalid file name %q", name)
	}
	if strings.ContainsAny(name, "\n\r/") {
		return fmt.Errorf("unsupported file name %q", name)
	}
	return nil
}

func isExecutable(path string) bool {
	info, err := os.Stat(path)
	if err != nil {
		return false
	}
	return info.Mode().Perm()&0111 != 0
}

// shellQuote wraps s for the remote shell, escaping embedded single quotes.
func shellQuote(s string) string {
	return "'" + strings.ReplaceAll(s, "'", "'\\''") + "'"
}

func humanBytes(n int64) string {
	const unit = 1024
	if n < unit {
		return fmt.Sprintf("%d B", n)
	}
	div, exp := int64(unit), 0
	for m := n / unit; m >= unit; m /= unit {
		div *= unit
		exp++
	}
	return fmt.Sprintf("%.1f %ciB", float64(n)/float64(div), "KMGTPE"[exp])
}