refactor(client): use streaming download with deferred file creation
Instead of downloading the entire file to memory before writing to disk, stream directly to a DeferredFileWriter that only creates the local file after receiving the first DATA packet. This provides the same guarantee (no local file created if remote doesn't exist) while being more memory efficient for large files. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
@@ -34,10 +34,61 @@
|
|||||||
//! tftp -m netascii put 192.168.1.1 readme.txt
|
//! tftp -m netascii put 192.168.1.1 readme.txt
|
||||||
//! ```
|
//! ```
|
||||||
|
|
||||||
use std::{env, fs::File, path::Path, process::ExitCode};
|
use std::{
|
||||||
|
env,
|
||||||
|
fs::File,
|
||||||
|
io::{self, Write},
|
||||||
|
path::{Path, PathBuf},
|
||||||
|
process::ExitCode,
|
||||||
|
};
|
||||||
|
|
||||||
use pfs_tftp::{Client, Mode};
|
use pfs_tftp::{Client, Mode};
|
||||||
|
|
||||||
|
/// A writer that defers file creation until the first write.
|
||||||
|
///
|
||||||
|
/// This ensures the local file is only created after we've confirmed the remote
|
||||||
|
/// file exists (i.e., after receiving the first DATA packet, not an ERROR).
|
||||||
|
struct DeferredFileWriter {
|
||||||
|
path: PathBuf,
|
||||||
|
file: Option<File>,
|
||||||
|
bytes_written: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl DeferredFileWriter {
|
||||||
|
fn new(path: PathBuf) -> Self {
|
||||||
|
Self {
|
||||||
|
path,
|
||||||
|
file: None,
|
||||||
|
bytes_written: 0,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn bytes_written(&self) -> u64 {
|
||||||
|
self.bytes_written
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Write for DeferredFileWriter {
|
||||||
|
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
|
||||||
|
if self.file.is_none() {
|
||||||
|
self.file = Some(File::create(&self.path)?);
|
||||||
|
}
|
||||||
|
// SAFETY: We just ensured self.file is Some above
|
||||||
|
let file = self.file.as_mut().expect("file is Some");
|
||||||
|
let n = file.write(buf)?;
|
||||||
|
self.bytes_written += n as u64;
|
||||||
|
Ok(n)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn flush(&mut self) -> io::Result<()> {
|
||||||
|
if let Some(ref mut file) = self.file {
|
||||||
|
file.flush()
|
||||||
|
} else {
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// Print usage information.
|
/// Print usage information.
|
||||||
fn print_usage(program: &str) {
|
fn print_usage(program: &str) {
|
||||||
eprintln!("TFTP Client (RFC 1350)");
|
eprintln!("TFTP Client (RFC 1350)");
|
||||||
@@ -248,28 +299,19 @@ fn main() -> ExitCode {
|
|||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
// Download to memory first - only create local file on success
|
// Use deferred writer - file is only created after first DATA packet
|
||||||
let data = match client.get(&remote_file, args.mode) {
|
let mut writer = DeferredFileWriter::new(PathBuf::from(&local_file));
|
||||||
Ok(data) => data,
|
|
||||||
Err(e) => {
|
|
||||||
eprintln!("Error: {e}");
|
|
||||||
return ExitCode::FAILURE;
|
|
||||||
}
|
|
||||||
};
|
|
||||||
|
|
||||||
// Write to local file
|
if let Err(e) = client.get_to_writer(&remote_file, args.mode, &mut writer) {
|
||||||
if let Err(e) = std::fs::write(&local_file, &data) {
|
eprintln!("Error: {e}");
|
||||||
eprintln!("Error writing file '{local_file}': {e}");
|
|
||||||
return ExitCode::FAILURE;
|
return ExitCode::FAILURE;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
let bytes = writer.bytes_written();
|
||||||
if args.verbose {
|
if args.verbose {
|
||||||
eprintln!("Received {} bytes", data.len());
|
eprintln!("Received {bytes} bytes");
|
||||||
}
|
}
|
||||||
println!(
|
println!("Downloaded '{remote_file}' -> '{local_file}' ({bytes} bytes)");
|
||||||
"Downloaded '{remote_file}' -> '{local_file}' ({} bytes)",
|
|
||||||
data.len()
|
|
||||||
);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
Command::Put {
|
Command::Put {
|
||||||
|
|||||||
Reference in New Issue
Block a user