diff options
Diffstat (limited to 'src')
-rw-r--r-- | src/main/kotlin/moe/nea/shot/ClassRef.kt | 17 | ||||
-rw-r--r-- | src/main/kotlin/moe/nea/shot/ClassShots.kt | 9 | ||||
-rw-r--r-- | src/main/kotlin/moe/nea/shot/FieldRef.kt | 5 | ||||
-rw-r--r-- | src/main/kotlin/moe/nea/shot/FieldShots.kt | 5 | ||||
-rw-r--r-- | src/main/kotlin/moe/nea/shot/InjectionApplicator.kt | 69 | ||||
-rw-r--r-- | src/main/kotlin/moe/nea/shot/MethodRef.kt | 5 | ||||
-rw-r--r-- | src/main/kotlin/moe/nea/shot/MethodShots.kt | 8 | ||||
-rw-r--r-- | src/main/kotlin/moe/nea/shot/ShotParser.kt | 149 | ||||
-rw-r--r-- | src/main/kotlin/moe/nea/shot/Shots.kt | 55 |
9 files changed, 322 insertions, 0 deletions
diff --git a/src/main/kotlin/moe/nea/shot/ClassRef.kt b/src/main/kotlin/moe/nea/shot/ClassRef.kt new file mode 100644 index 0000000..6fdd9d1 --- /dev/null +++ b/src/main/kotlin/moe/nea/shot/ClassRef.kt @@ -0,0 +1,17 @@ +package moe.nea.shot + +import java.io.Serializable + + +data class ClassRef(val className: String) : Serializable { + val path: String get() = "$nudeJvmRef.class" + val jvmRef: String get() = "L$nudeJvmRef;" + val nudeJvmRef: String get() = className.replace(".", "/") + + companion object { + fun fromPath(path: String): ClassRef? { + if (!path.endsWith(".class")) return null + return ClassRef(path.removeSuffix(".class").replace("/", ".")) + } + } +}
\ No newline at end of file diff --git a/src/main/kotlin/moe/nea/shot/ClassShots.kt b/src/main/kotlin/moe/nea/shot/ClassShots.kt new file mode 100644 index 0000000..066ba82 --- /dev/null +++ b/src/main/kotlin/moe/nea/shot/ClassShots.kt @@ -0,0 +1,9 @@ +package moe.nea.shot + +import java.io.Serializable + +data class ClassShots( + val annotations: List<ClassRef>, + val methodShots: Map<MethodRef, MethodShots>, + val fieldShots: Map<FieldRef, FieldShots>, +) : Serializable
\ No newline at end of file diff --git a/src/main/kotlin/moe/nea/shot/FieldRef.kt b/src/main/kotlin/moe/nea/shot/FieldRef.kt new file mode 100644 index 0000000..007c79b --- /dev/null +++ b/src/main/kotlin/moe/nea/shot/FieldRef.kt @@ -0,0 +1,5 @@ +package moe.nea.shot + +import java.io.Serializable + +data class FieldRef(val name: String) : Serializable
\ No newline at end of file diff --git a/src/main/kotlin/moe/nea/shot/FieldShots.kt b/src/main/kotlin/moe/nea/shot/FieldShots.kt new file mode 100644 index 0000000..36217c5 --- /dev/null +++ b/src/main/kotlin/moe/nea/shot/FieldShots.kt @@ -0,0 +1,5 @@ +package moe.nea.shot + +import java.io.Serializable + +data class FieldShots(val annotations: List<ClassRef>) : Serializable
\ No newline at end of file diff --git a/src/main/kotlin/moe/nea/shot/InjectionApplicator.kt b/src/main/kotlin/moe/nea/shot/InjectionApplicator.kt new file mode 100644 index 0000000..82b05a6 --- /dev/null +++ b/src/main/kotlin/moe/nea/shot/InjectionApplicator.kt @@ -0,0 +1,69 @@ +package moe.nea.shot + +import org.objectweb.asm.* + +class InjectionApplicator( + val injections: ClassShots, + writer: ClassWriter +) : ClassVisitor(Opcodes.ASM9, writer) { + + override fun visit( + version: Int, + access: Int, + name: String?, + signature: String?, + superName: String?, + interfaces: Array<out String>? + ) { + super.visit(version, access, name, signature, superName, interfaces) + for (annotation in injections.annotations) { + super.visitAnnotation(annotation.jvmRef, false).visitEnd() + } + } + + override fun visitField( + access: Int, + name: String, + descriptor: String?, + signature: String?, + value: Any? + ): FieldVisitor { + val fieldVisitor = super.visitField(access, name, descriptor, signature, value) + val fieldShots = injections.fieldShots[FieldRef(name)] ?: return fieldVisitor + for (annotation in fieldShots.annotations) { + fieldVisitor.visitAnnotation(annotation.jvmRef, false) + .visitEnd() + } + return fieldVisitor + } + + override fun visitMethod( + access: Int, + name: String, + descriptor: String, + signature: String?, + exceptions: Array<out String>? + ): MethodVisitor { + val parentVisitor = super.visitMethod(access, name, descriptor, signature, exceptions) + val methodShots = injections.methodShots[MethodRef(name, descriptor.replaceAfter(")", ""))] ?: return parentVisitor + for (injection in methodShots.annotations) { + parentVisitor.visitAnnotation(injection.jvmRef, false) + .visitEnd() + } + var maxParameterCount = 0 + for ((parameter, annotations) in methodShots.parameterAnnotations) { + maxParameterCount = maxOf(maxParameterCount, parameter + 1) + for (annotation in annotations) { + parentVisitor + .visitParameterAnnotation(parameter, annotation.jvmRef, false) + .visitEnd() + } + } + return object : MethodVisitor(Opcodes.ASM8, parentVisitor) { + override fun visitAnnotableParameterCount(parameterCount: Int, visible: Boolean) { + super.visitAnnotableParameterCount( + if (visible) parameterCount else maxOf(parameterCount, maxParameterCount), visible) + } + } + } +}
\ No newline at end of file diff --git a/src/main/kotlin/moe/nea/shot/MethodRef.kt b/src/main/kotlin/moe/nea/shot/MethodRef.kt new file mode 100644 index 0000000..7b07efc --- /dev/null +++ b/src/main/kotlin/moe/nea/shot/MethodRef.kt @@ -0,0 +1,5 @@ +package moe.nea.shot + +import java.io.Serializable + +data class MethodRef(val name: String, val argumentDescriptor: String) : Serializable diff --git a/src/main/kotlin/moe/nea/shot/MethodShots.kt b/src/main/kotlin/moe/nea/shot/MethodShots.kt new file mode 100644 index 0000000..59893e2 --- /dev/null +++ b/src/main/kotlin/moe/nea/shot/MethodShots.kt @@ -0,0 +1,8 @@ +package moe.nea.shot + +import java.io.Serializable + +data class MethodShots( + val annotations: List<ClassRef>, + val parameterAnnotations: Map<Int, List<ClassRef>>, +) : Serializable
\ No newline at end of file diff --git a/src/main/kotlin/moe/nea/shot/ShotParser.kt b/src/main/kotlin/moe/nea/shot/ShotParser.kt new file mode 100644 index 0000000..648baa4 --- /dev/null +++ b/src/main/kotlin/moe/nea/shot/ShotParser.kt @@ -0,0 +1,149 @@ +package moe.nea.shot + +import java.util.* + +class ShotParser { + + private data class Builder( + val indentation: Int, + val handler: Handler + ) + + private interface Handler { + fun handleLine(line: String): Handler? + } + + private inner class MethodInjectionBuilder : Handler { + val annotations = mutableListOf<ClassRef>() + val parameterAnnotations = mutableMapOf<Int, MutableList<ClassRef>>() + override fun handleLine(line: String): Handler? { + if (line.startsWith("annotate ")) { + annotations.add(ClassRef(line.substring("annotate ".length).trim())) + return null + } + if (line.startsWith("annotateParameter ")) { + val (idx, name) = line.substring("annotateParameter ".length).split(" ") + parameterAnnotations.getOrPut(idx.toInt()) { mutableListOf() }.add(ClassRef(name)) + return null + } + error("Unknown directive") + } + } + + private inner class FieldInjectionBuilder : Handler { + val annotations = mutableListOf<ClassRef>() + override fun handleLine(line: String): Handler? { + if (line.startsWith("annotate ")) { + annotations.add(ClassRef(line.substring("annotate ".length).trim())) + return null + } + error("Unknown directive") + } + } + + private inner class ClassInjectionBuilder : Handler { + val methodHandlers = mutableMapOf<MethodRef, MethodInjectionBuilder>() + val fieldHandlers = mutableMapOf<FieldRef, FieldInjectionBuilder>() + val annotations = mutableListOf<ClassRef>() + override fun handleLine(line: String): Handler? { + if (line.endsWith(":")) { + val data = line.substring(0, line.length - 1).trim() + if (data.contains("(")) { + return methodHandlers.getOrPut(parseMethod(data)) { MethodInjectionBuilder() } + } + if (" " !in data) { + return fieldHandlers.getOrPut(FieldRef(data)) { FieldInjectionBuilder() } + } + error("Unknown condition") + } + if (line.startsWith("annotate ")) { + annotations.add(ClassRef(line.substring("annotate ".length).trim())) + return null + } + error("Unknown line $line") + } + + private fun parseMethod(data: String): MethodRef { + require(data.endsWith(")")) + val name = data.substringBefore("(") + val parameterDescriptor = data.substringAfter("(").dropLast(1) + val parameters = parameterDescriptor.split(",").filter { it.isNotBlank() }.map { + mapTypeToDescriptor(false, it.trim()) + } + return MethodRef(name, parameters.joinToString("", "(", ")")) + } + } + + private fun mapTypeToDescriptor(allowVoid: Boolean, name: String): String { + require(" " !in name) + if (name.endsWith("[]")) { + return "[" + mapTypeToDescriptor(allowVoid, name.substring(0, name.length - 2)) + } + return when (name) { + "void" -> if (allowVoid) "V" else error("Void not allowed") + "boolean" -> "Z" + "byte" -> "B" + "int" -> "I" + "double" -> "D" + "long" -> "J" + "short" -> "S" + "float" -> "F" + "char" -> "C" + else -> "L" + name.replace(".", "/") + ";" + } + } + + private inner class Root : Handler { + val map = mutableMapOf<ClassRef, ClassInjectionBuilder>() + override fun handleLine(line: String): Handler? { + require(line.endsWith(":")) + val className = ClassRef(line.substring(0, line.length - 1).trim()) + return map.getOrPut(className) { ClassInjectionBuilder() } + } + } + + private val root = Root() + private val indentations = Stack<Builder>().also { + it.add(Builder(0, root)) + } + private var nextHandler: Handler? = null + + fun parse(lines: Iterable<String>): Map<ClassRef, ClassShots> { + lines.forEach { parseLine(it) } + return root.map.mapValues { + ClassShots( + it.value.annotations, + it.value.methodHandlers.mapValues { + MethodShots( + it.value.annotations, + it.value.parameterAnnotations + ) + }, + it.value.fieldHandlers.mapValues { + FieldShots( + it.value.annotations + ) + } + ) + } + } + + fun parseLine(line: String) { + val indentation = line.takeWhile { it == ' ' }.length + + if (indentation > indentations.peek().indentation) { + indentations.push(Builder(indentation, nextHandler ?: error("Illegal increase in indentation"))) + nextHandler = null + } + while (indentation < indentations.peek().indentation) { + indentations.pop() + } + + val nonCommentedLine = line.substringBefore('#') + if (nonCommentedLine.isBlank()) { + return + } + nextHandler = indentations.peek().handler.handleLine(line.trim()) + } + +}
\ No newline at end of file diff --git a/src/main/kotlin/moe/nea/shot/Shots.kt b/src/main/kotlin/moe/nea/shot/Shots.kt new file mode 100644 index 0000000..b0956b4 --- /dev/null +++ b/src/main/kotlin/moe/nea/shot/Shots.kt @@ -0,0 +1,55 @@ +package moe.nea.shot + +import org.objectweb.asm.ClassReader +import org.objectweb.asm.ClassWriter +import java.io.Serializable +import java.nio.file.Files +import java.nio.file.Path +import java.util.zip.ZipFile +import java.util.zip.ZipOutputStream +import kotlin.io.path.createDirectories +import kotlin.io.path.exists + +data class Shots( + val injections: Map<ClassRef, ClassShots>, +) : Serializable { + + fun processEntry(toInject: ClassShots, bytes: ByteArray): ByteArray { + val classReader = ClassReader(bytes) + val classWriter = ClassWriter(0) + val visitor = InjectionApplicator(toInject, classWriter) + classReader.accept(visitor, 0) + return classWriter.toByteArray() + } + + fun processZipFile(input: ZipFile, output: ZipOutputStream) { + for (entry in input.entries()) { + val classRef = ClassRef.fromPath(entry.name) + val toInject = classRef?.let(injections::get) + output.putNextEntry(entry) + if (toInject == null) { + input.getInputStream(entry).copyTo(output) + } else { + val bytes = input.getInputStream(entry).use { it.readBytes() } + val modifiedBytes = processEntry(toInject, bytes) + output.write(modifiedBytes) + } + } + } + + + /** + * Run on an extracted zip, (or alternatively, a jar file systems root path) + */ + fun process(sourceFiles: Path, targetFiles: Path) { + for ((classRef, toInject) in injections.entries) { + val classSource = sourceFiles.resolve(classRef.path) + val classTarget = targetFiles.resolve(classRef.path) + if (!classSource.exists()) { + continue + } + classTarget.parent.createDirectories() + Files.write(classTarget, processEntry(toInject, Files.readAllBytes(classSource))) + } + } +}
\ No newline at end of file |