From 00378bc6f2b782af32a249082d1806ee13850ab0 Mon Sep 17 00:00:00 2001
From: Chris Hennick <hennickc@amazon.com>
Date: Tue, 23 Apr 2024 12:07:30 -0700
Subject: [PATCH] fix: Alignment was previously handled incorrectly (#33)

---
 src/read.rs  |  2 +-
 src/write.rs | 33 ++++++++++++++++++++++-----------
 2 files changed, 23 insertions(+), 12 deletions(-)

diff --git a/src/read.rs b/src/read.rs
index dac4df58..fdebcad9 100644
--- a/src/read.rs
+++ b/src/read.rs
@@ -1447,7 +1447,7 @@ mod test {
         v.extend_from_slice(include_bytes!("../tests/data/deflate64_issue_25.zip"));
         ZipArchive::new(Cursor::new(v)).expect_err("Invalid file");
     }
-    
+
     #[test]
     fn test_read_with_data_descriptor() {
         let mut v = Vec::new();
diff --git a/src/write.rs b/src/write.rs
index 5d85d5ad..aaa6b3a5 100644
--- a/src/write.rs
+++ b/src/write.rs
@@ -761,25 +761,36 @@ impl<W: Write + Seek> ZipWriter<W> {
             let mut header_end = writer.stream_position()?;
             if options.alignment > 1 {
                 let align = options.alignment as u64;
-                if header_end % align != 0 {
-                    let pad_length = (align - (header_end + 4) % align) % align;
-                    if pad_length + extra_field_length as u64 > u16::MAX as u64 {
+                let unaligned_header_bytes = header_end % align;
+                if unaligned_header_bytes != 0 {
+                    let pad_length = (align - unaligned_header_bytes) as usize;
+                    let Some(new_extra_field_length) =
+                        (pad_length as u16).checked_add(extra_field_length)
+                    else {
                         let _ = self.abort_file();
                         return Err(InvalidArchive(
                             "Extra data field would be larger than allowed after aligning",
                         ));
+                    };
+                    if pad_length >= 4 {
+                        // Add an extra field to the extra_data
+                        let pad_body = vec![0; pad_length - 4];
+                        writer.write_all(b"za").map_err(ZipError::from)?; // 0x617a
+                        writer
+                            .write_u16::<LittleEndian>(pad_body.len() as u16)
+                            .map_err(ZipError::from)?;
+                        writer.write_all(&pad_body).map_err(ZipError::from)?;
+                    } else {
+                        // extra_data padding is too small for an extra field header, so pad with
+                        // zeroes
+                        let pad = vec![0; pad_length];
+                        writer.write_all(&pad).map_err(ZipError::from)?;
                     }
-                    let pad = vec![0; pad_length as usize];
-                    writer.write_all(b"za").map_err(ZipError::from)?; // 0x617a
-                    writer
-                        .write_u16::<LittleEndian>(pad.len() as u16)
-                        .map_err(ZipError::from)?;
-                    writer.write_all(&pad).map_err(ZipError::from)?;
                     header_end = writer.stream_position()?;
 
                     // Update extra field length in local file header.
                     writer.seek(SeekFrom::Start(file.header_start + 28))?;
-                    writer.write_u16::<LittleEndian>(pad_length as u16 + extra_field_length)?;
+                    writer.write_u16::<LittleEndian>(new_extra_field_length)?;
                     writer.seek(SeekFrom::Start(header_end))?;
                     debug_assert_eq!(header_end % align, 0);
                 }
@@ -2210,7 +2221,7 @@ mod test {
         writer.finish()?;
         Ok(())
     }
-    
+
     #[test]
     fn test_alignment() {
         let page_size = 4096;