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:
2025-12-21 13:48:51 +01:00
parent 27d891111f
commit ba358cbd79

View File

@@ -34,10 +34,61 @@
//! 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};
/// 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.
fn print_usage(program: &str) {
eprintln!("TFTP Client (RFC 1350)");
@@ -248,28 +299,19 @@ fn main() -> ExitCode {
}
};
// Download to memory first - only create local file on success
let data = match client.get(&remote_file, args.mode) {
Ok(data) => data,
Err(e) => {
// Use deferred writer - file is only created after first DATA packet
let mut writer = DeferredFileWriter::new(PathBuf::from(&local_file));
if let Err(e) = client.get_to_writer(&remote_file, args.mode, &mut writer) {
eprintln!("Error: {e}");
return ExitCode::FAILURE;
}
};
// Write to local file
if let Err(e) = std::fs::write(&local_file, &data) {
eprintln!("Error writing file '{local_file}': {e}");
return ExitCode::FAILURE;
}
let bytes = writer.bytes_written();
if args.verbose {
eprintln!("Received {} bytes", data.len());
eprintln!("Received {bytes} bytes");
}
println!(
"Downloaded '{remote_file}' -> '{local_file}' ({} bytes)",
data.len()
);
println!("Downloaded '{remote_file}' -> '{local_file}' ({bytes} bytes)");
}
Command::Put {