Avoid infinite loop for certain corrupt axmls

CalebFenton · Jun 4, 2017 · 5b6ce0d · 5b6ce0d
1 parent d48e745
commit 5b6ce0d
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 12 deletions.
diff --git a/src/main/java/org/cf/apkfile/apk/Resources.java b/src/main/java/org/cf/apkfile/apk/Resources.java
@@ -1,17 +1,12 @@
 package org.cf.apkfile.apk;
 
+import org.cf.apkfile.res.Chunk;
 import org.cf.apkfile.res.ResourceFile;
 import org.cf.apkfile.res.ResourceTableChunk;
 
 import java.io.IOException;
 import java.io.InputStream;
-import java.io.StringWriter;
-import java.util.Arrays;
-import java.util.Collection;
-import java.util.HashMap;
-import java.util.LinkedList;
-import java.util.List;
-import java.util.Map;
+import java.util.*;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
@@ -30,8 +25,12 @@ public class Resources {
     private final Map<String, Map<String, Object>> resourceConfigs;
 
     public Resources(InputStream resourcesStream) throws IOException {
-        ResourceFile arscRF = ResourceFile.fromInputStream(resourcesStream);
-        resourceTable = (ResourceTableChunk) arscRF.getChunks().get(0);
+        ResourceFile resourceFile = ResourceFile.fromInputStream(resourcesStream);
+        Chunk chunk = resourceFile.getChunks().get(0);
+        if (!(chunk instanceof ResourceTableChunk)) {
+            throw new RuntimeException("Resources file appears to be invalid");
+        }
+        resourceTable = (ResourceTableChunk) chunk;
 
         entries = new HashMap<>();
         resourceConfigs = new HashMap<>();

diff --git a/src/main/java/org/cf/apkfile/res/ChunkWithChunks.java b/src/main/java/org/cf/apkfile/res/ChunkWithChunks.java
@@ -40,7 +40,7 @@ protected void init(ByteBuffer buffer) {
         chunks.clear();
         int start = this.offset + getHeaderSize();
         int offset = start;
-        int end = this.offset + getOriginalChunkSize();
+        int end = this.offset + Math.min(getOriginalChunkSize(), buffer.remaining());
         int position = buffer.position();
         buffer.position(start);
 

diff --git a/src/main/java/org/cf/apkfile/res/UnknownChunk.java b/src/main/java/org/cf/apkfile/res/UnknownChunk.java
@@ -42,14 +42,14 @@ protected UnknownChunk(ByteBuffer buffer, @Nullable Chunk parent) throws Illegal
             Logger.warn("Header size (" + headerSize + ") > chunk size (" + chunkSize + ")");
             header = new byte[0];
             payload = new byte[0];
-            dummyChunkSize = Math.min(headerSize + chunkSize, buffer.remaining());
+            dummyChunkSize = Math.min(headerSize + chunkSize, buffer.remaining() + Chunk.METADATA_SIZE);
             return;
         }
         if (headerSize + chunkSize > buffer.remaining()) {
             Logger.warn("Chunk size (" + (headerSize + chunkSize) + ") greater than remaining buffer (" + buffer.remaining() + ")");
             header = new byte[0];
             payload = new byte[0];
-            dummyChunkSize = buffer.remaining();
+            dummyChunkSize = Math.max(buffer.remaining(), Chunk.METADATA_SIZE);
             return;
         }
         if (headerSize == 0) {