RedefineUtils.kt

package com.greenapple.glacia.utils

import com.ea.agentloader.AgentLoader
import javassist.*
import java.io.Serializable
import java.lang.instrument.ClassDefinition
import java.lang.instrument.Instrumentation
import java.util.concurrent.atomic.AtomicLong
import kotlin.reflect.KCallable
import kotlin.reflect.KClass
import kotlin.reflect.KFunction

private typealias TMethodNoReturn<O> = TMethod<O, Unit>
private typealias TMethod<O, R> = O.(@ParameterName("args") Array<out Any?>) -> R

object RedefineUtils {
    fun interface IMethod<O, R>: (O, Array<out Any?>) -> R, Serializable {
        override fun invoke(receiver: O, vararg args: Any?): R
    }
    private class RedefineAgent {
        companion object {
            @JvmStatic private lateinit var instrumentationInternal: Instrumentation
            @JvmStatic val instrumentation by lazy {
                ClassLoader.getSystemClassLoader().loadClass(RedefineAgent::class.java.name).getDeclaredField(::instrumentationInternal.name).apply {isAccessible = true}[null] as Instrumentation
            }
            @JvmStatic fun agentmain(agentArgs: String?, instr: Instrumentation) {
                instrumentationInternal = instr
            }
        }
    }

    @JvmStatic val instrumentation by lazy {
        AgentLoader.loadAgentClass(RedefineAgent::class.java.name, null)
        RedefineAgent.instrumentation
    }

    class FallBackException : Exception()
    fun fallback(): Nothing = throw FallBackException()
}

private val nextIndex = AtomicLong(0)

private inline fun <O : Any> KClass<O>.editClassDef(block: CtClass.(@ParameterName("instrumentation") Instrumentation) -> Unit) = try {
    ClassPool.getDefault().apply {appendClassPath(LoaderClassPath(Thread.currentThread().contextClassLoader))}[qualifiedName]?.apply {
        defrost()
        block(this, RedefineUtils.instrumentation)
        RedefineUtils.instrumentation.redefineClasses(ClassDefinition(this@editClassDef.java, this.toBytecode()))
    } ?: System.err.println("Class $qualifiedName not found")
} catch (ex: Exception) {ex.printStackTrace()}

private inline fun <O : Any> KClass<O>.editMethodDef(methodName: String, block: CtMethod.(@ParameterName("instrumentation") Instrumentation) -> Unit) = editClassDef { instrumentation->
    block(getDeclaredMethod(methodName), instrumentation)
}

private fun <O : Any> CtMethod.newStaticMethodCall(function: TMethod<O, Any>, shouldReturn: Boolean = false) = StringBuffer().apply {
    val extrasClass = "${this@newStaticMethodCall.declaringClass.name}\$extras_${nextIndex.getAndIncrement()}"
    ClassPool.getDefault().makeClass(extrasClass).apply {
        val iMethod = RedefineUtils.IMethod(function)
        addField(CtField.make("public static ${RedefineUtils.IMethod::class.java.name} __callable__ = null;\n", this))
        toClass().getDeclaredField("__callable__").set(null, iMethod)
    }
    append("{\nObject __returnValue__ = $extrasClass.__callable__.invoke($0, \$args);\n")
    if (shouldReturn && returnType.name != "void")
        append("return ((${(returnType as? CtPrimitiveType)?.wrapperName ?: returnType.name})(__returnValue__))${if (returnType.isPrimitive) ".${returnType.name}Value()" else ""};\n")
    else if (shouldReturn) append("return;\n")
    append("}")
}.toString()

private fun String.wrapTryOrFallback() = "try $this catch (${RedefineUtils.FallBackException::class.java.name} __fallbackException__) {}"

/**
 * Appends a function at end of a method
 */
fun <O : Any> KClass<O>.addMethodAfter(method: KCallable<*>, function: TMethodNoReturn<O>) = addMethodAfter(method.name, function)
fun <O : Any> KClass<O>.addMethodAfter(methodName: String, function: TMethodNoReturn<O>) = editMethodDef(methodName) {
    insertAfter(newStaticMethodCall(function))
}

/**
 * Appends a function at the start of a method
 */
fun <O : Any> KClass<O>.addMethodBefore(method: KCallable<*>, function: TMethodNoReturn<O>) = addMethodBefore(method.name, function)
fun <O : Any> KClass<O>.addMethodBefore(methodName: String, function: TMethodNoReturn<O>) = editMethodDef(methodName) {
    insertBefore(newStaticMethodCall(function))
}

/**
 * Replaces a method
 * Must call RedefineUtils.fallback() to revert to default behaviour
 */
fun <O : Any, R : Any> KClass<O>.replaceMethodOrFallback(method: KCallable<*>, function: TMethod<O, R>) = replaceMethodOrFallback(method.name, function)
fun <O : Any> KClass<O>.replaceMethodOrFallback(methodName: String, function: TMethod<O, Any>) = editMethodDef(methodName) {
    insertBefore(newStaticMethodCall(function, true).wrapTryOrFallback())
}

/**
 * Replaces a method
 */
fun <O : Any, R : Any> KClass<O>.replaceMethod(method: KFunction<R>, function: TMethod<O, R>) = replaceMethod(method.name, function)
fun <O : Any> KClass<O>.replaceMethod(methodName: String, function: TMethod<O, Any>) = editMethodDef(methodName) {
    setBody(newStaticMethodCall(function, true))
}