add whitelist and log switch

jokermonn · Sep 14, 2021 · c021ec8 · c021ec8
1 parent 923df33
commit c021ec8
Show file tree

Hide file tree

Showing 5 changed files with 49 additions and 15 deletions.
diff --git a/lib/src/main/kotlin/com/joker/codelinker/ClassLinker.kt b/lib/src/main/kotlin/com/joker/codelinker/ClassLinker.kt
@@ -5,6 +5,8 @@ import com.joker.codelinker.model.ClzMethod
 import com.joker.codelinker.model.ClzMethodInfo
 import com.joker.codelinker.model.METHODS_ALL
 import com.joker.codelinker.model.METHODS_CLINIT
+import com.joker.codelinker.util.logEnabled
+import com.joker.codelinker.util.printLog
 import com.joker.codelinker.util.toMethod
 import org.objectweb.asm.ClassReader
 import org.objectweb.asm.Opcodes
@@ -14,7 +16,9 @@ class ClassLinker constructor(builder: Builder) {
     private val linkerConfig = LinkerConfig()
 
     init {
-        linkerConfig.filterBlock = builder.filterBlock
+        linkerConfig.blockList = builder.blockList
+        linkerConfig.whiteList = builder.whiteList
+        logEnabled = builder.logEnabled
         builder.list.forEach { linkerConfig.addMethod(it) }
     }
 
@@ -25,7 +29,7 @@ class ClassLinker constructor(builder: Builder) {
             val clzMethod = linkerConfig.clzMethods.poll()
             val clzName = clzMethod.className
             clzs.add(clzName)
-            println("$clzName ${clzMethod.method.methodName}${clzMethod.method.methodDesc} is reading...")
+            printLog("$clzName ${clzMethod.method.methodName}${clzMethod.method.methodDesc} is reading...")
             val reader = ClassReader(clzName)
             reader.accept(clzNode, 0)
             clzNode.start(clzMethod.method, linkerConfig)
@@ -39,7 +43,9 @@ class ClassLinker constructor(builder: Builder) {
 
     class Builder internal constructor() {
         internal val list = arrayListOf<ClzMethod>()
-        internal var filterBlock: (String) -> Boolean = { false }
+        internal var logEnabled = false
+        internal var blockList: ((String) -> Boolean)? = null
+        internal var whiteList: ((String) -> Boolean)? = null
 
         fun methods(vararg clzMethods: ClzMethodInfo): Builder {
             val onlyClzSelf = clzMethods.filter { it.methodInfo == null }
@@ -56,8 +62,24 @@ class ClassLinker constructor(builder: Builder) {
             return this
         }
 
-        fun filter(block: (String) -> Boolean): Builder {
-            filterBlock = block
+        fun enableLog(enabled: Boolean): Builder {
+            this.logEnabled = enabled
+            return this
+        }
+
+        fun blockList(block: (String) -> Boolean): Builder {
+            if (whiteList != null) {
+                throw IllegalArgumentException("The blocklist does not take effect when the whitelist is used")
+            }
+            blockList = block
+            return this
+        }
+
+        fun whiteList(block: (String) -> Boolean): Builder {
+            if (blockList != null) {
+                throw IllegalArgumentException("The blocklist does not take effect when the whitelist is used")
+            }
+            whiteList = block
             return this
         }
 

diff --git a/lib/src/main/kotlin/com/joker/codelinker/ClzNode.kt b/lib/src/main/kotlin/com/joker/codelinker/ClzNode.kt
@@ -8,6 +8,7 @@ import com.joker.codelinker.util.TraceSignatureVisitor2
 import com.joker.codelinker.util.containsAllMethods
 import com.joker.codelinker.util.isClinitMethod
 import com.joker.codelinker.util.isStatic
+import com.joker.codelinker.util.printLog
 import org.objectweb.asm.signature.SignatureReader
 import org.objectweb.asm.tree.ClassNode
 import org.objectweb.asm.tree.FieldInsnNode
@@ -53,13 +54,13 @@ internal class ClzNode(api: Int) : ClassNode(api) {
     private fun analyseMethods(methods: List<MethodNode>, linkerConfig: LinkerConfig) {
         val clzName = name
         methods.forEach {
-            println("================method start: $clzName ${it.name}${it.desc} ================")
+            printLog("================method start: $clzName ${it.name}${it.desc} ================")
             analyseInstructions(it.instructions, linkerConfig)
             analyseLocalVariables(it.localVariables, linkerConfig)
             analyseMethodDescriptor(it.desc, linkerConfig)
             analyseExceptions(it.exceptions, linkerConfig)
             analyseTryCatchBlocks(it.tryCatchBlocks, linkerConfig)
-            println("================method end: $clzName ${it.name}${it.desc} ================")
+            printLog("================method end: $clzName ${it.name}${it.desc} ================")
         }
     }
 

diff --git a/lib/src/main/kotlin/com/joker/codelinker/LinkerConfig.kt b/lib/src/main/kotlin/com/joker/codelinker/LinkerConfig.kt
@@ -5,10 +5,12 @@ import com.joker.codelinker.model.METHODS_ALL
 import com.joker.codelinker.model.METHODS_CLINIT
 import com.joker.codelinker.model.Method
 import com.joker.codelinker.model.NoSuchMethod
+import com.joker.codelinker.util.printLog
 import java.util.LinkedList
 
 class LinkerConfig {
-    internal var filterBlock: (String) -> Boolean = { false }
+    internal var blockList: ((String) -> Boolean)? = null
+    internal var whiteList: ((String) -> Boolean)? = null
 
     internal val noSuchMethods: MutableList<NoSuchMethod> = arrayListOf()
 
@@ -30,7 +32,7 @@ class LinkerConfig {
 
     internal fun addMethod(clzMethod: ClzMethod) {
         val clzName = clzMethod.className
-        if (filterBlock.invoke(clzName)) {
+        if (blockList?.invoke(clzName) == true || whiteList?.invoke(clzName) == false) {
             return
         }
         if (recordSet.contains(ClzMethod(clzName, METHODS_ALL))) {
@@ -43,7 +45,7 @@ class LinkerConfig {
             noSuchMethods.removeIf { it.clzName == clzMethod.startClzName }
         }
         recordSet.add(clzMethod)
-        println("clzMethods add: $clzMethod")
+        printLog("clzMethods add: $clzMethod")
         clzMethods.add(clzMethod)
     }
 }
diff --git a/lib/src/main/kotlin/com/joker/codelinker/util/Log.kt b/lib/src/main/kotlin/com/joker/codelinker/util/Log.kt
@@ -0,0 +1,9 @@
+package com.joker.codelinker.util
+
+var logEnabled = false
+
+fun printLog(string: String?) {
+    if (logEnabled) {
+        println(string)
+    }
+}
diff --git a/src/main/kotlin/com/joker/Test.kt b/src/main/kotlin/com/joker/Test.kt
@@ -25,14 +25,14 @@ fun main() {
                 )
             )
         )
-        .filter { it.startsWith("java/") }
+        .blockList { it.startsWith("java/") }
         .build()
         .link()
 
-//    println("wantSet: $wantSet")
-//    println("actual class loader: $actualSet")
-//    println("want exist but actual not exist: ${wantSet.subtract(actualSet)}")
-//    println("actual exist but want not exist: ${actualSet.subtract(wantSet).filter { !it.contains("$") }}")
+//    printLog("wantSet: $wantSet")
+//    printLog("actual class loader: $actualSet")
+//    printLog("want exist but actual not exist: ${wantSet.subtract(actualSet)}")
+//    printLog("actual exist but want not exist: ${actualSet.subtract(wantSet).filter { !it.contains("$") }}")
 }
 
 private fun logClassLoader(block: () -> Unit): Set<String> {