diff --git a/src/main/kotlin/org/moe/tools/classvalidator/ClassModifier.kt b/src/main/kotlin/org/moe/tools/classvalidator/ClassModifier.kt new file mode 100644 index 0000000..1bee01b --- /dev/null +++ b/src/main/kotlin/org/moe/tools/classvalidator/ClassModifier.kt @@ -0,0 +1,16 @@ +package org.moe.tools.classvalidator + +import org.objectweb.asm.ClassVisitor + +/** + * A [ClassVisitor] that could modify a class + */ +abstract class ClassModifier(api: Int, classVisitor: ClassVisitor?) : ClassVisitor(api, classVisitor) { + + var modified: Boolean = false + private set + + fun markAsModified() { + modified = true + } +} diff --git a/src/main/kotlin/org/moe/tools/classvalidator/ClassSaver.kt b/src/main/kotlin/org/moe/tools/classvalidator/ClassSaver.kt index f4c3ea0..2782ee4 100644 --- a/src/main/kotlin/org/moe/tools/classvalidator/ClassSaver.kt +++ b/src/main/kotlin/org/moe/tools/classvalidator/ClassSaver.kt @@ -1,21 +1,32 @@ package org.moe.tools.classvalidator -import org.moe.common.utils.prepareDir import org.objectweb.asm.ClassReader +import java.nio.file.FileSystems import java.nio.file.Files import java.nio.file.Path internal class ClassSaver( private val outputDir: Path ) { - fun save(bytecode: ByteArray?) { + + private val classNameBytecodeMap = mutableMapOf() + fun add(bytecode: ByteArray?) { if (bytecode == null) { return } val cr = ClassReader(bytecode) - val outputFile = outputDir.resolve(cr.className + ".class") - outputFile.parent.prepareDir() - Files.write(outputFile, bytecode) + classNameBytecodeMap[cr.className + ".class"] = bytecode + } + + fun save() { + // In jar replacement + val fileSystem = FileSystems.newFileSystem(outputDir, null) + fileSystem.use { + classNameBytecodeMap.forEach { (fileName, bytecode) -> + val outputFile = it.getPath(fileName) + Files.write(outputFile, bytecode) + } + } } } diff --git a/src/main/kotlin/org/moe/tools/classvalidator/ClassValidator.kt b/src/main/kotlin/org/moe/tools/classvalidator/ClassValidator.kt index 4b867a3..7d42667 100644 --- a/src/main/kotlin/org/moe/tools/classvalidator/ClassValidator.kt +++ b/src/main/kotlin/org/moe/tools/classvalidator/ClassValidator.kt @@ -1,37 +1,49 @@ package org.moe.tools.classvalidator -import org.moe.common.utils.classAndJarInputIterator +import org.moe.common.utils.classpathIterator import org.moe.tools.classvalidator.natj.AddMissingAnnotations import org.moe.tools.classvalidator.natj.AddMissingNatJRegister import org.objectweb.asm.ClassReader import org.objectweb.asm.ClassVisitor import org.objectweb.asm.ClassWriter import java.io.File -import java.nio.file.Path object ClassValidator { fun process( - inputFiles: Set, - outputDir: Path, - classpath: Set, + inputFiles: Set, + classpath: Set, ) { - ContextClassLoaderHolder( - ChildFirstClassLoader(classpath.map { it.toURI().toURL() }.toTypedArray()) - ).use { - val classSaver = ClassSaver(outputDir.resolve(OUTPUT_CLASSES)) - - inputFiles.classAndJarInputIterator { _, inputStream -> - val cr = ClassReader(inputStream) - - val byteCode = processClass(cr) { next -> - next - .let(::AddMissingAnnotations) - .let(::AddMissingNatJRegister) + val classSavers = mutableListOf() + ChildFirstClassLoader(classpath.map { it.toURI().toURL() }.toTypedArray()).use { loader -> + ContextClassLoaderHolder(loader).use { + inputFiles.forEach { jar -> + val classSaver = ClassSaver(jar.absoluteFile.toPath()) + jar.classpathIterator({ _, inputStream -> + val cr = ClassReader(inputStream) + + val chain = mutableListOf() + fun ClassVisitor.chain(nextBuilder: (ClassVisitor) -> ClassVisitor): ClassVisitor { + val next = nextBuilder(this) + chain.add(next) + return next + } + + val byteCode = processClass(cr) { next -> + next + .chain(::AddMissingAnnotations) + .chain(::AddMissingNatJRegister) + } + + // Only save modified class + if (chain.any { it is ClassModifier && it.modified }) { + classSaver.add(byteCode) + } + }, { it.endsWith(".class") }) + classSavers.add(classSaver) } - - classSaver.save(byteCode) } } + classSavers.forEach { it.save() } } private inline fun processClass(reader: ClassReader, chain: (ClassVisitor) -> ClassVisitor): ByteArray { @@ -47,6 +59,4 @@ object ClassValidator { return writer.toByteArray() } - - const val OUTPUT_CLASSES = "classes" } diff --git a/src/main/kotlin/org/moe/tools/classvalidator/natj/AddMissingAnnotations.kt b/src/main/kotlin/org/moe/tools/classvalidator/natj/AddMissingAnnotations.kt index 20581a0..14ecc91 100644 --- a/src/main/kotlin/org/moe/tools/classvalidator/natj/AddMissingAnnotations.kt +++ b/src/main/kotlin/org/moe/tools/classvalidator/natj/AddMissingAnnotations.kt @@ -2,6 +2,7 @@ package org.moe.tools.classvalidator.natj import org.moe.tools.classvalidator.getParentImplementation import org.moe.tools.classvalidator.toClass +import org.moe.tools.classvalidator.ClassModifier import org.objectweb.asm.AnnotationVisitor import org.objectweb.asm.ClassVisitor import org.objectweb.asm.MethodVisitor @@ -11,7 +12,7 @@ import java.lang.reflect.AnnotatedElement class AddMissingAnnotations( next: ClassVisitor? -) : ClassVisitor(Opcodes.ASM5, next) { +) : ClassModifier(Opcodes.ASM5, next) { private var skip: Boolean = false private var superName: String? = null private var interfaces: Array? = null @@ -130,6 +131,7 @@ class AddMissingAnnotations( } } } + markAsModified() } private fun convertClassToType(value: Any): Any { diff --git a/src/main/kotlin/org/moe/tools/classvalidator/natj/AddMissingNatJRegister.kt b/src/main/kotlin/org/moe/tools/classvalidator/natj/AddMissingNatJRegister.kt index c3a892b..145dc63 100644 --- a/src/main/kotlin/org/moe/tools/classvalidator/natj/AddMissingNatJRegister.kt +++ b/src/main/kotlin/org/moe/tools/classvalidator/natj/AddMissingNatJRegister.kt @@ -1,5 +1,6 @@ package org.moe.tools.classvalidator.natj +import org.moe.tools.classvalidator.ClassModifier import org.objectweb.asm.AnnotationVisitor import org.objectweb.asm.ClassVisitor import org.objectweb.asm.MethodVisitor @@ -13,7 +14,7 @@ import org.objectweb.asm.tree.MethodNode */ class AddMissingNatJRegister( next: ClassVisitor? -) : ClassVisitor(Opcodes.ASM5, next) { +) : ClassModifier(Opcodes.ASM5, next) { private var skip: Boolean = false private var visit: Boolean = false @@ -89,6 +90,7 @@ class AddMissingNatJRegister( NatJRuntime.NATJ_OWNER, NatJRuntime.NATJ_REGISTER_NAME, NatJRuntime.NATJ_REGISTER_DESC, false) CLI.accept(mv) + markAsModified() println("Injected NatJ.register() into $name") } NATJREG_LEAVE_ALONE -> { @@ -107,6 +109,7 @@ class AddMissingNatJRegister( mv.visitInsn(Opcodes.RETURN) mv.visitMaxs(-1, -1) mv.visitEnd() + markAsModified() println("Injected NatJ.register() into $name") }