1 files changed, 89 insertions, 90 deletions

diff --git a/src/zip/driver.rs b/src/zip/driver.rs
index e7878cf..8dc902f 100644
--- a/src/zip/driver.rs
+++ b/src/zip/driver.rs
@@ -1,48 +1,51 @@
 use crate::driver::{ArchiveRead, ArchiveWrite, Driver};
 use crate::utils::ReadUtils;
 use crate::zip::cp437::FromCp437;
-use crate::zip::structs::{deserialize, Cdr, Eocdr, Eocdr64, Eocdr64Locator, ExtraHeader};
+use crate::zip::datetime::DosDateTime;
+use crate::zip::structs::{
+    deserialize, AesField, Cdr, Eocdr, Eocdr64, Eocdr64Locator, ExtraHeader, CDR_SIGNATURE,
+    EOCDR64_LOCATOR_SIGNATURE, EOCDR64_SIGNATURE, EOCDR_SIGNATURE,
+};
 use crate::zip::{
     BitFlag, CompressionMethod, EncryptionMethod, ZipError, ZipFileInfo, ZipFileReader,
     ZipFileWriter, ZipResult,
 };
-use chrono::{DateTime, Local, NaiveDate, NaiveDateTime, NaiveTime};
+use chrono::{DateTime, Local};
 use std::collections::HashMap as Map;
 use std::fs::File;
 use std::io::{BufReader, Read, Seek, SeekFrom, Write};
 
