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