Improve ZipFile

5 files changed, 149 insertions, 79 deletions

diff --git a/src/zip/driver.rs b/src/zip/driver.rs
index e5aa58d..8e8c27c 100644
--- a/src/zip/driver.rs
+++ b/src/zip/driver.rs
@@ -1,7 +1,8 @@
 use crate::driver::{ArchiveRead, ArchiveWrite, Driver};
-use crate::zip::error::{ZipError, ZipResult};
-use crate::zip::structs::{deserialize, EOCDR64Locator, CDR, EOCDR, EOCDR64};
-use crate::zip::ZipFile;
+use crate::zip::file::CompressionMethod;
+use crate::zip::structs::{deserialize, EOCDR64Locator, ExtraHeader, CDR, EOCDR, EOCDR64};
+use crate::zip::{ZipError, ZipFile, ZipResult};
+use chrono::{Local, NaiveDate, NaiveDateTime, NaiveTime};
 use std::collections::HashMap as Map;
 use std::fs::File;
 use std::io::{Read, Seek, SeekFrom, Write};
@@ -41,7 +42,7 @@ impl<IO: Read + Seek> ArchiveRead for Zip<IO> {
             io.read(&mut buf)?;
             buf
         };
-        let eocdr: EOCDR = deserialize(&buf).map_err(|_| ZipError::InvalidEOCDR)?;
+        let eocdr: EOCDR = deserialize(&buf).unwrap();
         let comment = {
             let mut buf: Vec<u8> = vec![0; eocdr.comment_len as usize];
             io.read(&mut buf)?;
@@ -55,31 +56,29 @@ impl<IO: Read + Seek> ArchiveRead for Zip<IO> {
             io.read(&mut buf)?;
             buf
         };
-        let (cd_pointer, cd_size, cd_records) = if u32::from_le_bytes(buf[0..4].try_into().unwrap())
-            == 0x07064b50
-        {
-            let eocdr64locator: EOCDR64Locator =
-                deserialize(&buf[4..]).map_err(|_| ZipError::InvalidEOCDR64Locator)?;
-
-            io.seek(SeekFrom::Start(eocdr64locator.eocdr64_pointer))?;
-            let buf = {
-                let mut buf = [0; 56];
-                io.read(&mut buf)?;
-                buf
+        let (cd_pointer, cd_size, cd_records) =
+            if u32::from_le_bytes(buf[0..4].try_into().unwrap()) == 0x07064b50 {
+                let eocdr64locator: EOCDR64Locator = deserialize(&buf[4..]).unwrap();
+
+                io.seek(SeekFrom::Start(eocdr64locator.eocdr64_pointer))?;
+                let buf = {
+                    let mut buf = [0; 56];
+                    io.read(&mut buf)?;
+                    buf
+                };
+                if u32::from_le_bytes(buf[0..4].try_into().unwrap()) != 0x06064b50 {
+                    return Err(ZipError::InvalidEOCDR64Signature.into());
+                }
+                let eocdr64: EOCDR64 = deserialize(&buf[4..]).unwrap();
+
+                (eocdr64.cd_pointer, eocdr64.cd_size, eocdr64.cd_records)
+            } else {
+                (
+                    eocdr.cd_pointer as u64,
+                    eocdr.cd_size as u64,
+                    eocdr.cd_records as u64,
+                )
             };
-            if u32::from_le_bytes(buf[0..4].try_into().unwrap()) != 0x06064b50 {
-                return Err(ZipError::InvalidEOCDR64Signature.into());
-            }
-            let eocdr64: EOCDR64 = deserialize(&buf[4..]).map_err(|_| ZipError::InvalidEOCDR64)?;
-
-            (eocdr64.cd_pointer, eocdr64.cd_size, eocdr64.cd_records)
-        } else {
-            (
-                eocdr.cd_pointer as u64,
-                eocdr.cd_size as u64,
-                eocdr.cd_records as u64,
-            )
-        };
 
         // Read cd records
         let mut files = Map::with_capacity(cd_records as usize);
@@ -96,7 +95,7 @@ impl<IO: Read + Seek> ArchiveRead for Zip<IO> {
                 return Err(ZipError::InvalidCDRSignature.into());
             }
             p += 4;
-            let cdr: CDR = deserialize(&buf[p..p + 42]).map_err(|_| ZipError::InvalidCDR)?;
+            let cdr: CDR = deserialize(&buf[p..p + 42]).unwrap();
             p += 42;
             let name = String::from_utf8(buf[p..p + cdr.name_len as usize].into()).unwrap();
             p += cdr.name_len as usize;
@@ -105,15 +104,61 @@ impl<IO: Read + Seek> ArchiveRead for Zip<IO> {
             let comment = String::from_utf8(buf[p..p + cdr.comment_len as usize].into()).unwrap();
             p += cdr.comment_len as usize;
 
+            let mut compressed_size = cdr.compressed_size as u64;
+            let mut size = cdr.size as u64;
+            let mut header_pointer = cdr.header_pointer as u64;
+
+            let mut ep: usize = 0;
+            while ep < cdr.extra_field_len as usize {
+                let header: ExtraHeader = deserialize(&extra_fields[ep..ep + 4]).unwrap();
+                ep += 4;
+                match header.id {
+                    0x0001 => {
+                        if size == 0xFFFFFFFF {
+                            compressed_size = deserialize(&extra_fields[ep..ep + 8]).unwrap();
+                            ep += 8;
+                        }
+                        if compressed_size == 0xFFFFFFFF {
+                            size = deserialize(&extra_fields[ep..ep + 8]).unwrap();
+                            ep += 8;
+                        }
+                        if header_pointer == 0xFFFFFFFF {
+                            header_pointer = deserialize(&extra_fields[ep..ep + 8]).unwrap();
+                            ep += 8;
+                        }
+                        if cdr.disk == 0xFFFF {
+                            ep += 4
+                        }
+                    }
+                    _ => ep += header.size as usize,
+                }
+            }
+
             files.insert(
                 name.clone(),
                 ZipFile::new(
+                    CompressionMethod::from_struct_id(cdr.compression_method)?,
+                    NaiveDateTime::new(
+                        NaiveDate::from_ymd_opt(
+                            (cdr.dos_date as i32 >> 9 & 0x7F) + 1980,
+                            cdr.dos_date as u32 >> 5 & 0xF,
+                            cdr.dos_date as u32 & 0x1F,
+                        )
+                        .ok_or(ZipError::InvalidDate)?,
+                        NaiveTime::from_hms_opt(
+                            (cdr.dos_time as u32 >> 11) & 0x1F,
+                            (cdr.dos_time as u32 >> 5) & 0x3F,
+                            (cdr.dos_time as u32 & 0x1F) * 2,
+                        )
+                        .ok_or(ZipError::InvalidTime)?,
+                    )
+                    .and_local_timezone(Local)
+                    .unwrap(),
+                    cdr.crc,
+                    compressed_size,
+                    size,
+                    header_pointer,
                     name,
-                    cdr.dos_date,
-                    cdr.dos_time,
-                    cdr.compression_method,
-                    cdr.compressed_size as u64,
-                    cdr.size as u64,
                     comment,
                 ),
             );
diff --git a/src/zip/error.rs b/src/zip/error.rs
index 18bbb22..1c5527c 100644
--- a/src/zip/error.rs
+++ b/src/zip/error.rs
@@ -7,15 +7,14 @@ pub type ZipResult<T> = ArchiveResult<T, ZipError>;
 #[derive(Debug)]
 pub enum ZipError {
     EOCDRNotFound,
-    InvalidEOCDR,
-    InvalidArchiveComment,
-
-    InvalidEOCDR64Locator,
     InvalidEOCDR64Signature,
-    InvalidEOCDR64,
-
     InvalidCDRSignature,
-    InvalidCDR,
+
+    InvalidArchiveComment,
+    InvalidCompressionMethod,
+    UnsupportedCompressionMethod,
+    InvalidDate,
+    InvalidTime,
     InvalidFileName,
     InvalidFileComment,
 }
@@ -32,25 +31,24 @@ impl From<ZipError> for ArchiveError<ZipError> {
 impl Display for ZipError {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
-            ZipError::EOCDRNotFound => write!(f, "End of central directory record not found"),
-            ZipError::InvalidEOCDR => write!(f, "Invalid end of central directory record"),
-            ZipError::InvalidArchiveComment => write!(f, "Invalid archive comment"),
-            ZipError::InvalidEOCDR64Locator => {
-                write!(f, "Invalid zip64 end of central directory locator")
-            }
-            ZipError::InvalidEOCDR64Signature => {
+            Self::EOCDRNotFound => write!(f, "End of central directory record not found"),
+            Self::InvalidEOCDR64Signature => {
                 write!(
                     f,
                     "Invalid signature of zip64 end of central directory record"
                 )
             }
-            ZipError::InvalidEOCDR64 => write!(f, "Invalid zip64 end of central directory record"),
-            ZipError::InvalidCDRSignature => {
+            Self::InvalidCDRSignature => {
                 write!(f, "Invalid signature of central directory record")
             }
-            ZipError::InvalidCDR => write!(f, "Invalid central directory record"),
-            ZipError::InvalidFileName => write!(f, "Invalid file name"),
-            ZipError::InvalidFileComment => write!(f, "Invalid file comment"),
+
+            Self::InvalidArchiveComment => write!(f, "Invalid archive comment"),
+            Self::InvalidCompressionMethod => writeln!(f, "Invalid compression method"),
+            Self::UnsupportedCompressionMethod => writeln!(f, "Unsupported compression method"),
+            Self::InvalidDate => write!(f, "Invalid date"),
+            Self::InvalidTime => write!(f, "Invalid time"),
+            Self::InvalidFileName => write!(f, "Invalid file name"),
+            Self::InvalidFileComment => write!(f, "Invalid file comment"),
         }
     }
 }
diff --git a/src/zip/file.rs b/src/zip/file.rs
index d5b3327..f735d65 100644
--- a/src/zip/file.rs
+++ b/src/zip/file.rs
@@ -1,45 +1,66 @@
 use crate::driver::ArchiveFile;
-use chrono::{NaiveDate, NaiveDateTime, NaiveTime};
+use crate::zip::{ZipError, ZipResult};
+use chrono::{DateTime, Local};
+
+pub enum CompressionMethod {
+    Store,
+    Deflate,
+    BZIP2,
+    LZMA,
+    ZStd,
+    XZ,
+}
+
+impl CompressionMethod {
+    pub(crate) fn from_struct_id(id: u16) -> ZipResult<Self> {
+        match id {
+            0 => Ok(Self::Store),
+            8 => Ok(Self::Deflate),
+            12 => Ok(Self::BZIP2),
+            14 => Ok(Self::LZMA),
+            93 => Ok(Self::ZStd),
+            95 => Ok(Self::XZ),
+            1..=7 | 9..=11 | 13 | 15..=20 | 94 | 96..=99 => {
+                Err(ZipError::UnsupportedCompressionMethod.into())
+            }
+            21..=92 | 100.. => Err(ZipError::InvalidCompressionMethod.into()),
+        }
+    }
+}
 
 pub struct ZipFile {
-    pub name: String,
-    pub datetime: NaiveDateTime,
-    pub compression_method: u16,
+    pub compression_method: CompressionMethod,
+    pub datetime: DateTime<Local>,
+    pub crc: u32,
     pub compressed_size: u64,
     pub size: u64,
+    pub header_pointer: u64,
+    pub name: String,
     pub comment: String,
 }
 
-impl ArchiveFile for ZipFile {}
-
 impl ZipFile {
-    pub(crate) fn new(
-        name: String,
-        dos_date: u16,
-        dos_time: u16,
-        compression_method: u16,
+    pub fn new(
+        compression_method: CompressionMethod,
+        datetime: DateTime<Local>,
+        crc: u32,
         compressed_size: u64,
         size: u64,
+        header_pointer: u64,
+        name: String,
         comment: String,
     ) -> Self {
-        let year = (dos_date >> 9 & 0x7F) + 1980;
-        let month = dos_date >> 5 & 0xF;
-        let day = dos_date & 0x1F;
-
-        let hour = (dos_time >> 11) & 0x1F;
-        let minute = (dos_time >> 5) & 0x3F;
-        let seconds = (dos_time & 0x1F) * 2;
-
         Self {
-            name,
-            datetime: NaiveDateTime::new(
-                NaiveDate::from_ymd_opt(year as i32, month as u32, day as u32).unwrap(),
-                NaiveTime::from_hms_opt(hour as u32, minute as u32, seconds as u32).unwrap(),
-            ),
             compression_method,
+            datetime,
+            crc,
             compressed_size,
             size,
+            header_pointer,
+            name,
             comment,
         }
     }
 }
+
+impl ArchiveFile for ZipFile {}
diff --git a/src/zip/mod.rs b/src/zip/mod.rs
index 5aeca97..8601526 100644
--- a/src/zip/mod.rs
+++ b/src/zip/mod.rs
@@ -4,5 +4,5 @@ mod file;
 mod structs;
 
 pub use driver::Zip;
-pub use error::ZipError;
+pub use error::{ZipError, ZipResult};
 pub use file::ZipFile;
diff --git a/src/zip/structs.rs b/src/zip/structs.rs
index 0f9579e..3a93c1b 100644
--- a/src/zip/structs.rs
+++ b/src/zip/structs.rs
@@ -40,7 +40,7 @@ pub struct CDR {
     pub compression_method: u16,
     pub dos_time: u16,
     pub dos_date: u16,
-    pub crc32: u32,
+    pub crc: u32,
     pub compressed_size: u32,
     pub size: u32,
     pub name_len: u16,
@@ -52,6 +52,12 @@ pub struct CDR {
     pub header_pointer: u32,
 }
 
+#[derive(Serialize, Deserialize)]
+pub struct ExtraHeader {
+    pub id: u16,
+    pub size: u16,
+}
+
 pub fn serialize<T: Serialize>(object: &mut T) -> StructResult<Vec<u8>> {
     Settings::default().serialize(object)
 }