use tokio::io::{AsyncReadExt, AsyncSeekExt, AsyncWriteExt, SeekFrom};
use tokio::net::TcpStream;
use tokio::time::{Duration, timeout};
use tracing::warn;
use super::connection::{FtpConnection, FtpResponseClass};
use super::listing::parse_ftp_list_response;
const DEFAULT_BUFFER_SIZE: usize = 65536;
#[derive(Debug, Clone)]
pub struct FtpDownloadOptions {
pub buffer_size: usize,
pub resume_offset: Option<u64>,
pub max_retries: u32,
pub binary_mode: bool,
pub data_connect_timeout: Duration,
pub recursive_download: bool,
}
impl Default for FtpDownloadOptions {
fn default() -> Self {
Self {
buffer_size: DEFAULT_BUFFER_SIZE,
resume_offset: None,
max_retries: 3,
binary_mode: true,
data_connect_timeout: Duration::from_secs(30),
recursive_download: false,
}
}
}
fn is_transient_io_error(e: &std::io::Error) -> bool {
use std::io::ErrorKind;
matches!(
e.kind(),
ErrorKind::Interrupted
| ErrorKind::WouldBlock
| ErrorKind::ConnectionReset
| ErrorKind::ConnectionAborted
| ErrorKind::BrokenPipe
| ErrorKind::TimedOut
) || e.to_string().to_lowercase().contains("temporary")
}
#[derive(Debug, Clone)]
pub struct DownloadProgress {
pub downloaded_bytes: u64,
pub total_bytes: Option<u64>,
pub speed_bytes_per_sec: f64,
}
#[derive(Debug, Clone)]
pub struct DownloadResult {
pub file_path: String,
pub bytes_downloaded: u64,
pub total_size: Option<u64>,
pub success: bool,
pub average_speed_bps: f64,
pub duration_secs: f64,
}
impl DownloadResult {
pub fn is_complete(&self) -> bool {
self.success
&& match self.total_size {
Some(total) => self.bytes_downloaded >= total,
None => self.bytes_downloaded > 0,
}
}
pub fn human_readable_size(bytes: u64) -> String {
const UNITS: &[&str] = &["B", "KB", "MB", "GB", "TB"];
let mut size = bytes as f64;
let mut unit_idx = 0;
while size >= 1024.0 && unit_idx < UNITS.len() - 1 {
size /= 1024.0;
unit_idx += 1;
}
if unit_idx == 0 {
format!("{} {}", bytes, UNITS[unit_idx])
} else {
format!("{:.2} {}", size, UNITS[unit_idx])
}
}
}
pub struct FtpDownload<'a> {
conn: &'a mut FtpConnection,
options: FtpDownloadOptions,
}
impl<'a> FtpDownload<'a> {
pub fn new(conn: &'a mut FtpConnection, options: Option<FtpDownloadOptions>) -> Self {
Self {
conn,
options: options.unwrap_or_default(),
}
}
pub async fn download_file(
&mut self,
remote_path: &str,
local_path: &str,
progress_callback: Option<fn(DownloadProgress)>,
) -> Result<DownloadResult, String> {
if self.options.binary_mode {
self.conn.type_image().await?;
} else {
self.conn.type_ascii().await?;
}
let file_size = self.conn.size(remote_path).await.ok();
if let Some(offset) = self.options.resume_offset
&& offset > 0
{
self.conn.rest(offset).await?;
}
let (data_host, data_port) = self.establish_data_connection().await?;
self.conn.retr(remote_path).await?;
let result = self
.receive_data_to_file(
&data_host,
data_port,
local_path,
file_size,
progress_callback,
)
.await?;
Ok(result)
}
pub async fn download_to_memory(&mut self, remote_path: &str) -> Result<Vec<u8>, String> {
self.conn.type_image().await?;
let file_size = self.conn.size(remote_path).await.ok();
if let Some(offset) = self.options.resume_offset
&& offset > 0
{
self.conn.rest(offset).await?;
}
let (data_host, data_port) = self.establish_data_connection().await?;
self.conn.retr(remote_path).await?;
let data = self
.receive_data_to_memory(&data_host, data_port, file_size)
.await?;
Ok(data)
}
pub async fn download_directory(
&mut self,
remote_dir: &str,
local_base_dir: &str,
progress_callback: Option<fn(DownloadProgress)>,
) -> Result<Vec<DownloadResult>, String> {
if !self.options.recursive_download {
return Err("Recursive download not enabled in options".to_string());
}
self.conn.cwd(remote_dir).await?;
let list_resp = self.conn.list(None).await?;
if list_resp.code != 150 && list_resp.code != 125 && list_resp.code != 226 {
}
let (data_host, data_port) = self.establish_data_connection().await?;
let listing_data = self
.receive_data_to_memory(&data_host, data_port, None)
.await?;
let listing_str = String::from_utf8_lossy(&listing_data);
let entries = parse_ftp_list_response(&listing_str);
std::fs::create_dir_all(local_base_dir)
.map_err(|e| format!("Failed to create local directory: {}", e))?;
let mut results = Vec::new();
for entry in entries {
if entry.is_directory {
let sub_remote = format!("{}/{}", remote_dir.trim_end_matches('/'), entry.name);
let sub_local = format!("{}/{}", local_base_dir.trim_end_matches('/'), entry.name);
let sub_results =
Box::pin(self.download_directory(&sub_remote, &sub_local, progress_callback))
.await?;
results.extend(sub_results);
} else {
let remote_file = format!("{}/{}", remote_dir.trim_end_matches('/'), entry.name);
let local_file = format!("{}/{}", local_base_dir.trim_end_matches('/'), entry.name);
let result = self
.download_file(&remote_file, &local_file, progress_callback)
.await?;
results.push(result);
}
}
Ok(results)
}
async fn establish_data_connection(&mut self) -> Result<(String, u16), String> {
if self.conn.options.passive_mode {
match self.conn.epsv().await {
Ok(port) => {
Ok((self.conn.host.clone(), port))
}
Err(_) => {
self.conn.pasv().await
}
}
} else {
match self.conn.eprt_active().await {
Ok((host, port)) => Ok((host, port)),
Err(_) => {
let port = self.conn.port_active().await?;
Ok(("127.0.0.1".to_string(), port))
}
}
}
}
async fn receive_data_to_file(
&mut self,
data_host: &str,
data_port: u16,
local_path: &str,
file_size: Option<u64>,
progress_callback: Option<fn(DownloadProgress)>,
) -> Result<DownloadResult, String> {
let mut data_stream: TcpStream = timeout(
self.options.data_connect_timeout,
TcpStream::connect((data_host, data_port)),
)
.await
.map_err(|_| {
format!(
"FTP data connection timeout ({}s)",
self.options.data_connect_timeout.as_secs()
)
})?
.map_err(|e| format!("FTP data connection failed: {}", e))?;
let mut file = tokio::fs::File::create(local_path)
.await
.map_err(|e| format!("Failed to create local file: {}", e))?;
if let Some(offset) = self.options.resume_offset
&& offset > 0
{
file.seek(SeekFrom::Start(offset))
.await
.map_err(|e| format!("Failed to seek file: {}", e))?;
}
let mut buffer = vec![0u8; self.options.buffer_size];
let mut total_downloaded = self.options.resume_offset.unwrap_or(0);
let start_time = std::time::Instant::now();
let mut read_retry_count = 0u32;
const MAX_READ_RETRIES: u32 = 3;
loop {
let read_result = data_stream.read(&mut buffer).await;
match read_result {
Ok(bytes_read) => {
read_retry_count = 0;
if bytes_read == 0 {
break; }
file.write_all(&buffer[..bytes_read])
.await
.map_err(|e| format!("Failed to write to local file: {}", e))?;
total_downloaded += bytes_read as u64;
if let Some(cb) = progress_callback {
let elapsed = start_time.elapsed().as_secs_f64();
let speed = if elapsed > 0.0 {
total_downloaded as f64 / elapsed
} else {
0.0
};
cb(DownloadProgress {
downloaded_bytes: total_downloaded,
total_bytes: file_size,
speed_bytes_per_sec: speed,
});
}
}
Err(ref e) if is_transient_io_error(e) && read_retry_count < MAX_READ_RETRIES => {
read_retry_count += 1;
let wait_ms = 1000u64 * (1 << (read_retry_count - 1));
warn!(
"FTP read error (#{}), retrying in {}ms...",
read_retry_count, wait_ms
);
tokio::time::sleep(Duration::from_millis(wait_ms)).await;
continue;
}
Err(e) => {
return Err(format!(
"FTP data read failed after {} retries: {}",
read_retry_count, e
));
}
}
}
file.flush()
.await
.map_err(|e| format!("Failed to flush file: {}", e))?;
drop(data_stream);
let final_resp = self.conn.read_response().await?;
if final_resp.class() != FtpResponseClass::PositiveCompletion
&& final_resp.class() != FtpResponseClass::PositivePreliminary
{
return Err(format!(
"Download completed but server reported error: {} {}",
final_resp.code, final_resp.message
));
}
let elapsed = start_time.elapsed().as_secs_f64();
let avg_speed = if elapsed > 0.0 {
total_downloaded as f64 / elapsed
} else {
0.0
};
Ok(DownloadResult {
file_path: local_path.to_string(),
bytes_downloaded: total_downloaded,
total_size: file_size,
success: true,
average_speed_bps: avg_speed,
duration_secs: elapsed,
})
}
async fn receive_data_to_memory(
&mut self,
data_host: &str,
data_port: u16,
expected_size: Option<u64>,
) -> Result<Vec<u8>, String> {
let mut data_stream = timeout(
self.options.data_connect_timeout,
TcpStream::connect((data_host, data_port)),
)
.await
.map_err(|_| "FTP data connection timeout")?
.map_err(|e| format!("FTP data connection failed: {}", e))?;
let capacity = expected_size.unwrap_or(1024 * 1024) as usize;
let mut result = Vec::with_capacity(capacity);
let mut buffer = vec![0u8; self.options.buffer_size];
loop {
let bytes_read = data_stream
.read(&mut buffer)
.await
.map_err(|e| format!("FTP data read error: {}", e))?;
if bytes_read == 0 {
break;
}
result.extend_from_slice(&buffer[..bytes_read]);
}
drop(data_stream);
let _final_resp = self.conn.read_response().await.ok();
Ok(result)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_human_readable_size() {
assert_eq!(DownloadResult::human_readable_size(500), "500 B");
assert_eq!(DownloadResult::human_readable_size(1024), "1.00 KB");
assert_eq!(DownloadResult::human_readable_size(1536), "1.50 KB");
assert_eq!(DownloadResult::human_readable_size(1048576), "1.00 MB");
assert_eq!(DownloadResult::human_readable_size(1073741824), "1.00 GB");
}
#[test]
fn test_download_result_complete() {
let full = DownloadResult {
file_path: "test.bin".into(),
bytes_downloaded: 1000,
total_size: Some(1000),
success: true,
average_speed_bps: 1000.0,
duration_secs: 1.0,
};
assert!(full.is_complete());
let partial = DownloadResult {
file_path: "test.bin".into(),
bytes_downloaded: 500,
total_size: Some(1000),
success: true,
average_speed_bps: 500.0,
duration_secs: 1.0,
};
assert!(!partial.is_complete());
let unknown_total = DownloadResult {
file_path: "test.bin".into(),
bytes_downloaded: 100,
total_size: None,
success: true,
average_speed_bps: 100.0,
duration_secs: 1.0,
};
assert!(unknown_total.is_complete());
let failed = DownloadResult {
file_path: "test.bin".into(),
bytes_downloaded: 1000,
total_size: Some(1000),
success: false, average_speed_bps: 1000.0,
duration_secs: 1.0,
};
assert!(!failed.is_complete());
}
#[test]
fn test_ftp_download_options_default() {
let opts = FtpDownloadOptions::default();
assert_eq!(opts.buffer_size, DEFAULT_BUFFER_SIZE);
assert!(opts.resume_offset.is_none());
assert_eq!(opts.max_retries, 3);
assert!(opts.binary_mode);
assert_eq!(opts.data_connect_timeout, Duration::from_secs(30));
assert!(!opts.recursive_download);
}
#[test]
fn test_is_transient_io_error() {
use std::io::ErrorKind;
let interrupted = std::io::Error::new(ErrorKind::Interrupted, "interrupted");
assert!(is_transient_io_error(&interrupted));
let would_block = std::io::Error::new(ErrorKind::WouldBlock, "would block");
assert!(is_transient_io_error(&would_block));
let connection_reset = std::io::Error::new(ErrorKind::ConnectionReset, "connection reset");
assert!(is_transient_io_error(&connection_reset));
let timed_out = std::io::Error::new(ErrorKind::TimedOut, "timed out");
assert!(is_transient_io_error(&timed_out));
let not_found = std::io::Error::new(ErrorKind::NotFound, "not found");
assert!(!is_transient_io_error(¬_found));
let permission_denied =
std::io::Error::new(ErrorKind::PermissionDenied, "permission denied");
assert!(!is_transient_io_error(&permission_denied));
let temp_error = std::io::Error::other("temporary failure");
assert!(is_transient_io_error(&temp_error));
}
}