aboutsummaryrefslogtreecommitdiff
path: root/src/main/kotlin/moe/nea/shot
diff options
context:
space:
mode:
Diffstat (limited to 'src/main/kotlin/moe/nea/shot')
-rw-r--r--src/main/kotlin/moe/nea/shot/ClassRef.kt17
-rw-r--r--src/main/kotlin/moe/nea/shot/ClassShots.kt9
-rw-r--r--src/main/kotlin/moe/nea/shot/FieldRef.kt5
-rw-r--r--src/main/kotlin/moe/nea/shot/FieldShots.kt5
-rw-r--r--src/main/kotlin/moe/nea/shot/InjectionApplicator.kt69
-rw-r--r--src/main/kotlin/moe/nea/shot/MethodRef.kt5
-rw-r--r--src/main/kotlin/moe/nea/shot/MethodShots.kt8
-rw-r--r--src/main/kotlin/moe/nea/shot/ShotParser.kt149
-rw-r--r--src/main/kotlin/moe/nea/shot/Shots.kt55
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