feat: add FTP and SFTP upload support

2026-07-12 10:04:21 +08:00
parent f1998fb23c
commit dd12521be2
22 changed files with 1466 additions and 44 deletions
@@ -0,0 +1,213 @@
+package ssh
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"io"
+	"os"
+	"path"
+	"strings"
+	"time"
+
+	"github.com/pkg/sftp"
+
+	"github.com/mmmy/snapgo/internal/domain"
+)
+
+// SFTPUploader transfers screenshots through the SSH File Transfer Protocol.
+// It lives beside the SCP client so both protocols share the same SSH
+// authentication and host-key implementation without duplicating credential
+// discovery code.
+type SFTPUploader struct {
+	cfg domain.FTPConfig
+}
+
+// NewSFTPUploader returns a short-lived SFTP adapter. A fresh SSH/SFTP
+// connection is opened for every upload or connection probe.
+func NewSFTPUploader(cfg domain.FTPConfig) *SFTPUploader {
+	return &SFTPUploader{cfg: cfg}
+}
+
+func (u *SFTPUploader) sshConfig() domain.SSHConfig {
+	return domain.SSHConfig{
+		Host:               u.cfg.Host,
+		Port:               u.cfg.Port,
+		User:               u.cfg.User,
+		AuthMethod:         u.cfg.AuthMethod,
+		Password:           u.cfg.Password,
+		StrictHostKey:      u.cfg.StrictHostKey,
+		KnownHostsPath:     u.cfg.KnownHostsPath,
+		ConnectTimeoutSecs: u.cfg.ConnectTimeoutSecs,
+	}
+}
+
+func (u *SFTPUploader) dial(ctx context.Context) (*Client, *sftp.Client, error) {
+	sshClient, err := Dial(ctx, u.sshConfig())
+	if err != nil {
+		return nil, nil, fmt.Errorf("sftp: %w", err)
+	}
+	sftpClient, err := sftp.NewClient(sshClient.client)
+	if err != nil {
+		_ = sshClient.Close()
+		return nil, nil, fmt.Errorf("sftp: start subsystem: %w", err)
+	}
+	return sshClient, sftpClient, nil
+}
+
+// Upload writes data to a temporary file and atomically exposes it at the
+// final relative path after the transfer completes.
+func (u *SFTPUploader) Upload(ctx context.Context, remoteRelPath string, data []byte) error {
+	cleaned, err := validateSFTPRemotePath(remoteRelPath)
+	if err != nil {
+		return err
+	}
+	sshClient, sftpClient, err := u.dial(ctx)
+	if err != nil {
+		return err
+	}
+	done := closeSFTPOnCancel(ctx, sshClient, sftpClient)
+	defer sshClient.Close()
+	defer sftpClient.Close()
+	defer close(done)
+
+	dir := path.Dir(cleaned)
+	if dir != "." {
+		if err := sftpClient.MkdirAll(dir); err != nil {
+			return fmt.Errorf("sftp: create directory %s: %w", dir, err)
+		}
+	}
+
+	tempPath := path.Join(dir, temporarySFTPName(path.Base(cleaned)))
+	removeTemp := true
+	defer func() {
+		if removeTemp {
+			_ = sftpClient.Remove(tempPath)
+		}
+	}()
+
+	file, err := sftpClient.OpenFile(tempPath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC)
+	if err != nil {
+		return fmt.Errorf("sftp: create temporary file: %w", err)
+	}
+	reader := &sftpContextReader{ctx: ctx, reader: bytes.NewReader(data)}
+	_, copyErr := io.Copy(file, reader)
+	if copyErr == nil {
+		copyErr = file.Chmod(0o644)
+	}
+	closeErr := file.Close()
+	if copyErr != nil {
+		if ctx.Err() != nil {
+			return ctx.Err()
+		}
+		return fmt.Errorf("sftp: write %s: %w", cleaned, copyErr)
+	}
+	if closeErr != nil {
+		return fmt.Errorf("sftp: close %s: %w", cleaned, closeErr)
+	}
+	if err := sftpClient.Rename(tempPath, cleaned); err != nil {
+		return fmt.Errorf("sftp: commit %s: %w", cleaned, err)
+	}
+	removeTemp = false
+	sshLog().Info("SFTP upload ok", "remote_path", cleaned, "size", len(data))
+	return nil
+}
+
+// TestConnection verifies SFTP write/delete permissions with a small probe.
+func (u *SFTPUploader) TestConnection(ctx context.Context) error {
+	probe, err := sftpProbePath(u.cfg.PathPrefix)
+	if err != nil {
+		return err
+	}
+	sshClient, sftpClient, err := u.dial(ctx)
+	if err != nil {
+		return err
+	}
+	done := closeSFTPOnCancel(ctx, sshClient, sftpClient)
+	defer sshClient.Close()
+	defer sftpClient.Close()
+	defer close(done)
+
+	dir := path.Dir(probe)
+	if dir != "." {
+		if err := sftpClient.MkdirAll(dir); err != nil {
+			return fmt.Errorf("sftp probe: create directory: %w", err)
+		}
+	}
+	file, err := sftpClient.OpenFile(probe, os.O_WRONLY|os.O_CREATE|os.O_TRUNC)
+	if err != nil {
+		return fmt.Errorf("sftp probe: create: %w", err)
+	}
+	_, writeErr := file.Write([]byte("snapgo"))
+	closeErr := file.Close()
+	if writeErr != nil {
+		_ = sftpClient.Remove(probe)
+		return fmt.Errorf("sftp probe: write: %w", writeErr)
+	}
+	if closeErr != nil {
+		_ = sftpClient.Remove(probe)
+		return fmt.Errorf("sftp probe: close: %w", closeErr)
+	}
+	if err := sftpClient.Remove(probe); err != nil {
+		return fmt.Errorf("sftp probe: delete: %w", err)
+	}
+	return nil
+}
+
+func closeSFTPOnCancel(ctx context.Context, sshClient *Client, sftpClient *sftp.Client) chan struct{} {
+	done := make(chan struct{})
+	go func() {
+		select {
+		case <-ctx.Done():
+			_ = sftpClient.Close()
+			_ = sshClient.Close()
+		case <-done:
+		}
+	}()
+	return done
+}
+
+func validateSFTPRemotePath(value string) (string, error) {
+	value = strings.TrimSpace(value)
+	if value == "" || strings.HasPrefix(value, "/") || strings.HasPrefix(value, "~") {
+		return "", fmt.Errorf("sftp: remote path must be relative to the login directory")
+	}
+	if strings.ContainsAny(value, "\x00\r\n\\") {
+		return "", fmt.Errorf("sftp: remote path contains an invalid character")
+	}
+	for _, segment := range strings.Split(value, "/") {
+		if segment == ".." {
+			return "", fmt.Errorf("sftp: remote path must not contain '..'")
+		}
+	}
+	cleaned := path.Clean(value)
+	if cleaned == "." || path.Base(cleaned) == "." || path.Base(cleaned) == ".." {
+		return "", fmt.Errorf("sftp: remote filename is empty")
+	}
+	return cleaned, nil
+}
+
+func sftpProbePath(prefix string) (string, error) {
+	prefix = strings.TrimSpace(prefix)
+	prefix = strings.TrimPrefix(prefix, "~")
+	prefix = strings.TrimLeft(prefix, "/")
+	return validateSFTPRemotePath(path.Join(prefix, fmt.Sprintf(".snapgo-probe-%d", time.Now().UnixNano())))
+}
+
+func temporarySFTPName(base string) string {
+	return fmt.Sprintf(".%s.snapgo-upload-%d", base, time.Now().UnixNano())
+}
+
+type sftpContextReader struct {
+	ctx    context.Context
+	reader io.Reader
+}
+
+func (r *sftpContextReader) Read(p []byte) (int, error) {
+	select {
+	case <-r.ctx.Done():
+		return 0, r.ctx.Err()
+	default:
+		return r.reader.Read(p)
+	}
+}