-fn dos_to_local(date: u16, time: u16) -> ZipResult<DateTime<Local>> {
-    Ok(NaiveDateTime::new(
-        NaiveDate::from_ymd_opt(
-            (date as i32 >> 9 & 0x7F) + 1980,
-            date as u32 >> 5 & 0xF,
-            date as u32 & 0x1F,
-        )
-        .ok_or(ZipError::InvalidDate)?,
-        NaiveTime::from_hms_opt(
-            (time as u32 >> 11) & 0x1F,
-            (time as u32 >> 5) & 0x3F,
-            (time as u32 & 0x1F) * 2,
-        )
-        .ok_or(ZipError::InvalidTime)?,
-    )
-    .and_local_timezone(Local)
-    .unwrap())
+#[inline]
+fn split_fields(bytes: &[u8]) -> Option<Vec<(u16, &[u8])>> {
+    let mut fields = Vec::new();
+
+    let mut p = 0;
+    while p < bytes.len() {
+        let header: ExtraHeader = deserialize(bytes.get(p..p + 4)?).unwrap();
+        p += 4;
+        let data = bytes.get(p..p + header.size as usize)?;
+        p += header.size as usize;
+
+        fields.push((header.id, data));
+    }
+
+    Some(fields)
 }
 
-fn ntfs_to_local(time: u64) -> ZipResult<DateTime<Local>> {
-    Ok(DateTime::from_timestamp(
-        (time / 10000000 - 11644473600) as i64,
-        (time % 10000000) as u32,
+#[inline]
+fn ntfs_to_local(time: u64) -> Option<DateTime<Local>> {
+    Some(
+        DateTime::from_timestamp(
+            (time / 10000000 - 11644473600) as i64,
+            (time % 10000000) as u32,
+        )?
+        .with_timezone(&Local),
     )
-    .ok_or(ZipError::InvalidTime)?
-    .with_timezone(&Local))
 }
 
-fn timestamp_to_local(time: i32) -> ZipResult<DateTime<Local>> {
-    Ok(DateTime::from_timestamp(time as i64, 0)
-        .ok_or(ZipError::InvalidTime)?
-        .with_timezone(&Local))
+#[inline]
+fn timestamp_to_local(time: i32) -> Option<DateTime<Local>> {
+    Some(DateTime::from_timestamp(time as i64, 0)?.with_timezone(&Local))
 }
 
 pub struct Zip<Io = File> {
@@ -74,7 +77,7 @@ impl<Io: Read + Seek> ArchiveRead for Zip<Io> {
                     .ok_or(ZipError::EocdrNotFound)?,
             )?
             .windows(4)
-            .rposition(|v| u32::from_le_bytes(v.try_into().unwrap()) == 0x06054b50)
+            .rposition(|v| v == EOCDR_SIGNATURE)
             .ok_or(ZipError::EocdrNotFound)? as u64;
 
         // Read eocdr
@@ -88,13 +91,13 @@ impl<Io: Read + Seek> ArchiveRead for Zip<Io> {
         let buf = io.read_arr::<20>()?;
         let (cd_pointer, cd_size, cd_records) =
             // If locator found then read eocdr64
-            if u32::from_le_bytes(buf[0..4].try_into().unwrap()) == 0x07064b50 {
+            if buf[0..4] == EOCDR64_LOCATOR_SIGNATURE {
                 let eocdr64locator: Eocdr64Locator = deserialize(&buf[4..]).unwrap();
 
                 io.seek(SeekFrom::Start(eocdr64locator.eocdr64_pointer))?;
                 let buf = io.read_arr::<56>()?;
-                if u32::from_le_bytes(buf[0..4].try_into().unwrap()) != 0x06064b50 {
-                    return Err(ZipError::InvalidEOCDR64Signature.into());
+                if buf[0..4] != EOCDR64_SIGNATURE {
+                    return Err(ZipError::InvalidEOCDR64Signature);
                 }
                 let eocdr64: Eocdr64 = deserialize(&buf[4..]).unwrap();
 
@@ -116,8 +119,8 @@ impl<Io: Read + Seek> ArchiveRead for Zip<Io> {
         for i in 0..cd_records as usize {
             let buf = buf_reader.read_arr::<46>()?;
 
-            if u32::from_le_bytes(buf[..4].try_into().unwrap()) != 0x02014b50 {
-                return Err(ZipError::InvalidCDRSignature.into());
+            if buf[..4] != CDR_SIGNATURE {
+                return Err(ZipError::InvalidCDRSignature);
             }
             let cdr: Cdr = deserialize(&buf[4..46]).unwrap();
             let bit_flag = BitFlag::new(cdr.bit_flag);
@@ -138,91 +141,87 @@ impl<Io: Read + Seek> ArchiveRead for Zip<Io> {
                 String::from_cp437(comment)
             };
 
+            let mut compression_method = cdr.compression_method;
+            let mut encryption_method = if !bit_flag.is_encrypted() {
+                EncryptionMethod::None
+            } else if !bit_flag.is_strong_encryption() {
+                EncryptionMethod::Weak
+            } else {
+                EncryptionMethod::Unsupported
+            };
+
             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 mtime = dos_to_local(cdr.dos_date, cdr.dos_time)?;
+            let mut mtime = DateTime::from_dos_date_time(cdr.dos_date, cdr.dos_time, Local)?;
             let mut atime = None;
             let mut ctime = None;
 
             // Parse extensible data fields
-            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 {
+            for (id, mut data) in split_fields(&extra_fields).ok_or(ZipError::InvalidExtraFields)? {
+                match id {
                     // Zip64
                     0x0001 => {
                         if size == 0xFFFFFFFF {
-                            compressed_size =
-                                u64::from_le_bytes(extra_fields[ep..ep + 8].try_into().unwrap());
-                            ep += 8;
+                            size = u64::from_le_bytes(data.read_arr()?);
                         }
                         if compressed_size == 0xFFFFFFFF {
-                            size = u64::from_le_bytes(extra_fields[ep..ep + 8].try_into().unwrap());
-                            ep += 8;
+                            compressed_size = u64::from_le_bytes(data.read_arr()?);
                         }
                         if header_pointer == 0xFFFFFFFF {
-                            header_pointer =
-                                u64::from_le_bytes(extra_fields[ep..ep + 8].try_into().unwrap());
-                            ep += 8;
-                        }
-                        if cdr.disk == 0xFFFF {
-                            ep += 4
+                            header_pointer = u64::from_le_bytes(data.read_arr()?);
                         }
                     }
                     // NTFS
                     0x000a => {
-                        let mut tp = ep + 4;
-                        ep += header.size as usize;
-
-                        while tp < ep {
-                            let header: ExtraHeader =
-                                deserialize(&extra_fields[tp..tp + 4]).unwrap();
-                            tp += 4;
-
-                            match header.id {
+                        for (id, mut data) in
+                            split_fields(&data).ok_or(ZipError::InvalidExtraFields)?
+                        {
+                            match id {
                                 0x0001 => {
-                                    mtime = ntfs_to_local(u64::from_le_bytes(
-                                        extra_fields[tp..tp + 8].try_into().unwrap(),
-                                    ))?;
-                                    tp += 8;
-                                    atime = Some(ntfs_to_local(u64::from_le_bytes(
-                                        extra_fields[tp..tp + 8].try_into().unwrap(),
-                                    ))?);
-                                    tp += 8;
-                                    ctime = Some(ntfs_to_local(u64::from_le_bytes(
-                                        extra_fields[tp..tp + 8].try_into().unwrap(),
-                                    ))?);
-                                    tp += 8;
-                                }
-                                _ => {
-                                    tp += header.size as usize;
+                                    mtime = ntfs_to_local(u64::from_le_bytes(data.read_arr()?))
+                                        .unwrap_or(mtime);
+                                    atime = ntfs_to_local(u64::from_le_bytes(data.read_arr()?));
+                                    ctime = ntfs_to_local(u64::from_le_bytes(data.read_arr()?));
                                 }
+                                _ => {}
                             }
                         }
                     }
                     // Unix
                     0x000d => {
-                        atime = Some(timestamp_to_local(i32::from_le_bytes(
-                            extra_fields[ep..ep + 4].try_into().unwrap(),
-                        ))?);
-                        mtime = timestamp_to_local(i32::from_le_bytes(
-                            extra_fields[ep + 4..ep + 8].try_into().unwrap(),
-                        ))?;
-                        ep += header.size as usize
+                        atime = timestamp_to_local(i32::from_le_bytes(data.read_arr()?));
+                        mtime = timestamp_to_local(i32::from_le_bytes(data.read_arr()?))
+                            .unwrap_or(mtime);
+                    }
+                    // AES
+                    0x9901 => {
+                        let aes: AesField = deserialize(&data.read_arr::<7>()?).unwrap();
+                        if aes.id != 0x4541 {
+                            return Err(ZipError::InvalidExtraFields);
+                        }
+                        encryption_method = match aes.strength {
+                            0x01 => EncryptionMethod::Aes128,
+                            0x02 => EncryptionMethod::Aes192,
+                            0x03 => EncryptionMethod::Aes256,
+                            _ => EncryptionMethod::Unsupported,
+                        };
+                        compression_method = aes.compression_method;
                     }
                     // Skip unrecognized header
-                    _ => ep += header.size as usize,
+                    _ => {}
                 }
             }
 
+            if compression_method == 99 {
+                return Err(ZipError::AesExtraFieldNotFound);
+            }
+
             indexes.insert(name.clone(), i);
             files.push(ZipFileInfo::new(
-                CompressionMethod::from_struct_id(cdr.compression_method)?,
-                EncryptionMethod::from_bif_flag(bit_flag, cdr.crc, cdr.dos_time),
+                CompressionMethod::from_struct_id(compression_method),
+                encryption_method,
                 bit_flag,
                 mtime,
                 atime,
@@ -248,15 +247,15 @@ impl<Io: Read + Seek> ArchiveRead for Zip<Io> {
         &self.files
     }
 
-    fn get_file_index(&self, name: &str) -> crate::ArchiveResult<usize, Self::Error> {
+    fn get_file_index(&self, name: &str) -> ZipResult<usize> {
         self.indexes
             .get(name)
-            .ok_or(ZipError::FileNotFound.into())
+            .ok_or(ZipError::FileNotFound)
             .copied()
     }
 
     fn get_file_info(&self, index: usize) -> ZipResult<&Self::FileInfo> {
-        self.files.get(index).ok_or(ZipError::FileNotFound.into())
+        self.files.get(index).ok_or(ZipError::FileNotFound)
     }
 
     #[inline]