diff --git a/proc_macros/proc_macros.rs b/proc_macros/proc_macros.rs index 7e3931a..4ae53ca 100644 --- a/proc_macros/proc_macros.rs +++ b/proc_macros/proc_macros.rs @@ -43,3 +43,37 @@ pub fn impl_struct_reader_read(item: TokenStream) -> TokenStream { }; stream.into() } + +#[proc_macro] +pub fn define_struct_writer_fn(item: TokenStream) -> TokenStream { + let i = parse_macro_input!(item as Ident); + let lefname = format!("write_le_{}", i); + let lefname = Ident::new(&lefname, i.span()); + let befname = format!("write_be_{}", i); + let befname = Ident::new(&befname, i.span()); + let stream = quote! { + #[doc = concat!("Write [", stringify!(#i), "] in little endian.")] + fn #lefname(&mut self, data: #i) -> Result<(), Self::Error>; + #[doc = concat!("Write [", stringify!(#i), "] in big endian.")] + fn #befname(&mut self, data: #i) -> Result<(), Self::Error>; + }; + stream.into() +} + +#[proc_macro] +pub fn impl_struct_writer_write(item: TokenStream) -> TokenStream { + let i = parse_macro_input!(item as Ident); + let lefname = format!("write_le_{}", i); + let lefname = Ident::new(&lefname, i.span()); + let befname = format!("write_be_{}", i); + let befname = Ident::new(&befname, i.span()); + let stream = quote! { + fn #lefname(&mut self, data: #i) -> Result<(), Self::Error> { + self.write_all(&data.to_le_bytes()) + } + fn #befname(&mut self, data: #i) -> Result<(), Self::Error> { + self.write_all(&data.to_be_bytes()) + } + }; + stream.into() +} diff --git a/src/downloader/pd_file/file.rs b/src/downloader/pd_file/file.rs index e0df9da..deeb636 100644 --- a/src/downloader/pd_file/file.rs +++ b/src/downloader/pd_file/file.rs @@ -6,6 +6,7 @@ use crate::downloader::pd_file::part_status::PdFilePartStatus; use crate::downloader::pd_file::version::PdFileVersion; use crate::ext::atomic::AtomicQuick; use crate::ext::io::StructRead; +use crate::ext::io::StructWrite; use crate::ext::replace::ReplaceWith2; use crate::ext::rw_lock::GetRwLock; use crate::ext::try_err::TryErr; @@ -37,6 +38,8 @@ lazy_static! { static ref MAGIC_WORDS: Vec = vec![0x50, 0x44, 0xff, 0xff]; } +const FILE_SIZE_OFFSET: SeekFrom = SeekFrom::Start(12); + #[derive(Debug)] /// The pd file pub struct PdFile { @@ -325,6 +328,21 @@ impl PdFile { } } + /// Set the target size of the file. If unknown, set this to 0. + /// * `file_size` - The target size of the file. + pub fn set_file_size(&self, file_size: u64) -> Result<(), PdFileError> { + self.file_size.qstore(file_size); + if !self.is_mem_only() { + self.need_saved.qstore(true); + let mut f = self.file.get_mut(); + let f = f.as_mut().unwrap(); + f.seek(FILE_SIZE_OFFSET)?; + f.write_le_u64(file_size)?; + self.need_saved.qstore(false); + } + Ok(()) + } + /// Write all data to the file. pub fn write(&self) -> Result<(), PdFileError> { let mut f = self.file.get_mut(); @@ -334,19 +352,19 @@ impl PdFile { self.version.write_to(&mut f)?; let file_name = self.file_name.get_ref().try_err2(gettext("File name is not set."))?; let file_name = file_name.as_bytes(); - f.write_all(&(file_name.len() as u32).to_le_bytes())?; - f.write_all(&self.status.get_ref().int_value().to_le_bytes())?; + f.write_le_u32(file_name.len() as u32)?; + f.write_le_u8(self.status.get_ref().int_value())?; let ftype = self.ftype.get_ref(); - f.write_all(&ftype.int_value().to_le_bytes())?; + f.write_le_u8(ftype.int_value())?; let file_size = self.file_size.qload(); - f.write_all(&file_size.to_le_bytes())?; - f.write_all(&self.downloaded_file_size.qload().to_le_bytes())?; + f.write_le_u64(file_size)?; + f.write_le_u64(self.downloaded_file_size.qload())?; let part_size = if ftype.is_multi() { self.part_size.qload() } else { 0 }; - f.write_all(&part_size.to_le_bytes())?; + f.write_le_u32(part_size)?; f.write_all(file_name)?; if ftype.is_multi() && file_size != 0 && part_size != 0 { let part_counts = (file_size + (part_size as u64) - 1) / (part_size as u64); diff --git a/src/downloader/tasks.rs b/src/downloader/tasks.rs index b3348fd..a7115ed 100644 --- a/src/downloader/tasks.rs +++ b/src/downloader/tasks.rs @@ -64,5 +64,13 @@ pub async fn create_download_tasks_simple= 400 { return Err(DownloaderError::from(status)); } + if file_size == 0 && status != 206 { + match result.content_length() { + Some(len) => { + d.pd.set_file_size(len)?; + } + None => {} + } + } Ok(()) } diff --git a/src/ext/io.rs b/src/ext/io.rs index 02bb0a5..5a1141c 100644 --- a/src/ext/io.rs +++ b/src/ext/io.rs @@ -1,6 +1,9 @@ use proc_macros::define_struct_reader_fn; +use proc_macros::define_struct_writer_fn; use proc_macros::impl_struct_reader_read; +use proc_macros::impl_struct_writer_write; use std::io::Read; +use std::io::Write; /// Read number. pub trait StructRead { @@ -57,3 +60,37 @@ pub trait ClearFile { /// Clear all datas in file fn clear_file(&mut self) -> std::io::Result<()>; } + +/// Write number. +pub trait StructWrite { + /// The error type + type Error; + define_struct_writer_fn!(u8); + define_struct_writer_fn!(i8); + define_struct_writer_fn!(u16); + define_struct_writer_fn!(i16); + define_struct_writer_fn!(u32); + define_struct_writer_fn!(i32); + define_struct_writer_fn!(u64); + define_struct_writer_fn!(i64); + define_struct_writer_fn!(usize); + define_struct_writer_fn!(isize); + define_struct_writer_fn!(u128); + define_struct_writer_fn!(i128); +} + +impl StructWrite for T { + type Error = std::io::Error; + impl_struct_writer_write!(u8); + impl_struct_writer_write!(i8); + impl_struct_writer_write!(u16); + impl_struct_writer_write!(i16); + impl_struct_writer_write!(u32); + impl_struct_writer_write!(i32); + impl_struct_writer_write!(u64); + impl_struct_writer_write!(i64); + impl_struct_writer_write!(usize); + impl_struct_writer_write!(isize); + impl_struct_writer_write!(u128); + impl_struct_writer_write!(i128); +}