feat: initial goget release — modern IPv6-first download utility
This commit is contained in:
@@ -0,0 +1,817 @@
|
||||
//go:build linux || freebsd
|
||||
// +build linux freebsd
|
||||
|
||||
package ftp
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"net/textproto"
|
||||
"net/url"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"codeberg.org/petrbalvin/goget/internal/output"
|
||||
)
|
||||
|
||||
type Client struct {
|
||||
conn net.Conn
|
||||
reader *textproto.Reader
|
||||
writer *bufio.Writer
|
||||
host string
|
||||
port int
|
||||
user string
|
||||
password string
|
||||
url string
|
||||
tlsConfig *tls.Config
|
||||
useTLS bool
|
||||
passive bool
|
||||
timeout time.Duration
|
||||
}
|
||||
|
||||
func NewClient(rawURL string, timeout time.Duration) (*Client, error) {
|
||||
u, err := url.Parse(rawURL)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("invalid url: %w", err)
|
||||
}
|
||||
|
||||
host := u.Hostname()
|
||||
port := 21
|
||||
if u.Port() != "" {
|
||||
p, err := strconv.Atoi(u.Port())
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("invalid port: %w", err)
|
||||
}
|
||||
port = p
|
||||
}
|
||||
|
||||
user := "anonymous"
|
||||
password := "anonymous@"
|
||||
if u.User != nil {
|
||||
user = u.User.Username()
|
||||
if pass, ok := u.User.Password(); ok {
|
||||
password = pass
|
||||
}
|
||||
}
|
||||
|
||||
useTLS := u.Scheme == "ftps"
|
||||
if !useTLS && (user != "anonymous" || (u.User != nil && u.User.Username() != "anonymous")) {
|
||||
fmt.Fprintf(os.Stderr, "Warning: FTP credentials sent in plaintext. Use ftps:// for encrypted transfers.\n")
|
||||
}
|
||||
|
||||
c := &Client{
|
||||
host: host,
|
||||
port: port,
|
||||
user: user,
|
||||
password: password,
|
||||
url: rawURL,
|
||||
useTLS: useTLS,
|
||||
passive: true, // Default to passive mode
|
||||
timeout: timeout,
|
||||
tlsConfig: &tls.Config{
|
||||
MinVersion: tls.VersionTLS12,
|
||||
ServerName: host,
|
||||
InsecureSkipVerify: false,
|
||||
},
|
||||
}
|
||||
|
||||
return c, nil
|
||||
}
|
||||
|
||||
func (c *Client) Connect() error {
|
||||
addr := net.JoinHostPort(c.host, fmt.Sprintf("%d", c.port))
|
||||
conn, err := net.DialTimeout("tcp", addr, c.timeout)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to connect: %w", err)
|
||||
}
|
||||
|
||||
c.conn = conn
|
||||
c.reader = textproto.NewReader(bufio.NewReader(conn))
|
||||
c.writer = bufio.NewWriter(conn)
|
||||
|
||||
// Read greeting
|
||||
code, _, err := c.readResponse()
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to read greeting: %w", err)
|
||||
}
|
||||
if code != 220 {
|
||||
return fmt.Errorf("unexpected greeting code: %d", code)
|
||||
}
|
||||
|
||||
// Upgrade to TLS if FTPS (explicit)
|
||||
if c.useTLS {
|
||||
if err := c.sendCommand("AUTH TLS"); err != nil {
|
||||
return err
|
||||
}
|
||||
code, _, err := c.readResponse()
|
||||
if err != nil {
|
||||
return fmt.Errorf("auth tls failed: %w", err)
|
||||
}
|
||||
if code != 234 && code != 334 {
|
||||
// Try AUTH SSL as fallback
|
||||
if err := c.sendCommand("AUTH SSL"); err != nil {
|
||||
return err
|
||||
}
|
||||
code, _, err = c.readResponse()
|
||||
if err != nil {
|
||||
return fmt.Errorf("auth ssl failed: %w", err)
|
||||
}
|
||||
if code != 234 && code != 334 {
|
||||
return fmt.Errorf("server does not support tls/ssl")
|
||||
}
|
||||
}
|
||||
|
||||
// Wrap connection with TLS
|
||||
tlsConn := tls.Client(c.conn, c.tlsConfig)
|
||||
if err := tlsConn.Handshake(); err != nil {
|
||||
return fmt.Errorf("tls handshake failed: %w", err)
|
||||
}
|
||||
c.conn = tlsConn
|
||||
c.reader = textproto.NewReader(bufio.NewReader(tlsConn))
|
||||
c.writer = bufio.NewWriter(tlsConn)
|
||||
}
|
||||
|
||||
// Send USER
|
||||
if err := c.sendCommand("USER %s", c.user); err != nil {
|
||||
return err
|
||||
}
|
||||
code, _, err = c.readResponse()
|
||||
if err != nil {
|
||||
return fmt.Errorf("user command failed: %w", err)
|
||||
}
|
||||
|
||||
// 331 means password required, 230 means already logged in
|
||||
if code == 331 {
|
||||
// Send PASS
|
||||
if err := c.sendCommand("PASS %s", c.password); err != nil {
|
||||
return err
|
||||
}
|
||||
code, _, err = c.readResponse()
|
||||
if err != nil {
|
||||
return fmt.Errorf("pass command failed: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
if code != 230 && code != 200 {
|
||||
return fmt.Errorf("login failed with code: %d", code)
|
||||
}
|
||||
|
||||
// Set binary mode
|
||||
if err := c.sendCommand("TYPE I"); err != nil {
|
||||
return err
|
||||
}
|
||||
code, _, err = c.readResponse()
|
||||
if err != nil {
|
||||
return fmt.Errorf("type command failed: %w", err)
|
||||
}
|
||||
if code != 200 {
|
||||
return fmt.Errorf("failed to set binary mode")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *Client) sendCommand(format string, args ...interface{}) error {
|
||||
cmd := fmt.Sprintf(format, args...)
|
||||
if _, err := c.writer.WriteString(cmd); err != nil {
|
||||
return fmt.Errorf("failed to send command: %w", err)
|
||||
}
|
||||
if _, err := c.writer.WriteString("\r\n"); err != nil {
|
||||
return fmt.Errorf("failed to send command terminator: %w", err)
|
||||
}
|
||||
if err := c.writer.Flush(); err != nil {
|
||||
return fmt.Errorf("failed to flush command: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *Client) readResponse() (int, string, error) {
|
||||
code, msg, err := c.reader.ReadResponse(0)
|
||||
if err != nil {
|
||||
return 0, "", err
|
||||
}
|
||||
return code, msg, nil
|
||||
}
|
||||
|
||||
func (c *Client) getFileSize(path string) (int64, error) {
|
||||
if err := c.sendCommand("SIZE %s", path); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
code, msg, err := c.readResponse()
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("size command failed: %w", err)
|
||||
}
|
||||
if code != 213 {
|
||||
return 0, fmt.Errorf("size command failed with code: %d", code)
|
||||
}
|
||||
|
||||
size, err := strconv.ParseInt(strings.TrimSpace(msg), 10, 64)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("failed to parse size: %w", err)
|
||||
}
|
||||
return size, nil
|
||||
}
|
||||
|
||||
func (c *Client) establishDataConnection() (net.Conn, error) {
|
||||
if c.passive {
|
||||
return c.enterPassiveMode()
|
||||
}
|
||||
return c.enterActiveMode()
|
||||
}
|
||||
|
||||
func (c *Client) enterPassiveMode() (net.Conn, error) {
|
||||
if err := c.sendCommand("PASV"); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
code, msg, err := c.readResponse()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("pasv command failed: %w", err)
|
||||
}
|
||||
if code != 227 {
|
||||
return nil, fmt.Errorf("pasv command failed with code: %d", code)
|
||||
}
|
||||
|
||||
// Parse PASV response: 227 Entering Passive Mode (h1,h2,h3,h4,p1,p2)
|
||||
start := strings.Index(msg, "(")
|
||||
end := strings.Index(msg, ")")
|
||||
if start == -1 || end == -1 {
|
||||
return nil, fmt.Errorf("invalid pasv response")
|
||||
}
|
||||
|
||||
parts := strings.Split(msg[start+1:end], ",")
|
||||
if len(parts) != 6 {
|
||||
return nil, fmt.Errorf("invalid pasv response format")
|
||||
}
|
||||
|
||||
h1, _ := strconv.Atoi(parts[0])
|
||||
h2, _ := strconv.Atoi(parts[1])
|
||||
h3, _ := strconv.Atoi(parts[2])
|
||||
h4, _ := strconv.Atoi(parts[3])
|
||||
p1, _ := strconv.Atoi(parts[4])
|
||||
p2, _ := strconv.Atoi(parts[5])
|
||||
|
||||
dataHost := fmt.Sprintf("%d.%d.%d.%d", h1, h2, h3, h4)
|
||||
dataPort := p1*256 + p2
|
||||
|
||||
// Use original host for TLS connections to match certificate
|
||||
if c.useTLS {
|
||||
dataHost = c.host
|
||||
}
|
||||
|
||||
conn, err := net.DialTimeout("tcp", fmt.Sprintf("%s:%d", dataHost, dataPort), c.timeout)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to connect to data port: %w", err)
|
||||
}
|
||||
|
||||
if c.useTLS {
|
||||
tlsConn := tls.Client(conn, c.tlsConfig)
|
||||
if err := tlsConn.Handshake(); err != nil {
|
||||
conn.Close()
|
||||
return nil, fmt.Errorf("tls handshake on data connection failed: %w", err)
|
||||
}
|
||||
return tlsConn, nil
|
||||
}
|
||||
|
||||
return conn, nil
|
||||
}
|
||||
|
||||
func (c *Client) enterActiveMode() (net.Conn, error) {
|
||||
// Get local IP from the control connection to bind only to that interface
|
||||
ctrlAddr, ok := c.conn.LocalAddr().(*net.TCPAddr)
|
||||
if !ok || ctrlAddr.IP == nil {
|
||||
return nil, fmt.Errorf("failed to get local address for active mode")
|
||||
}
|
||||
ctrlIP := ctrlAddr.IP.To4()
|
||||
if ctrlIP == nil {
|
||||
return nil, fmt.Errorf("no ipv4 address available for active mode")
|
||||
}
|
||||
|
||||
// Bind listener only to the interface of the control connection, not 0.0.0.0
|
||||
listener, err := net.Listen("tcp", net.JoinHostPort(ctrlIP.String(), "0"))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create listener: %w", err)
|
||||
}
|
||||
defer listener.Close()
|
||||
|
||||
// Set a deadline so Accept respects the timeout
|
||||
if tcpListener, ok := listener.(*net.TCPListener); ok {
|
||||
tcpListener.SetDeadline(time.Now().Add(c.timeout))
|
||||
}
|
||||
|
||||
listenAddr := listener.Addr().(*net.TCPAddr)
|
||||
|
||||
p1 := listenAddr.Port / 256
|
||||
p2 := listenAddr.Port % 256
|
||||
|
||||
portCmd := fmt.Sprintf("PORT %d,%d,%d,%d,%d,%d",
|
||||
ctrlIP[0], ctrlIP[1], ctrlIP[2], ctrlIP[3], p1, p2)
|
||||
|
||||
if err := c.sendCommand("%s", portCmd); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
code, _, err := c.readResponse()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("port command failed: %w", err)
|
||||
}
|
||||
if code != 200 {
|
||||
return nil, fmt.Errorf("port command failed with code: %d", code)
|
||||
}
|
||||
|
||||
// Accept incoming connection
|
||||
connChan := make(chan net.Conn, 1)
|
||||
errChan := make(chan error, 1)
|
||||
|
||||
go func() {
|
||||
conn, err := listener.Accept()
|
||||
if err != nil {
|
||||
errChan <- err
|
||||
return
|
||||
}
|
||||
connChan <- conn
|
||||
}()
|
||||
|
||||
select {
|
||||
case conn := <-connChan:
|
||||
return conn, nil
|
||||
case err := <-errChan:
|
||||
return nil, fmt.Errorf("failed to accept data connection: %w", err)
|
||||
case <-time.After(c.timeout):
|
||||
return nil, fmt.Errorf("timeout waiting for data connection")
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Client) DownloadFile(ctx context.Context, remotePath, localPath string, resumeOffset int64, progressFunc func(current, total int64)) error {
|
||||
// Get file size
|
||||
totalSize, err := c.getFileSize(remotePath)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to get file size: %w", err)
|
||||
}
|
||||
|
||||
// Check if we need to resume
|
||||
var startOffset int64 = 0
|
||||
if resumeOffset > 0 && resumeOffset < totalSize {
|
||||
startOffset = resumeOffset
|
||||
if err := c.sendCommand("REST %d", startOffset); err != nil {
|
||||
return err
|
||||
}
|
||||
code, _, err := c.readResponse()
|
||||
if err != nil {
|
||||
return fmt.Errorf("rest command failed: %w", err)
|
||||
}
|
||||
if code != 350 {
|
||||
return fmt.Errorf("rest command failed with code: %d", code)
|
||||
}
|
||||
} else if resumeOffset >= totalSize {
|
||||
// File already complete
|
||||
return nil
|
||||
}
|
||||
|
||||
// Open local file for writing
|
||||
mode := os.O_CREATE | os.O_WRONLY
|
||||
if startOffset > 0 {
|
||||
mode |= os.O_APPEND
|
||||
} else {
|
||||
mode |= os.O_TRUNC
|
||||
}
|
||||
|
||||
localFile, err := os.OpenFile(localPath, mode, 0644)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to open local file: %w", err)
|
||||
}
|
||||
defer localFile.Close()
|
||||
|
||||
// Seek to start position if resuming
|
||||
if startOffset > 0 {
|
||||
if _, err := localFile.Seek(startOffset, io.SeekStart); err != nil {
|
||||
return fmt.Errorf("failed to seek: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Establish data connection FIRST (before RETR command)
|
||||
dataConn, err := c.establishDataConnection()
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to establish data connection: %w", err)
|
||||
}
|
||||
defer dataConn.Close()
|
||||
|
||||
// Send RETR command AFTER data connection is established
|
||||
if err := c.sendCommand("RETR %s", remotePath); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Read response (should be 150 or 125)
|
||||
code, _, err := c.readResponse()
|
||||
if err != nil {
|
||||
return fmt.Errorf("retr command failed: %w", err)
|
||||
}
|
||||
if code != 150 && code != 125 {
|
||||
// Try waiting a bit for the server to respond
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
code, _, err = c.readResponse()
|
||||
if err != nil {
|
||||
return fmt.Errorf("retr command failed with code: %d", code)
|
||||
}
|
||||
if code != 150 && code != 125 {
|
||||
return fmt.Errorf("retr command failed with code: %d", code)
|
||||
}
|
||||
}
|
||||
|
||||
// Copy data with progress. Polls ctx between reads so a SIGINT
|
||||
// arriving mid-transfer is detected on the next iteration and the
|
||||
// partial progress is saved to ResumeMetadata before returning.
|
||||
var copied int64 = startOffset
|
||||
buf := make([]byte, 32*1024)
|
||||
for {
|
||||
if ctx.Err() != nil {
|
||||
if err := saveFTPResume(c.url, localPath, copied, totalSize); err != nil {
|
||||
return fmt.Errorf("failed to save resume metadata: %w", err)
|
||||
}
|
||||
return ctx.Err()
|
||||
}
|
||||
n, err := dataConn.Read(buf)
|
||||
if n > 0 {
|
||||
written, werr := localFile.Write(buf[:n])
|
||||
copied += int64(written)
|
||||
if werr != nil {
|
||||
return fmt.Errorf("failed to write: %w", err)
|
||||
}
|
||||
if progressFunc != nil {
|
||||
progressFunc(copied, totalSize)
|
||||
}
|
||||
}
|
||||
if err == io.EOF {
|
||||
break
|
||||
}
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to read from data connection: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Read final response
|
||||
code, _, err = c.readResponse()
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to read final response: %w", err)
|
||||
}
|
||||
if code != 226 && code != 250 {
|
||||
return fmt.Errorf("transfer failed with code: %d", code)
|
||||
}
|
||||
|
||||
// Successful transfer — drop any stale resume sidecar.
|
||||
_ = output.DeleteResumeMetadata(localPath)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *Client) ListRemoteDir(path string) ([]string, error) {
|
||||
// Establish data connection
|
||||
dataConn, err := c.establishDataConnection()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to establish data connection: %w", err)
|
||||
}
|
||||
defer dataConn.Close()
|
||||
|
||||
// Send LIST command
|
||||
cmd := "LIST"
|
||||
if path != "" {
|
||||
cmd = fmt.Sprintf("LIST %s", path)
|
||||
}
|
||||
if err := c.sendCommand("%s", cmd); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Read response (should be 150 or 125)
|
||||
code, _, err := c.readResponse()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("list command failed: %w", err)
|
||||
}
|
||||
if code != 150 && code != 125 {
|
||||
return nil, fmt.Errorf("list command failed with code: %d", code)
|
||||
}
|
||||
|
||||
// Read directory listing
|
||||
var lines []string
|
||||
scanner := bufio.NewScanner(dataConn)
|
||||
for scanner.Scan() {
|
||||
lines = append(lines, scanner.Text())
|
||||
}
|
||||
if err := scanner.Err(); err != nil {
|
||||
return nil, fmt.Errorf("failed to read directory listing: %w", err)
|
||||
}
|
||||
|
||||
// Read final response
|
||||
code, _, err = c.readResponse()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read final response: %w", err)
|
||||
}
|
||||
if code != 226 && code != 250 {
|
||||
return nil, fmt.Errorf("list command failed with code: %d", code)
|
||||
}
|
||||
|
||||
return lines, nil
|
||||
}
|
||||
|
||||
func (c *Client) ChangeDir(path string) error {
|
||||
if err := c.sendCommand("CWD %s", path); err != nil {
|
||||
return err
|
||||
}
|
||||
code, _, err := c.readResponse()
|
||||
if err != nil {
|
||||
return fmt.Errorf("cwd command failed: %w", err)
|
||||
}
|
||||
if code != 250 {
|
||||
return fmt.Errorf("cwd command failed with code: %d", code)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *Client) GetCurrentDir() (string, error) {
|
||||
if err := c.sendCommand("PWD"); err != nil {
|
||||
return "", err
|
||||
}
|
||||
code, msg, err := c.readResponse()
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("pwd command failed: %w", err)
|
||||
}
|
||||
if code != 257 {
|
||||
return "", fmt.Errorf("pwd command failed with code: %d", code)
|
||||
}
|
||||
|
||||
// Extract path from response (usually quoted)
|
||||
start := strings.Index(msg, "\"")
|
||||
end := strings.LastIndex(msg, "\"")
|
||||
if start != -1 && end > start {
|
||||
return msg[start+1 : end], nil
|
||||
}
|
||||
return strings.TrimSpace(msg), nil
|
||||
}
|
||||
|
||||
func (c *Client) Close() error {
|
||||
if err := c.sendCommand("QUIT"); err != nil {
|
||||
return err
|
||||
}
|
||||
_, _, _ = c.readResponse() // Ignore errors on close
|
||||
if c.conn != nil {
|
||||
return c.conn.Close()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Helper function to download entire directory recursively
|
||||
func (c *Client) DownloadDirectory(ctx context.Context, remotePath, localPath string, progressFunc func(file string, current, total int64)) error {
|
||||
// Create local directory
|
||||
if err := os.MkdirAll(localPath, 0755); err != nil {
|
||||
return fmt.Errorf("failed to create local directory: %w", err)
|
||||
}
|
||||
|
||||
// Save current directory
|
||||
origDir, err := c.GetCurrentDir()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer func() {
|
||||
c.ChangeDir(origDir)
|
||||
}()
|
||||
|
||||
// Change to remote directory
|
||||
if err := c.ChangeDir(remotePath); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// List directory contents
|
||||
entries, err := c.ListRemoteDir("")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Parse entries and download files
|
||||
for _, entry := range entries {
|
||||
fields := strings.Fields(entry)
|
||||
if len(fields) < 9 {
|
||||
continue
|
||||
}
|
||||
|
||||
name := fields[len(fields)-1]
|
||||
if name == "." || name == ".." {
|
||||
continue
|
||||
}
|
||||
|
||||
isDir := fields[0][0] == 'd'
|
||||
remoteFilePath := filepath.Join(remotePath, name)
|
||||
localFilePath := filepath.Join(localPath, name)
|
||||
|
||||
if isDir {
|
||||
if err := c.DownloadDirectory(ctx, remoteFilePath, localFilePath, progressFunc); err != nil {
|
||||
return err
|
||||
}
|
||||
} else {
|
||||
if progressFunc != nil {
|
||||
progressFunc(name, 0, 0)
|
||||
}
|
||||
if err := c.DownloadFile(ctx, remoteFilePath, localFilePath, 0, nil); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// saveFTPResume persists partial FTP download progress to the standard
|
||||
// `.goget.meta` sidecar so a subsequent `goget --resume` can continue
|
||||
// from `downloaded` bytes. Errors are wrapped to surface failures
|
||||
// without aborting the cancellation handling.
|
||||
func saveFTPResume(url, localPath string, downloaded, total int64) error {
|
||||
if localPath == "" || localPath == "-" || downloaded <= 0 {
|
||||
return nil
|
||||
}
|
||||
meta := output.NewResumeMetadata(url, "", "", downloaded, total)
|
||||
return meta.Save(localPath)
|
||||
}
|
||||
|
||||
// clientPool holds N pre-connected FTP clients. A single shared pool
|
||||
// removes the need to open/close connections for every file in a
|
||||
// recursive download, and is required for the parallel recursive
|
||||
// implementation because each *Client owns its own TCP control
|
||||
// connection and is not safe for concurrent use.
|
||||
type clientPool struct {
|
||||
available chan *Client
|
||||
all []*Client
|
||||
}
|
||||
|
||||
// newClientPool opens `size` independent FTP connections to the same
|
||||
// server. Connections are kept open until close() is called.
|
||||
func newClientPool(rawURL string, timeout time.Duration, size int) (*clientPool, error) {
|
||||
if size < 1 {
|
||||
size = 1
|
||||
}
|
||||
all := make([]*Client, 0, size)
|
||||
avail := make(chan *Client, size)
|
||||
for i := 0; i < size; i++ {
|
||||
c, err := NewClient(rawURL, timeout)
|
||||
if err != nil {
|
||||
for _, existing := range all {
|
||||
_ = existing.Close()
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
if err := c.Connect(); err != nil {
|
||||
for _, existing := range all {
|
||||
_ = existing.Close()
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
all = append(all, c)
|
||||
avail <- c
|
||||
}
|
||||
return &clientPool{available: avail, all: all}, nil
|
||||
}
|
||||
|
||||
// acquire blocks until a client is available, then returns it. The
|
||||
// caller must call release() when done.
|
||||
func (p *clientPool) acquire() *Client {
|
||||
return <-p.available
|
||||
}
|
||||
|
||||
// release returns a client to the pool. It must be called exactly once
|
||||
// per acquire().
|
||||
func (p *clientPool) release(c *Client) {
|
||||
p.available <- c
|
||||
}
|
||||
|
||||
// size returns the number of clients in the pool.
|
||||
func (p *clientPool) size() int {
|
||||
return len(p.all)
|
||||
}
|
||||
|
||||
// close shuts down every client in the pool. The pool must not be used
|
||||
// after this call.
|
||||
func (p *clientPool) close() {
|
||||
for _, c := range p.all {
|
||||
_ = c.Close()
|
||||
}
|
||||
}
|
||||
|
||||
// DownloadDirectoryConcurrent downloads a remote directory to localPath
|
||||
// using a pool of N FTP connections. All operations (LIST, RETR) use
|
||||
// absolute paths so concurrent workers never share or mutate CWD state
|
||||
// on the server side. The returned error is the first error encountered
|
||||
// by any worker; other workers continue in the background until they
|
||||
// hit ctx cancellation or finish their tasks.
|
||||
func DownloadDirectoryConcurrent(ctx context.Context, rawURL string, timeout time.Duration, remotePath, localPath string, parallel int, progressFunc func(file string, current, total int64)) error {
|
||||
if parallel < 1 {
|
||||
parallel = 1
|
||||
}
|
||||
pool, err := newClientPool(rawURL, timeout, parallel)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer pool.close()
|
||||
|
||||
sem := make(chan struct{}, parallel)
|
||||
var wg sync.WaitGroup
|
||||
var firstErr error
|
||||
var errMu sync.Mutex
|
||||
|
||||
setErr := func(err error) {
|
||||
errMu.Lock()
|
||||
if firstErr == nil {
|
||||
firstErr = err
|
||||
}
|
||||
errMu.Unlock()
|
||||
}
|
||||
|
||||
processEntry(ctx, pool, sem, &wg, setErr, remotePath, localPath, progressFunc)
|
||||
|
||||
wg.Wait()
|
||||
return firstErr
|
||||
}
|
||||
|
||||
// processEntry processes one remote path: lists it, then dispatches
|
||||
// its child entries (file → download, subdir → recursive processEntry)
|
||||
// through the shared semaphore-bounded goroutine pool. A single
|
||||
// sem is threaded through the recursion so the total in-flight
|
||||
// goroutines never exceed `parallel`, regardless of directory depth.
|
||||
func processEntry(ctx context.Context, pool *clientPool, sem chan struct{}, wg *sync.WaitGroup, setErr func(error), remotePath, localPath string, progressFunc func(file string, current, total int64)) {
|
||||
defer wg.Done()
|
||||
defer func() { <-sem }()
|
||||
|
||||
if ctx.Err() != nil {
|
||||
return
|
||||
}
|
||||
|
||||
if err := os.MkdirAll(localPath, 0755); err != nil {
|
||||
setErr(fmt.Errorf("failed to create local dir %s: %w", localPath, err))
|
||||
return
|
||||
}
|
||||
|
||||
c := pool.acquire()
|
||||
entries, err := c.ListRemoteDir(remotePath)
|
||||
pool.release(c)
|
||||
if err != nil {
|
||||
setErr(fmt.Errorf("list %s: %w", remotePath, err))
|
||||
return
|
||||
}
|
||||
|
||||
for _, entry := range entries {
|
||||
fields := strings.Fields(entry)
|
||||
if len(fields) < 9 {
|
||||
continue
|
||||
}
|
||||
name := fields[len(fields)-1]
|
||||
if name == "." || name == ".." {
|
||||
continue
|
||||
}
|
||||
isDir := fields[0][0] == 'd'
|
||||
remoteFilePath := filepath.Join(remotePath, name)
|
||||
localFilePath := filepath.Join(localPath, name)
|
||||
|
||||
if isDir {
|
||||
wg.Add(1)
|
||||
select {
|
||||
case sem <- struct{}{}:
|
||||
go processEntry(ctx, pool, sem, wg, setErr, remoteFilePath, localFilePath, progressFunc)
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
} else {
|
||||
wg.Add(1)
|
||||
select {
|
||||
case sem <- struct{}{}:
|
||||
go func(remoteFilePath, localFilePath, name string) {
|
||||
defer wg.Done()
|
||||
defer func() { <-sem }()
|
||||
if ctx.Err() != nil {
|
||||
return
|
||||
}
|
||||
c := pool.acquire()
|
||||
defer pool.release(c)
|
||||
if progressFunc != nil {
|
||||
progressFunc(name, 0, 0)
|
||||
}
|
||||
// DownloadFile's progress callback has a different
|
||||
// signature than our directory-level progressFunc, so
|
||||
// we wrap it here to keep the existing per-file
|
||||
// progress behaviour.
|
||||
var fileProgress func(current, total int64)
|
||||
if progressFunc != nil {
|
||||
fileProgress = func(current, total int64) {
|
||||
progressFunc(name, current, total)
|
||||
}
|
||||
}
|
||||
if err := c.DownloadFile(ctx, remoteFilePath, localFilePath, 0, fileProgress); err != nil {
|
||||
setErr(fmt.Errorf("download %s: %w", remoteFilePath, err))
|
||||
}
|
||||
}(remoteFilePath, localFilePath, name)
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user