Return class header detection back & some fixes

Col-E · Jan 7, 2022 · d0cf108 · d0cf108
1 parent 41bd61d
commit d0cf108
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 11 deletions.
diff --git a/src/main/java/me/coley/recaf/util/ClassUtil.java b/src/main/java/me/coley/recaf/util/ClassUtil.java
@@ -298,7 +298,7 @@ public static int getVersion(byte[] code) {
 	 * @return {@code true} if data has class magic prefix.
 	 */
 	public static boolean isClass(byte[] data) {
-		return data.length > 4 &&
+		return data.length >= 4 &&
 				0xCAFEBABEL == ((
 						(0xFF & data[0]) << 24L |
 						(0xFF & data[1]) << 16L |

diff --git a/src/main/java/me/coley/recaf/util/IOUtil.java b/src/main/java/me/coley/recaf/util/IOUtil.java
@@ -42,7 +42,8 @@ private IOUtil() {
     public static int transfer(InputStream in, OutputStream out, byte[] buffer, int max) throws IOException {
         int transferred = 0;
         int r;
-        while ((max == ANY || max > 0) && (r = in.read(buffer, 0, buffer.length)) != -1) {
+        while ((max == ANY || max > 0) && (r = in.read(buffer, 0,
+                max == ANY ? buffer.length : Math.min(buffer.length, max))) != -1) {
             transferred += r;
             out.write(buffer, 0, r);
             if (max != ANY) {

diff --git a/src/main/java/me/coley/recaf/workspace/JarResource.java b/src/main/java/me/coley/recaf/workspace/JarResource.java
@@ -36,20 +36,27 @@ protected Map<String, byte[]> loadClasses() throws IOException {
 		byte[] buffer = new byte[8192];
 		EntryLoader loader = getEntryLoader();
 
-		try {
-			ZipInputStream zis = new ZipInputStream(new FileInputStream(getPath().toFile()));
+		try (ZipInputStream zis = new ZipInputStream(new FileInputStream(getPath().toFile()))) {
 			ZipEntry entry;
 
 			while ((entry = zis.getNextEntry()) != null) {
 				// verify entries are classes and valid files
 				// - skip intentional garbage / zip file abnormalities
 				if (shouldSkip(entry.getName()))
 					continue;
-				if (!loader.isValidClassEntry(entry))
-					continue;
 
 				out.reset();
-				byte[] in = IOUtil.toByteArray(zis, out, buffer);
+				byte[] in;
+				if (!loader.isValidClassEntry(entry)) {
+					// The class file might not end with .class or .class/
+					// so we also check it's header.
+					in = IOUtil.toByteArray(zis, out, buffer, 4);
+					if (!loader.isValidClassFile(new ByteArrayInputStream(in))) {
+						continue;
+					}
+				}
+
+				in = IOUtil.toByteArray(zis, out, buffer);
 
 				// There is no possible way a "class" under 30 bytes is valid
 				if (in.length < 30)
@@ -70,11 +77,27 @@ protected Map<String, byte[]> loadClasses() throws IOException {
 
 						if (shouldSkip(entry.getName()))
 							continue;
-						if (!loader.isValidClassEntry(entry))
-							continue;
 
-						InputStream zis = zf.getInputStream(entry);
-						byte[] in = IOUtil.toByteArray(zis);
+						out.reset();
+						byte[] in;
+
+						if (!loader.isValidClassEntry(entry)) {
+							// The class file might not end with .class or .class/
+							// so we also check it's header.
+							out.reset();
+							try (InputStream zis = zf.getInputStream(entry)) {
+								in = IOUtil.toByteArray(zis, out, buffer, 4);
+							}
+							if (!loader.isValidClassFile(new ByteArrayInputStream(in))) {
+								continue;
+							}
+						}
+
+						out.reset();
+						try (InputStream zis = zf.getInputStream(entry)) {
+							in = IOUtil.toByteArray(zis, out, buffer);
+						}
+
 						loader.onClass(entry.getName(), in);
 					}
 				}