SnapGo/internal/infrastructure/ssh/sftp_uploader.go

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)
	}
}