#485 Calculate AES mac with cache and push back functionality

srikanth-lingala · Jan 25, 2023 · ddd8fdc · ddd8fdc
1 parent 597b31a
commit ddd8fdc
Show file tree

Hide file tree

Showing 11 changed files with 110 additions and 23 deletions.
diff --git a/src/main/java/net/lingala/zip4j/crypto/AESDecrypter.java b/src/main/java/net/lingala/zip4j/crypto/AESDecrypter.java
@@ -86,7  86,7 @@ public int decryptData(byte[] buff, int start, int len) throws ZipException {
     return len;
   }
 
-  public byte[] getCalculatedAuthenticationBytes() {
-    return mac.doFinal();
   public byte[] getCalculatedAuthenticationBytes(int numberOfBytesPushedBack) {
     return mac.doFinal(numberOfBytesPushedBack);
   }
 }
diff --git a/src/main/java/net/lingala/zip4j/crypto/PBKDF2/MacBasedPRF.java b/src/main/java/net/lingala/zip4j/crypto/PBKDF2/MacBasedPRF.java
@@ -16,11  16,16 @@
 
 package net.lingala.zip4j.crypto.PBKDF2;
 
 import net.lingala.zip4j.util.InternalZipConstants;
 
 import javax.crypto.Mac;
 import javax.crypto.spec.SecretKeySpec;
 import java.io.ByteArrayOutputStream;
 import java.security.InvalidKeyException;
 import java.security.NoSuchAlgorithmException;
 
 import static net.lingala.zip4j.util.InternalZipConstants.AES_BLOCK_SIZE;
 
 /*
  * Source referred from Matthias Gartner's PKCS#5 implementation -
  * see http://rtner.de/software/PBKDF2.html
@@ -30,9  35,11 @@ public class MacBasedPRF implements PRF {
   private Mac mac;
   private int hLen;
   private String macAlgorithm;
   private ByteArrayOutputStream macCache;
 
   public MacBasedPRF(String macAlgorithm) {
     this.macAlgorithm = macAlgorithm;
     this.macCache = new ByteArrayOutputStream(InternalZipConstants.BUFF_SIZE);
     try {
       mac = Mac.getInstance(macAlgorithm);
       hLen = mac.getMacLength();
@@ -42,10  49,20 @@ public MacBasedPRF(String macAlgorithm) {
   }
 
   public byte[] doFinal(byte[] M) {
     if (macCache.size() > 0) {
       doMacUpdate(0);
     }
     return mac.doFinal(M);
   }
 
   public byte[] doFinal() {
     return doFinal(0);
   }
 
   public byte[] doFinal(int numberOfBytesToPushbackForMac) {
     if (macCache.size() > 0) {
       doMacUpdate(numberOfBytesToPushbackForMac);
     }
     return mac.doFinal();
   }
 
@@ -61,19  78,29 @@ public void init(byte[] P) {
     }
   }
 
-  public void update(byte[] U) {
   public void update(byte[] u) {
     update(u, 0, u.length);
   }
 
   public void update(byte[] u, int start, int len) {
     try {
-      mac.update(U);
       if (macCache.size()   len > InternalZipConstants.BUFF_SIZE) {
         doMacUpdate(0);
       }
       macCache.write(u, start, len);
     } catch (IllegalStateException e) {
       throw new RuntimeException(e);
     }
   }
 
-  public void update(byte[] U, int start, int len) {
-    try {
-      mac.update(U, start, len);
-    } catch (IllegalStateException e) {
-      throw new RuntimeException(e);
   private void doMacUpdate(int numberOfBytesToPushBack) {
     byte[] macBytes = macCache.toByteArray();
     int numberOfBytesToRead = macBytes.length - numberOfBytesToPushBack;
     int updateLength;
     for (int i = 0; i < numberOfBytesToRead; i  = InternalZipConstants.AES_BLOCK_SIZE) {
       updateLength = (i   AES_BLOCK_SIZE) <= numberOfBytesToRead ? AES_BLOCK_SIZE : numberOfBytesToRead - i;
       mac.update(macBytes, i, updateLength);
     }
     macCache.reset();
   }
 }
diff --git a/src/main/java/net/lingala/zip4j/io/inputstream/AesCipherInputStream.java b/src/main/java/net/lingala/zip4j/io/inputstream/AesCipherInputStream.java
@@ -117,12  117,12 @@ private void copyBytesFromBuffer(byte[] b, int off) {
   }
 
   @Override
-  protected void endOfEntryReached(InputStream inputStream) throws IOException {
-    verifyContent(readStoredMac(inputStream));
   protected void endOfEntryReached(InputStream inputStream, int numberOfBytesPushedBack) throws IOException {
     verifyContent(readStoredMac(inputStream), numberOfBytesPushedBack);
   }
 
-  private void verifyContent(byte[] storedMac) throws IOException {
-    byte[] calculatedMac = getDecrypter().getCalculatedAuthenticationBytes();
   private void verifyContent(byte[] storedMac, int numberOfBytesPushedBack) throws IOException {
     byte[] calculatedMac = getDecrypter().getCalculatedAuthenticationBytes(numberOfBytesPushedBack);
     byte[] first10BytesOfCalculatedMac = new byte[AES_AUTH_LENGTH];
     System.arraycopy(calculatedMac, 0, first10BytesOfCalculatedMac, 0, InternalZipConstants.AES_AUTH_LENGTH);
 

diff --git a/src/main/java/net/lingala/zip4j/io/inputstream/CipherInputStream.java b/src/main/java/net/lingala/zip4j/io/inputstream/CipherInputStream.java
@@ -80,7  80,7 @@ public T getDecrypter() {
     return decrypter;
   }
 
-  protected void endOfEntryReached(InputStream inputStream) throws IOException {
   protected void endOfEntryReached(InputStream inputStream, int numberOfBytesPushedBack) throws IOException {
     // is optional but useful for AES
   }
 

diff --git a/src/main/java/net/lingala/zip4j/io/inputstream/DecompressedInputStream.java b/src/main/java/net/lingala/zip4j/io/inputstream/DecompressedInputStream.java
@@ -39,12  39,13 @@ public void close() throws IOException {
     cipherInputStream.close();
   }
 
-  public void endOfEntryReached(InputStream inputStream) throws IOException {
-    cipherInputStream.endOfEntryReached(inputStream);
   public void endOfEntryReached(InputStream inputStream, int numberOfBytesPushedBack) throws IOException {
     cipherInputStream.endOfEntryReached(inputStream, numberOfBytesPushedBack);
   }
 
-  public void pushBackInputStreamIfNecessary(PushbackInputStream pushbackInputStream) throws IOException {
   public int pushBackInputStreamIfNecessary(PushbackInputStream pushbackInputStream) throws IOException {
     // Do nothing by default
     return 0;
   }
 
   protected byte[] getLastReadRawDataCache() {

diff --git a/src/main/java/net/lingala/zip4j/io/inputstream/InflaterInputStream.java b/src/main/java/net/lingala/zip4j/io/inputstream/InflaterInputStream.java
@@ -55,21  55,22 @@ public int read(byte[] b, int off, int len) throws IOException {
   }
 
   @Override
-  public void endOfEntryReached(InputStream inputStream) throws IOException {
   public void endOfEntryReached(InputStream inputStream, int numberOfBytesPushedBack) throws IOException {
     if (inflater != null) {
       inflater.end();
       inflater = null;
     }
-    super.endOfEntryReached(inputStream);
     super.endOfEntryReached(inputStream, numberOfBytesPushedBack);
   }
 
   @Override
-  public void pushBackInputStreamIfNecessary(PushbackInputStream pushbackInputStream) throws IOException {
   public int pushBackInputStreamIfNecessary(PushbackInputStream pushbackInputStream) throws IOException {
     int n = inflater.getRemaining();
     if (n > 0) {
       byte[] rawDataCache = getLastReadRawDataCache();
       pushbackInputStream.unread(rawDataCache, len - n, n);
     }
     return n;
   }
 
   @Override

diff --git a/src/main/java/net/lingala/zip4j/io/inputstream/ZipInputStream.java b/src/main/java/net/lingala/zip4j/io/inputstream/ZipInputStream.java
@@ -231,10  231,10 @@ public void setPassword(char[] password) {
   private void endOfCompressedDataReached() throws IOException {
     //With inflater, without knowing the compressed or uncompressed size, we over read necessary data
     //In such cases, we have to push back the inputstream to the end of data
-    decompressedInputStream.pushBackInputStreamIfNecessary(inputStream);
     int numberOfBytesPushedBack = decompressedInputStream.pushBackInputStreamIfNecessary(inputStream);
 
     //First signal the end of data for this entry so that ciphers can read any header data if applicable
-    decompressedInputStream.endOfEntryReached(inputStream);
     decompressedInputStream.endOfEntryReached(inputStream, numberOfBytesPushedBack);
 
     readExtendedLocalFileHeaderIfPresent();
     verifyCrc();

diff --git a/src/test/java/net/lingala/zip4j/MiscZipFileIT.java b/src/test/java/net/lingala/zip4j/MiscZipFileIT.java
@@ -673,6  673,16 @@ public void testAddFileWithCustomLastModifiedFileTimeSetsInputTime() throws IOEx
     verifyLastModifiedFileTime(zipFile, fileToTestWith, expectedLastModifiedTimeInMillis);
   }
 
   @Test
   public void testExtractFileWithExtraDataRecordAndCorruptMac() throws ZipException {
     ZipFile zipFile = new ZipFile(getTestArchiveFromResources("aes_with_extra_data_record_and_corrupt_mac.zip"), PASSWORD);
 
     expectedException.expect(ZipException.class);
     expectedException.expectMessage("java.io.IOException: Reached end of data for this entry, but aes verification failed");
 
     zipFile.extractAll(outputFolder.getPath());
   }
 
   private void testAddAndExtractWithPasswordUtf8Encoding(boolean useUtf8ForPassword) throws IOException {
     char[] password = "hun 焰".toCharArray();
     ZipFile zipFile = new ZipFile(generatedZipFile, password);

diff --git a/src/test/java/net/lingala/zip4j/io/inputstream/ZipInputStreamIT.java b/src/test/java/net/lingala/zip4j/io/inputstream/ZipInputStreamIT.java
@@ -24,10  24,10 @@
 import java.io.InputStream;
 import java.io.OutputStream;
 import java.nio.file.Files;
 import java.security.SecureRandom;
 import java.util.ArrayList;
 import java.util.HashSet;
 import java.util.List;
-import java.security.SecureRandom;
 import java.util.Set;
 import java.util.zip.ZipEntry;
 import java.util.zip.ZipOutputStream;
@@ -351,6  351,15 @@ public void readingJarFileWithCompressedSizeNotZeroForDirectoryIsSuccessful() th
     }
   }
 
   @Test
   public void testExtractZipFileWithExtraDataRecordAndCorruptAesMacFails() throws IOException {
     expectedException.expect(IOException.class);
     expectedException.expectMessage("Reached end of data for this entry, but aes verification failed");
 
     extractZipFileWithInputStreams(TestUtils.getTestArchiveFromResources("aes_with_extra_data_record_and_corrupt_mac.zip"),
             PASSWORD, InternalZipConstants.BUFF_SIZE, false, 1);
   }
 
   private void extractZipFileWithInputStreams(File zipFile, char[] password) throws IOException {
     extractZipFileWithInputStreams(zipFile, password, InternalZipConstants.BUFF_SIZE);
   }

diff --git a/src/test/java/net/lingala/zip4j/io/outputstream/ZipOutputStreamIT.java b/src/test/java/net/lingala/zip4j/io/outputstream/ZipOutputStreamIT.java
@@ -30,6  30,8 @@
 import java.io.OutputStream;
 import java.nio.charset.Charset;
 import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
 import java.nio.file.Paths;
 import java.util.ArrayList;
 import java.util.List;
 
@@ -310,6  312,28 @@ public void testLastModifiedTimeIsSetWhenItIsNotExplicitlySet() throws IOExcepti
     }
   }
 
   @Test
   public void testZipInputStreamWithDeflateAndAesEncryption() throws IOException {
     byte[] buffer = new byte[InternalZipConstants.BUFF_SIZE];
     int readLen;
     File fileToAdd = getTestFileFromResources("file_PDF_1MB.pdf");
     try (ZipOutputStream zipOutputStream = new ZipOutputStream(Files.newOutputStream(generatedZipFile.toPath()), PASSWORD)) {
       ZipParameters zipParameters = new ZipParameters();
       zipParameters.setFileNameInZip(fileToAdd.getName());
       zipParameters.setEncryptFiles(true);
       zipParameters.setEncryptionMethod(EncryptionMethod.AES);
       zipOutputStream.putNextEntry(zipParameters);
       try (InputStream inputStream = Files.newInputStream(fileToAdd.toPath())) {
         while ((readLen = inputStream.read(buffer)) != -1) {
           zipOutputStream.write(buffer, 0, readLen);
         }
       }
     }
 
     verifyZipFileByExtractingAllFiles(generatedZipFile, PASSWORD, outputFolder, 1, true);
     extractZipFileWithInputStream(generatedZipFile);
   }
 
   private void testZipOutputStream(CompressionMethod compressionMethod, boolean encrypt,
                                    EncryptionMethod encryptionMethod, AesKeyStrength aesKeyStrength,
                                    AesVersion aesVersion)
@@ -513,4  537,19 @@ private void putNextEntryAndCloseEntry(ZipOutputStream zipOutputStream, String f
     zipOutputStream.putNextEntry(zipParameters);
     zipOutputStream.closeEntry();
   }
 
   private void extractZipFileWithInputStream(File zipFile) throws IOException {
     byte[] buffer = new byte[InternalZipConstants.BUFF_SIZE];
     int readLen;
     LocalFileHeader lfh;
     try (ZipInputStream zipInputStream = new ZipInputStream(Files.newInputStream(zipFile.toPath()), PASSWORD)) {
       while ((lfh = zipInputStream.getNextEntry()) != null) {
         while ((readLen = zipInputStream.read(buffer)) != -1) {
           try (OutputStream outputStream = Files.newOutputStream(Paths.get(outputFolder.getPath(), lfh.getFileName()))) {
             outputStream.write(buffer, 0, readLen);
           }
         }
       }
     }
   }
 }
diff --git a/src/test/resources/test-archives/aes_with_extra_data_record_and_corrupt_mac.zip b/src/test/resources/test-archives/aes_with_extra_data_record_and_corrupt_mac.zip