diff --git a/neqo-bin/src/bin/client/http09.rs b/neqo-bin/src/bin/client/http09.rs index a7dc2c21c7..6d9a26fec2 100644 --- a/neqo-bin/src/bin/client/http09.rs +++ b/neqo-bin/src/bin/client/http09.rs @@ -10,7 +10,7 @@ use std::{ cell::RefCell, collections::{HashMap, VecDeque}, fs::File, - io::Write, + io::{BufWriter, Write}, net::SocketAddr, path::PathBuf, rc::Rc, @@ -29,7 +29,7 @@ use super::{get_output_file, Args, KeyUpdateState, Res}; use crate::qlog_new; pub struct Handler<'a> { - streams: HashMap>, + streams: HashMap>>, url_queue: VecDeque, all_paths: Vec, args: &'a Args, @@ -219,7 +219,7 @@ impl<'b> Handler<'b> { client: &mut Connection, stream_id: StreamId, output_read_data: bool, - maybe_out_file: &mut Option, + maybe_out_file: &mut Option>, ) -> Res { let mut data = vec![0; 4096]; loop { @@ -246,8 +246,7 @@ impl<'b> Handler<'b> { } fn read(&mut self, client: &mut Connection, stream_id: StreamId) -> Res<()> { - let mut maybe_maybe_out_file = self.streams.get_mut(&stream_id); - match &mut maybe_maybe_out_file { + match self.streams.get_mut(&stream_id) { None => { println!("Data on unexpected stream: {stream_id}"); return Ok(()); @@ -261,7 +260,9 @@ impl<'b> Handler<'b> { )?; if fin_recvd { - if maybe_out_file.is_none() { + if let Some(mut out_file) = maybe_out_file.take() { + out_file.flush()?; + } else { println!(""); } self.streams.remove(&stream_id); diff --git a/neqo-bin/src/bin/client/http3.rs b/neqo-bin/src/bin/client/http3.rs index 754de9cb16..07cc0e4cde 100644 --- a/neqo-bin/src/bin/client/http3.rs +++ b/neqo-bin/src/bin/client/http3.rs @@ -11,7 +11,7 @@ use std::{ collections::{HashMap, VecDeque}, fmt::Display, fs::File, - io::Write, + io::{BufWriter, Write}, net::SocketAddr, path::PathBuf, rc::Rc, @@ -269,7 +269,7 @@ impl StreamHandlerType { } struct DownloadStreamHandler { - out_file: Option, + out_file: Option>, } impl StreamHandler for DownloadStreamHandler { @@ -300,8 +300,12 @@ impl StreamHandler for DownloadStreamHandler { println!("READ[{}]: 0x{}", stream_id, hex(&data)); } - if fin && self.out_file.is_none() { - println!(""); + if fin { + if let Some(mut out_file) = self.out_file.take() { + out_file.flush()?; + } else { + println!(""); + } } Ok(true) diff --git a/neqo-bin/src/bin/client/main.rs b/neqo-bin/src/bin/client/main.rs index d472dfb2bc..7b1a5928a6 100644 --- a/neqo-bin/src/bin/client/main.rs +++ b/neqo-bin/src/bin/client/main.rs @@ -8,7 +8,7 @@ use std::{ collections::{HashMap, VecDeque}, fmt::{self, Display}, fs::{create_dir_all, File, OpenOptions}, - io, + io::{self, BufWriter}, net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, ToSocketAddrs}, path::PathBuf, pin::Pin, @@ -36,6 +36,8 @@ use url::{Origin, Url}; mod http09; mod http3; +const BUFWRITER_BUFFER_SIZE: usize = 64 * 1024; + #[derive(Debug)] pub enum ClientError { ArgumentError(&'static str), @@ -252,7 +254,7 @@ fn get_output_file( url: &Url, output_dir: &Option, all_paths: &mut Vec, -) -> Option { +) -> Option> { if let Some(ref dir) = output_dir { let mut out_path = dir.clone(); @@ -284,7 +286,7 @@ fn get_output_file( .ok()?; all_paths.push(out_path); - Some(f) + Some(BufWriter::with_capacity(BUFWRITER_BUFFER_SIZE, f)) } else { None }