Add support for writing AES-encrypted files

Signed-off-by: Johannes Löthberg <johannes.loethberg@elokon.com>

examples/write_sample.rs

@@ -1,5 +1,7 @@
 use std::io::prelude::*;
 use zip::write::SimpleFileOptions;
+#[cfg(feature = "aes-crypto")]
+use zip::{AesMode, CompressionMethod};
 
 fn main() {
     std::process::exit(real_main());
@@ -38,6 +40,24 @@ fn doit(filename: &str) -> zip::result::ZipResult<()> {
     zip.start_file("test/lorem_ipsum.txt", options)?;
     zip.write_all(LOREM_IPSUM)?;
 
+    #[cfg(feature = "aes-crypto")]
+    {
+        zip.start_file(
+            "test/lorem_ipsum.aes.txt",
+            options
+                .compression_method(CompressionMethod::Zstd)
+                .with_aes_encryption(AesMode::Aes256, "password"),
+        )?;
+        zip.write_all(LOREM_IPSUM)?;
+
+        // This should use AE-1 due to the short file length.
+        zip.start_file(
+            "test/short.aes.txt",
+            options.with_aes_encryption(AesMode::Aes256, "password"),
+        )?;
+        zip.write_all(b"short text\n")?;
+    }
+
     zip.finish()?;
     Ok(())
 }

fuzz/fuzz_targets/fuzz_write.rs

@@ -7,38 +7,38 @@ use std::io::{Cursor, Read, Seek, Write};
 use std::path::PathBuf;
 
 #[derive(Arbitrary, Clone, Debug)]
-pub enum BasicFileOperation {
+pub enum BasicFileOperation<'k> {
     WriteNormalFile {
         contents: Vec<Vec<u8>>,
-        options: zip::write::FullFileOptions,
+        options: zip::write::FullFileOptions<'k>,
     },
-    WriteDirectory(zip::write::FullFileOptions),
+    WriteDirectory(zip::write::FullFileOptions<'k>),
     WriteSymlinkWithTarget {
         target: PathBuf,
-        options: zip::write::FullFileOptions,
+        options: zip::write::FullFileOptions<'k>,
     },
-    ShallowCopy(Box<FileOperation>),
+    ShallowCopy(Box<FileOperation<'k>>),
-    DeepCopy(Box<FileOperation>),
+    DeepCopy(Box<FileOperation<'k>>),
 }
 
 #[derive(Arbitrary, Clone, Debug)]
-pub struct FileOperation {
+pub struct FileOperation<'k> {
-    basic: BasicFileOperation,
+    basic: BasicFileOperation<'k>,
     path: PathBuf,
     reopen: bool,
     // 'abort' flag is separate, to prevent trying to copy an aborted file
 }
 
 #[derive(Arbitrary, Clone, Debug)]
-pub struct FuzzTestCase {
+pub struct FuzzTestCase<'k> {
     comment: Vec<u8>,
-    operations: Vec<(FileOperation, bool)>,
+    operations: Vec<(FileOperation<'k>, bool)>,
     flush_on_finish_file: bool,
 }
 
 fn do_operation<T>(
     writer: &mut zip::ZipWriter<T>,
-    operation: &FileOperation,
+    operation: &FileOperation<'k>,
     abort: bool,
     flush_on_finish_file: bool,
 ) -> Result<(), Box<dyn std::error::Error>>

src/aes.rs

@@ -318,7 +318,7 @@ mod tests {
         let mut read_buffer = vec![];
 
         {
-            let mut writer = AesWriter::new(&mut buf, aes_mode, &password)?;
+            let mut writer = AesWriter::new(&mut buf, aes_mode, password)?;
             writer.write_all(plaintext)?;
             writer.finish()?;
         }
@@ -329,7 +329,7 @@ mod tests {
         {
             let compressed_length = buf.get_ref().len() as u64;
             let mut reader =
-                match AesReader::new(&mut buf, aes_mode, compressed_length).validate(&password)? {
+                match AesReader::new(&mut buf, aes_mode, compressed_length).validate(password)? {
                     None => {
                         return Err(io::Error::new(
                             io::ErrorKind::InvalidData,
@@ -341,7 +341,7 @@ mod tests {
             reader.read_to_end(&mut read_buffer)?;
         }
 
-        return Ok(plaintext == read_buffer);
+        Ok(plaintext == read_buffer)
     }
 
     #[test]

src/lib.rs

@@ -29,7 +29,7 @@
 #![allow(unexpected_cfgs)] // Needed for cfg(fuzzing) on nightly as of 2024-05-06
 pub use crate::compression::{CompressionMethod, SUPPORTED_COMPRESSION_METHODS};
 pub use crate::read::ZipArchive;
-pub use crate::types::DateTime;
+pub use crate::types::{AesMode, DateTime};
 pub use crate::write::ZipWriter;
 
 #[cfg(feature = "aes-crypto")]

src/read.rs

@@ -927,11 +927,13 @@ fn central_header_to_zip_file_inner<R: Read>(
         central_extra_field: None,
         file_comment,
         header_start: offset,
+        extra_data_start: None,
         central_header_start,
         data_start: OnceLock::new(),
         external_attributes: external_file_attributes,
         large_file: false,
         aes_mode: None,
+        aes_extra_data_start: 0,
         extra_fields: Vec::new(),
     };
 
@@ -1316,6 +1318,7 @@ pub fn read_zipfile_from_stream<'a, R: Read>(reader: &'a mut R) -> ZipResult<Opt
         // header_start and data start are not available, but also don't matter, since seeking is
         // not available.
         header_start: 0,
+        extra_data_start: None,
         data_start: OnceLock::new(),
         central_header_start: 0,
         // The external_attributes field is only available in the central directory.
@@ -1324,6 +1327,7 @@ pub fn read_zipfile_from_stream<'a, R: Read>(reader: &'a mut R) -> ZipResult<Opt
         external_attributes: 0,
         large_file: false,
         aes_mode: None,
+        aes_extra_data_start: 0,
         extra_fields: Vec::new(),
     };
 

src/types.rs

@@ -352,6 +352,8 @@ pub struct ZipFileData {
     pub file_comment: Box<str>,
     /// Specifies where the local header of the file starts
     pub header_start: u64,
+    /// Specifies where the extra data of the file starts
+    pub extra_data_start: Option<u64>,
     /// Specifies where the central header of the file starts
     ///
     /// Note that when this is not known, it is set to 0
@@ -364,6 +366,8 @@ pub struct ZipFileData {
     pub large_file: bool,
     /// AES mode if applicable
     pub aes_mode: Option<(AesMode, AesVendorVersion, CompressionMethod)>,
+    /// Specifies where in the extra data the AES metadata starts
+    pub aes_extra_data_start: u64,
 
     /// extra fields, see <https://libzip.org/specifications/extrafld.txt>
     pub extra_fields: Vec<ExtraField>,
@@ -475,25 +479,33 @@ impl ZipFileData {
 /// According to the [specification](https://www.winzip.com/win/en/aes_info.html#winzip11) AE-2
 /// does not make use of the CRC check.
 #[derive(Copy, Clone, Debug)]
+#[repr(u16)]
 pub enum AesVendorVersion {
-    Ae1,
+    Ae1 = 0x0001,
-    Ae2,
+    Ae2 = 0x0002,
 }
 
 /// AES variant used.
 #[derive(Copy, Clone, Debug)]
+#[cfg_attr(fuzzing, derive(arbitrary::Arbitrary))]
+#[repr(u8)]
 pub enum AesMode {
-    Aes128,
+    /// 128-bit AES encryption.
-    Aes192,
+    Aes128 = 0x01,
-    Aes256,
+    /// 192-bit AES encryption.
+    Aes192 = 0x02,
+    /// 256-bit AES encryption.
+    Aes256 = 0x03,
 }
 
 #[cfg(feature = "aes-crypto")]
 impl AesMode {
+    /// Length of the salt for the given AES mode.
     pub const fn salt_length(&self) -> usize {
         self.key_length() / 2
     }
 
+    /// Length of the key for the given AES mode.
     pub const fn key_length(&self) -> usize {
         match self {
             Self::Aes128 => 16,
@@ -539,11 +551,13 @@ mod test {
             central_extra_field: None,
             file_comment: String::with_capacity(0).into_boxed_str(),
             header_start: 0,
+            extra_data_start: None,
             data_start: OnceLock::new(),
             central_header_start: 0,
             external_attributes: 0,
             large_file: false,
             aes_mode: None,
+            aes_extra_data_start: 0,
             extra_fields: Vec::new(),
         };
         assert_eq!(data.file_name_sanitized(), PathBuf::from("path/etc/passwd"));

src/unstable.rs

@@ -17,8 +17,8 @@ pub mod write {
         /// This is not recommended for new archives, as ZipCrypto is not secure.
         fn with_deprecated_encryption(self, password: &[u8]) -> Self;
     }
-    impl<T: FileOptionExtension> FileOptionsExt for FileOptions<T> {
+    impl<'k, T: FileOptionExtension> FileOptionsExt for FileOptions<'k, T> {
-        fn with_deprecated_encryption(self, password: &[u8]) -> Self {
+        fn with_deprecated_encryption(self, password: &[u8]) -> FileOptions<'static, T> {
             self.with_deprecated_encryption(password)
         }
     }

src/write.rs

@@ -1,10 +1,14 @@
 //! Types for creating ZIP archives
 
+#[cfg(feature = "aes-crypto")]
+use crate::aes::AesWriter;
 use crate::compression::CompressionMethod;
 use crate::read::{find_content, ZipArchive, ZipFile, ZipFileReader};
 use crate::result::{ZipError, ZipResult};
 use crate::spec;
-use crate::types::{ffi, DateTime, System, ZipFileData, DEFAULT_VERSION};
+#[cfg(feature = "aes-crypto")]
+use crate::types::AesMode;
+use crate::types::{ffi, AesVendorVersion, DateTime, System, ZipFileData, DEFAULT_VERSION};
 #[cfg(any(feature = "_deflate-any", feature = "bzip2", feature = "zstd",))]
 use core::num::NonZeroU64;
 use crc32fast::Hasher;
@@ -13,6 +17,7 @@ use std::default::Default;
 use std::io;
 use std::io::prelude::*;
 use std::io::{BufReader, SeekFrom};
+use std::marker::PhantomData;
 use std::mem;
 use std::str::{from_utf8, Utf8Error};
 use std::sync::{Arc, OnceLock};
@@ -42,19 +47,25 @@ use zstd::stream::write::Encoder as ZstdEncoder;
 
 enum MaybeEncrypted<W> {
     Unencrypted(W),
-    Encrypted(crate::zipcrypto::ZipCryptoWriter<W>),
+    #[cfg(feature = "aes-crypto")]
+    Aes(crate::aes::AesWriter<W>),
+    ZipCrypto(crate::zipcrypto::ZipCryptoWriter<W>),
 }
 impl<W: Write> Write for MaybeEncrypted<W> {
     fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
         match self {
             MaybeEncrypted::Unencrypted(w) => w.write(buf),
-            MaybeEncrypted::Encrypted(w) => w.write(buf),
+            #[cfg(feature = "aes-crypto")]
+            MaybeEncrypted::Aes(w) => w.write(buf),
+            MaybeEncrypted::ZipCrypto(w) => w.write(buf),
         }
     }
     fn flush(&mut self) -> io::Result<()> {
         match self {
             MaybeEncrypted::Unencrypted(w) => w.flush(),
-            MaybeEncrypted::Encrypted(w) => w.flush(),
+            #[cfg(feature = "aes-crypto")]
+            MaybeEncrypted::Aes(w) => w.flush(),
+            MaybeEncrypted::ZipCrypto(w) => w.flush(),
         }
     }
 }
@@ -177,24 +188,52 @@ mod sealed {
     }
 }
 
+#[derive(Copy, Clone, Debug)]
+enum EncryptWith<'k> {
+    #[cfg(feature = "aes-crypto")]
+    Aes {
+        mode: AesMode,
+        password: &'k str,
+    },
+    ZipCrypto(ZipCryptoKeys, PhantomData<&'k ()>),
+}
+
+#[cfg(fuzzing)]
+impl<'a> arbitrary::Arbitrary<'a> for EncryptWith<'a> {
+    fn arbitrary(u: &mut arbitrary::Unstructured<'a>) -> arbitrary::Result<Self> {
+        #[cfg(feature = "aes-crypto")]
+        if bool::arbitrary(u)? {
+            return Ok(EncryptWith::Aes {
+                mode: AesMode::arbitrary(u)?,
+                password: u.arbitrary::<&str>()?,
+            });
+        }
+
+        Ok(EncryptWith::ZipCrypto(
+            ZipCryptoKeys::arbitrary(u)?,
+            PhantomData,
+        ))
+    }
+}
+
 /// Metadata for a file to be written
 #[derive(Clone, Debug, Copy)]
-pub struct FileOptions<T: FileOptionExtension> {
+pub struct FileOptions<'k, T: FileOptionExtension> {
     pub(crate) compression_method: CompressionMethod,
     pub(crate) compression_level: Option<i64>,
     pub(crate) last_modified_time: DateTime,
     pub(crate) permissions: Option<u32>,
     pub(crate) large_file: bool,
-    encrypt_with: Option<ZipCryptoKeys>,
+    encrypt_with: Option<EncryptWith<'k>>,
     extended_options: T,
     alignment: u16,
     #[cfg(feature = "deflate-zopfli")]
     pub(super) zopfli_buffer_size: Option<usize>,
 }
 /// Simple File Options. Can be copied and good for simple writing zip files
-pub type SimpleFileOptions = FileOptions<()>;
+pub type SimpleFileOptions = FileOptions<'static, ()>;
 /// Adds Extra Data and Central Extra Data. It does not implement copy.
-pub type FullFileOptions = FileOptions<ExtendedFileOptions>;
+pub type FullFileOptions<'k> = FileOptions<'k, ExtendedFileOptions>;
 /// The Extension for Extra Data and Central Extra Data
 #[derive(Clone, Debug, Default)]
 pub struct ExtendedFileOptions {
@@ -203,15 +242,15 @@ pub struct ExtendedFileOptions {
 }
 
 #[cfg(fuzzing)]
-impl arbitrary::Arbitrary<'_> for FileOptions<ExtendedFileOptions> {
+impl<'a> arbitrary::Arbitrary<'a> for FileOptions<'a, ExtendedFileOptions> {
-    fn arbitrary(u: &mut arbitrary::Unstructured) -> arbitrary::Result<Self> {
+    fn arbitrary(u: &mut arbitrary::Unstructured<'a>) -> arbitrary::Result<Self> {
         let mut options = FullFileOptions {
             compression_method: CompressionMethod::arbitrary(u)?,
             compression_level: None,
             last_modified_time: DateTime::arbitrary(u)?,
             permissions: Option::<u32>::arbitrary(u)?,
             large_file: bool::arbitrary(u)?,
-            encrypt_with: Option::<ZipCryptoKeys>::arbitrary(u)?,
+            encrypt_with: Option::<EncryptWith>::arbitrary(u)?,
             alignment: u16::arbitrary(u)?,
             #[cfg(feature = "deflate-zopfli")]
             zopfli_buffer_size: None,
@@ -253,7 +292,7 @@ impl arbitrary::Arbitrary<'_> for FileOptions<ExtendedFileOptions> {
     }
 }
 
-impl<T: FileOptionExtension> FileOptions<T> {
+impl<'k, T: FileOptionExtension> FileOptions<'k, T> {
     /// Set the compression method for the new file
     ///
     /// The default is `CompressionMethod::Deflated` if it is enabled. If not,
@@ -317,9 +356,24 @@ impl<T: FileOptionExtension> FileOptions<T> {
         self.large_file = large;
         self
     }
-    pub(crate) fn with_deprecated_encryption(mut self, password: &[u8]) -> Self {
+
-        self.encrypt_with = Some(ZipCryptoKeys::derive(password));
+    pub(crate) fn with_deprecated_encryption(self, password: &[u8]) -> FileOptions<'static, T> {
-        self
+        FileOptions {
+            encrypt_with: Some(EncryptWith::ZipCrypto(
+                ZipCryptoKeys::derive(password),
+                PhantomData,
+            )),
+            ..self
+        }
+    }
+
+    /// Set the AES encryption parameters.
+    #[cfg(feature = "aes-crypto")]
+    pub fn with_aes_encryption(self, mode: AesMode, password: &str) -> FileOptions<'_, T> {
+        FileOptions {
+            encrypt_with: Some(EncryptWith::Aes { mode, password }),
+            ..self
+        }
     }
 
     /// Sets the size of the buffer used to hold the next block that Zopfli will compress. The
@@ -344,7 +398,7 @@ impl<T: FileOptionExtension> FileOptions<T> {
         self
     }
 }
-impl FileOptions<ExtendedFileOptions> {
+impl<'k> FileOptions<'k, ExtendedFileOptions> {
     /// Adds an extra data field.
     pub fn add_extra_data(
         &mut self,
@@ -396,7 +450,7 @@ impl FileOptions<ExtendedFileOptions> {
         self
     }
 }
-impl<T: FileOptionExtension> Default for FileOptions<T> {
+impl<'k, T: FileOptionExtension> Default for FileOptions<'k, T> {
     /// Construct a new FileOptions object
     fn default() -> Self {
         Self {
@@ -700,16 +754,57 @@ impl<W: Write + Seek> ZipWriter<W> {
             uncompressed_size: 0,
         });
 
+        #[allow(unused_mut)]
+        let mut extra_field = options.extended_options.extra_data().cloned();
+
+        // Write AES encryption extra data.
+        #[allow(unused_mut)]
+        let mut aes_extra_data_start = 0;
+        #[cfg(feature = "aes-crypto")]
+        if let Some(EncryptWith::Aes { .. }) = options.encrypt_with {
+            const AES_DUMMY_EXTRA_DATA: [u8; 11] = [
+                0x01, 0x99, 0x07, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+            ];
+
+            let extra_data = extra_field.get_or_insert_with(Default::default);
+            let extra_data = match Arc::get_mut(extra_data) {
+                Some(exclusive) => exclusive,
+                None => {
+                    let new = Arc::new(extra_data.to_vec());
+                    Arc::get_mut(extra_field.insert(new)).unwrap()
+                }
+            };
+
+            if extra_data.len() + AES_DUMMY_EXTRA_DATA.len() > u16::MAX as usize {
+                let _ = self.abort_file();
+                return Err(InvalidArchive("Extra data field is too large"));
+            }
+
+            aes_extra_data_start = extra_data.len() as u64;
+
+            // We write zero bytes for now since we need to update the data when finishing the
+            // file.
+            extra_data.write_all(&AES_DUMMY_EXTRA_DATA)?;
+        }
+
         {
             let header_start = self.inner.get_plain().stream_position()?;
 
             let permissions = options.permissions.unwrap_or(0o100644);
+            let (compression_method, aes_mode) = match options.encrypt_with {
+                #[cfg(feature = "aes-crypto")]
+                Some(EncryptWith::Aes { mode, .. }) => (
+                    CompressionMethod::Aes,
+                    Some((mode, AesVendorVersion::Ae2, options.compression_method)),
+                ),
+                _ => (options.compression_method, None),
+            };
             let file = ZipFileData {
                 system: System::Unix,
                 version_made_by: DEFAULT_VERSION,
                 encrypted: options.encrypt_with.is_some(),
                 using_data_descriptor: false,
-                compression_method: options.compression_method,
+                compression_method,
                 compression_level: options.compression_level,
                 last_modified_time: options.last_modified_time,
                 crc32: raw_values.crc32,
@@ -717,15 +812,18 @@ impl<W: Write + Seek> ZipWriter<W> {
                 uncompressed_size: raw_values.uncompressed_size,
                 file_name: name.into(),
                 file_name_raw: vec![].into_boxed_slice(), // Never used for saving
-                extra_field: options.extended_options.extra_data().cloned(),
+                extra_field,
                 central_extra_field: options.extended_options.central_extra_data().cloned(),
                 file_comment: String::with_capacity(0).into_boxed_str(),
                 header_start,
+                extra_data_start: None,
                 data_start: OnceLock::new(),
                 central_header_start: 0,
                 external_attributes: permissions << 16,
                 large_file: options.large_file,
-                aes_mode: None,
+                aes_mode,
+                aes_extra_data_start,
+
                 extra_fields: Vec::new(),
             };
             let index = self.insert_file_data(file)?;
@@ -777,6 +875,7 @@ impl<W: Write + Seek> ZipWriter<W> {
                 write_local_zip64_extra_field(writer, file)?;
             }
             if let Some(extra_field) = &file.extra_field {
+                file.extra_data_start = Some(writer.stream_position()?);
                 writer.write_all(extra_field)?;
             }
             let mut header_end = writer.stream_position()?;
@@ -816,17 +915,29 @@ impl<W: Write + Seek> ZipWriter<W> {
                     debug_assert_eq!(header_end % align, 0);
                 }
             }
-            if let Some(keys) = options.encrypt_with {
+            match options.encrypt_with {
-                let mut zipwriter = crate::zipcrypto::ZipCryptoWriter {
+                #[cfg(feature = "aes-crypto")]
-                    writer: mem::replace(&mut self.inner, Closed).unwrap(),
+                Some(EncryptWith::Aes { mode, password }) => {
-                    buffer: vec![],
+                    let aeswriter = AesWriter::new(
-                    keys,
+                        mem::replace(&mut self.inner, GenericZipWriter::Closed).unwrap(),
-                };
+                        mode,
-                let crypto_header = [0u8; 12];
+                        password.as_bytes(),
+                    )?;
+                    self.inner = GenericZipWriter::Storer(MaybeEncrypted::Aes(aeswriter));
+                }
+                Some(EncryptWith::ZipCrypto(keys, ..)) => {
+                    let mut zipwriter = crate::zipcrypto::ZipCryptoWriter {
+                        writer: mem::replace(&mut self.inner, Closed).unwrap(),
+                        buffer: vec![],
+                        keys,
+                    };
+                    let crypto_header = [0u8; 12];
 
-                zipwriter.write_all(&crypto_header)?;
+                    zipwriter.write_all(&crypto_header)?;
-                header_end = zipwriter.writer.stream_position()?;
+                    header_end = zipwriter.writer.stream_position()?;
-                self.inner = Storer(MaybeEncrypted::Encrypted(zipwriter));
+                    self.inner = Storer(MaybeEncrypted::ZipCrypto(zipwriter));
+                }
+                None => {}
             }
             self.stats.start = header_end;
             debug_assert!(file.data_start.get().is_none());
@@ -867,13 +978,28 @@ impl<W: Write + Seek> ZipWriter<W> {
                 None => return Ok(()),
                 Some((_, f)) => f,
             };
-            file.crc32 = self.stats.hasher.clone().finalize();
             file.uncompressed_size = self.stats.bytes_written;
 
             let file_end = writer.stream_position()?;
             debug_assert!(file_end >= self.stats.start);
             file.compressed_size = file_end - self.stats.start;
 
+            file.crc32 = self.stats.hasher.clone().finalize();
+            if let Some(aes_mode) = &mut file.aes_mode {
+                // We prefer using AE-1 which provides an extra CRC check, but for small files we
+                // switch to AE-2 to prevent being able to use the CRC value to to reconstruct the
+                // unencrypted contents.
+                //
+                // C.f. https://www.winzip.com/en/support/aes-encryption/#crc-faq
+                aes_mode.1 = if self.stats.bytes_written < 20 {
+                    file.crc32 = 0;
+                    AesVendorVersion::Ae2
+                } else {
+                    AesVendorVersion::Ae1
+                }
+            }
+
+            update_aes_extra_data(writer, file)?;
             update_local_file_header(writer, file)?;
             writer.seek(SeekFrom::Start(file_end))?;
         }
@@ -890,7 +1016,11 @@ impl<W: Write + Seek> ZipWriter<W> {
 
     fn switch_to_non_encrypting_writer(&mut self) -> Result<(), ZipError> {
         match mem::replace(&mut self.inner, Closed) {
-            Storer(MaybeEncrypted::Encrypted(writer)) => {
+            #[cfg(feature = "aes-crypto")]
+            Storer(MaybeEncrypted::Aes(writer)) => {
+                self.inner = Storer(MaybeEncrypted::Unencrypted(writer.finish()?));
+            }
+            Storer(MaybeEncrypted::ZipCrypto(writer)) => {
                 let crc32 = self.stats.hasher.clone().finalize();
                 self.inner = Storer(MaybeEncrypted::Unencrypted(writer.finish(crc32)?))
             }
@@ -1469,7 +1599,7 @@ impl<W: Write + Seek> GenericZipWriter<W> {
                     }))
                 }
                 CompressionMethod::AES => Err(ZipError::UnsupportedArchive(
-                    "AES compression is not supported for writing",
+                    "AES encryption is enabled through FileOptions::with_aes_encryption",
                 )),
                 #[cfg(feature = "zstd")]
                 CompressionMethod::Zstd => {
@@ -1607,6 +1737,50 @@ fn clamp_opt<T: Ord + Copy, U: Ord + Copy + TryFrom<T>>(
     }
 }
 
+fn update_aes_extra_data<W: Write + io::Seek>(
+    writer: &mut W,
+    file: &mut ZipFileData,
+) -> ZipResult<()> {
+    let Some((aes_mode, version, compression_method)) = file.aes_mode else {
+        return Ok(());
+    };
+
+    let extra_data_start = file.extra_data_start.unwrap();
+
+    writer.seek(io::SeekFrom::Start(
+        extra_data_start + file.aes_extra_data_start,
+    ))?;
+
+    let mut buf = Vec::new();
+
+    // Extra field header ID.
+    buf.write_u16_le(0x9901)?;
+    // Data size.
+    buf.write_u16_le(7)?;
+    // Integer version number.
+    buf.write_u16_le(version as u16)?;
+    // Vendor ID.
+    buf.write_all(b"AE")?;
+    // AES encryption strength.
+    buf.write_all(&[aes_mode as u8])?;
+    // Real compression method.
+    #[allow(deprecated)]
+    buf.write_u16_le(compression_method.to_u16())?;
+
+    writer.write_all(&buf)?;
+
+    let aes_extra_data_start = file.aes_extra_data_start as usize;
+    let extra_field = Arc::get_mut(file.extra_field.as_mut().unwrap()).unwrap();
+    extra_field
+        .splice(
+            aes_extra_data_start..(aes_extra_data_start + buf.len()),
+            buf,
+        )
+        .count();
+
+    Ok(())
+}
+
 fn update_local_file_header<T: Write + Seek>(writer: &mut T, file: &ZipFileData) -> ZipResult<()> {
     const CRC32_OFFSET: u64 = 14;
     writer.seek(SeekFrom::Start(file.header_start + CRC32_OFFSET))?;

tests/aes_encryption.rs

@@ -1,7 +1,7 @@
 #![cfg(feature = "aes-crypto")]
 
-use std::io::{self, Read};
+use std::io::{self, Read, Write};
-use zip::ZipArchive;
+use zip::{result::ZipError, write::SimpleFileOptions, AesMode, CompressionMethod, ZipArchive};
 
 const SECRET_CONTENT: &str = "Lorem ipsum dolor sit amet";
 
@@ -74,3 +74,81 @@ fn aes128_encrypted_file() {
         .expect("couldn't read encrypted file");
     assert_eq!(SECRET_CONTENT, content);
 }
+
+#[test]
+fn aes128_stored_roundtrip() {
+    let cursor = {
+        let mut zip = zip::ZipWriter::new(io::Cursor::new(Vec::new()));
+
+        zip.start_file(
+            "test.txt",
+            SimpleFileOptions::default().with_aes_encryption(AesMode::Aes128, "some password"),
+        )
+        .unwrap();
+        zip.write_all(SECRET_CONTENT.as_bytes()).unwrap();
+
+        zip.finish().unwrap()
+    };
+
+    let mut archive = ZipArchive::new(cursor).expect("couldn't open test zip file");
+    test_extract_encrypted_file(&mut archive, "test.txt", "some password", "other password");
+}
+
+#[test]
+fn aes256_deflated_roundtrip() {
+    let cursor = {
+        let mut zip = zip::ZipWriter::new(io::Cursor::new(Vec::new()));
+
+        zip.start_file(
+            "test.txt",
+            SimpleFileOptions::default()
+                .compression_method(CompressionMethod::Deflated)
+                .with_aes_encryption(AesMode::Aes256, "some password"),
+        )
+        .unwrap();
+        zip.write_all(SECRET_CONTENT.as_bytes()).unwrap();
+
+        zip.finish().unwrap()
+    };
+
+    let mut archive = ZipArchive::new(cursor).expect("couldn't open test zip file");
+    test_extract_encrypted_file(&mut archive, "test.txt", "some password", "other password");
+}
+
+fn test_extract_encrypted_file<R: io::Read + io::Seek>(
+    archive: &mut ZipArchive<R>,
+    file_name: &str,
+    correct_password: &str,
+    incorrect_password: &str,
+) {
+    {
+        let file = archive.by_name(file_name).map(|_| ());
+        match file {
+            Err(ZipError::UnsupportedArchive("Password required to decrypt file")) => {}
+            Err(err) => {
+                panic!("Failed to read file for unknown reason: {err:?}");
+            }
+            Ok(_) => {
+                panic!("Was able to successfully read encrypted file without password");
+            }
+        }
+    }
+
+    {
+        match archive.by_name_decrypt(file_name, incorrect_password.as_bytes()) {
+            Err(ZipError::InvalidPassword) => {}
+            Err(err) => panic!("Expected invalid password error, got: {err:?}"),
+            Ok(_) => panic!("Expected invalid password, got decrypted file"),
+        }
+    }
+
+    {
+        let mut content = String::new();
+        archive
+            .by_name_decrypt(file_name, correct_password.as_bytes())
+            .expect("couldn't read encrypted file")
+            .read_to_string(&mut content)
+            .expect("couldn't read encrypted file");
+        assert_eq!(SECRET_CONTENT, content);
+    }
+}