2 files changed, 107 insertions, 35 deletions

diff --git a/src/zip/file/info.rs b/src/zip/file/info.rs
index 599dcc3..38ea984 100644
--- a/src/zip/file/info.rs
+++ b/src/zip/file/info.rs
@@ -1,5 +1,5 @@
 use crate::driver::ArchiveFileInfo;
-use crate::zip::{ZipError, ZipResult};
+use crate::zip::datetime::DosDateTime;
 use chrono::{DateTime, Local};
 
 #[derive(Debug, PartialEq, Eq, Clone)]
@@ -14,42 +14,30 @@ pub enum CompressionMethod {
 }
 
 impl CompressionMethod {
-    pub(crate) fn from_struct_id(id: u16) -> ZipResult<Self> {
-        Ok(match id {
+    #[inline]
+    pub(crate) fn from_struct_id(id: u16) -> Self {
+        match id {
             0 => Self::Store,
             8 => Self::Deflate,
             12 => Self::BZip2,
             14 => Self::Lzma,
             93 => Self::Zstd,
             95 => Self::Xz,
-            1..=7 | 9..=11 | 13 | 15..=20 | 94 | 96..=99 => Self::Unsupported(id),
-            21..=92 | 100.. => return Err(ZipError::InvalidCompressionMethod(id).into()),
-        })
+            _ => Self::Unsupported(id),
+        }
     }
 }
 
 #[derive(Debug, PartialEq, Eq, Clone)]
 pub enum EncryptionMethod {
     None,
-    Weak(u8),
+    Weak,   // ZipCrypto
+    Aes128, // WinZip encryption
+    Aes192,
+    Aes256,
     Unsupported,
 }
 
-impl EncryptionMethod {
-    pub(crate) fn from_bif_flag(bit_flag: BitFlag, crc: u32, dos_time: u16) -> EncryptionMethod {
-        match (bit_flag.is_encrypted(), bit_flag.is_strong_encryption()) {
-            (false, false) => EncryptionMethod::None,
-            (true, false) => EncryptionMethod::Weak(if bit_flag.is_has_data_descriptor() {
-                (dos_time >> 8) as u8 // Info-ZIP modification
-            } else {
-                (crc >> 24) as u8
-            }),
-            (true, true) => EncryptionMethod::Unsupported,
-            _ => panic!("impossible"),
-        }
-    }
-}
-
 #[derive(Debug, PartialEq, Eq, Clone, Copy)]
 pub struct BitFlag {
     flag: u16,
@@ -159,7 +147,7 @@ pub struct ZipFileInfo {
 }
 
 impl ZipFileInfo {
-    pub fn new(
+    pub(crate) fn new(
         compression_method: CompressionMethod,
         encryption_method: EncryptionMethod,
         bit_flag: BitFlag,
@@ -189,6 +177,15 @@ impl ZipFileInfo {
         }
     }
 
+    #[inline]
+    pub(crate) fn password_check(&self) -> u8 {
+        if self.bit_flag.is_has_data_descriptor() {
+            (self.mtime.to_dos_time() >> 8) as u8
+        } else {
+            (self.crc >> 24) as u8
+        }
+    }
+
     pub fn is_dir(&self) -> bool {
         self.name.ends_with("/")
     }
diff --git a/src/zip/file/read.rs b/src/zip/file/read.rs
index c26b304..80bdcfb 100644
--- a/src/zip/file/read.rs
+++ b/src/zip/file/read.rs
@@ -1,11 +1,16 @@
 use crate::driver::FileDriver;
 use crate::utils::{IoCursor, ReadUtils};
-use crate::zip::encryption::WeakDecoder;
+use crate::zip::encryption::{AesDecoder, Keys, WeakDecoder};
+use crate::zip::structs::FILE_HEADER_SIGNATURE;
 use crate::zip::{CompressionMethod, EncryptionMethod, ZipError, ZipFileInfo, ZipResult};
+use aes::cipher::KeyInit;
+use aes::{Aes128, Aes192, Aes256};
 use bzip2::read::BzDecoder;
 use flate2::read::DeflateDecoder;
 use liblzma::read::XzDecoder;
 use liblzma::stream::{Filters, LzmaOptions, Stream};
+use pbkdf2::pbkdf2_hmac_array;
+use sha1::Sha1;
 use std::io::{
     BufReader, Error as IoError, ErrorKind as IoErrorKind, Read, Result as IoResult, Seek, SeekFrom,
 };
@@ -14,34 +19,101 @@ use zstd::stream::Decoder as ZstdDecoder;
 enum Encryption<Io: Read> {
     None(Io),
     Weak(WeakDecoder<Io>),
+    Aes128(AesDecoder<Io, Aes128>),
+    Aes192(AesDecoder<Io, Aes192>),
+    Aes256(AesDecoder<Io, Aes256>),
 }
 
 impl<Io: Read> Encryption<Io> {
-    pub fn new(io: Io, info: &ZipFileInfo, password: Option<&[u8]>) -> ZipResult<Self> {
+    #[inline]
+    pub fn new(mut io: Io, info: &ZipFileInfo, password: Option<&[u8]>) -> ZipResult<Self> {
         Ok(match info.encryption_method {
             EncryptionMethod::None => Self::None(io),
-            EncryptionMethod::Weak(check) => Self::Weak(WeakDecoder::new(
-                io,
-                check,
-                password.ok_or(ZipError::PasswordIsNotSpecified)?,
-            )?),
-            EncryptionMethod::Unsupported => {
-                return Err(ZipError::UnsupportedEncryptionMethod.into())
+            EncryptionMethod::Weak => {
+                let mut keys = Keys::new();
+                keys.set_password(password.ok_or(ZipError::PasswordIsNotSpecified)?);
+
+                let check = keys.set_header(io.read_arr()?);
+                if check != info.password_check() {
+                    return Err(ZipError::IncorrectPassword);
+                }
+
+                Self::Weak(WeakDecoder::new(io, keys))
+            }
+            EncryptionMethod::Aes128 => {
+                let header = io.read_arr::<10>()?;
+                let salt = &header[..8];
+                let check = &header[8..];
+
+                let hash = pbkdf2_hmac_array::<Sha1, 34>(
+                    password.ok_or(ZipError::PasswordIsNotSpecified)?,
+                    salt,
+                    1000,
+                );
+                let key = &hash[..16];
+
+                if check != &hash[32..] {
+                    return Err(ZipError::IncorrectPassword);
+                }
+
+                Self::Aes128(AesDecoder::new(io, Aes128::new(key.into()))?)
+            }
+            EncryptionMethod::Aes192 => {
+                let header = io.read_arr::<14>()?;
+                let salt = &header[..12];
+                let check = &header[12..];
+
+                let hash = pbkdf2_hmac_array::<Sha1, 50>(
+                    password.ok_or(ZipError::PasswordIsNotSpecified)?,
+                    salt,
+                    1000,
+                );
+                let key = &hash[..24];
+
+                if check != &hash[48..] {
+                    return Err(ZipError::IncorrectPassword);
+                }
+
+                Self::Aes192(AesDecoder::new(io, Aes192::new(key.into()))?)
+            }
+            EncryptionMethod::Aes256 => {
+                let header = io.read_arr::<18>()?;
+                let salt = &header[..16];
+                let check = &header[16..];
+
+                let hash = pbkdf2_hmac_array::<Sha1, 66>(
+                    password.ok_or(ZipError::PasswordIsNotSpecified)?,
+                    salt,
+                    1000,
+                );
+                let key = &hash[..32];
+
+                if check != &hash[64..] {
+                    return Err(ZipError::IncorrectPassword);
+                }
+
+                Self::Aes256(AesDecoder::new(io, Aes256::new(key.into()))?)
             }
+            EncryptionMethod::Unsupported => return Err(ZipError::UnsupportedEncryptionMethod),
         })
     }
 }
 
 impl<Io: Read> Read for Encryption<Io> {
+    #[inline]
     fn read(&mut self, buf: &mut [u8]) -> IoResult<usize> {
         match self {
             Self::None(io) => io.read(buf),
             Self::Weak(io) => io.read(buf),
+            Self::Aes128(io) => io.read(buf),
+            Self::Aes192(io) => io.read(buf),
+            Self::Aes256(io) => io.read(buf),
         }
     }
 }
 
 impl<Io: Read + Seek> Seek for Encryption<Io> {
+    #[inline]
     fn seek(&mut self, pos: SeekFrom) -> IoResult<u64> {
         match self {
             Self::None(io) => io.seek(pos),
@@ -62,6 +134,7 @@ enum Compression<Io: Read> {
 }
 
 impl<Io: Read + Seek> Compression<Io> {
+    #[inline]
     pub fn new(mut io: Io, info: &ZipFileInfo) -> ZipResult<Self> {
         Ok(match info.compression_method {
             CompressionMethod::Store => Self::Store(io),
@@ -88,13 +161,14 @@ impl<Io: Read + Seek> Compression<Io> {
             CompressionMethod::Zstd => Self::Zstd(ZstdDecoder::new(io)?),
             CompressionMethod::Xz => Self::Xz(XzDecoder::new(io)),
             CompressionMethod::Unsupported(id) => {
-                return Err(ZipError::UnsupportedCompressionMethod(id).into())
+                return Err(ZipError::UnsupportedCompressionMethod(id))
             }
         })
     }
 }
 
 impl<Io: Read> Read for Compression<Io> {
+    #[inline]
     fn read(&mut self, buf: &mut [u8]) -> IoResult<usize> {
         match self {
             Compression::Store(io) => io.read(buf),
@@ -107,6 +181,7 @@ impl<Io: Read> Read for Compression<Io> {
 }
 
 impl<Io: Read + Seek> Seek for Compression<Io> {
+    #[inline]
     fn seek(&mut self, pos: SeekFrom) -> IoResult<u64> {
         match self {
             Compression::Store(io) => io.seek(pos),
@@ -137,8 +212,8 @@ impl<'d, Io: Read + Seek> ZipFileReader<'d, Io> {
         io.seek(SeekFrom::Start(info.header_pointer))?;
 
         let buf = io.read_arr::<30>()?;
-        if u32::from_le_bytes(buf[..4].try_into().unwrap()) != 0x04034b50 {
-            return Err(ZipError::InvalidFileHeaderSignature.into());
+        if buf[..4] != FILE_HEADER_SIGNATURE {
+            return Err(ZipError::InvalidFileHeaderSignature);
         }
         let data_pointer = info.header_pointer
             + 30