fix(bytecode): fix import issues

archguard · Mar 29, 2022 · 3dae498 · 3dae498
1 parent 2bf2c04
commit 3dae498
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 8 deletions.
diff --git a/scan_bytecode/src/main/kotlin/org/archguard/scanner/bytecode/ByteCodeParser.kt b/scan_bytecode/src/main/kotlin/org/archguard/scanner/bytecode/ByteCodeParser.kt
@@ -11,7 +11,6 @@ import org.slf4j.LoggerFactory
 import java.io.File
 import java.io.FileInputStream
 import java.io.IOException
-import java.util.*
 import java.util.stream.Collectors
 
 class ByteCodeParser {
@@ -100,7 +99,7 @@ class ByteCodeParser {
 
     private fun createDataStruct(node: ClassNode, module: CodeModule): CodeDataStruct {
         val ds = CodeDataStruct()
-        val names = importCollector.splitPackageAndClassName(Type.getObjectType(node.name).className)
+        val names = importCollector.splitPackageAndClassName(classNameFromType(node.name, false))
         ds.Package = names.first
         ds.NodeName = names.second
 
@@ -115,9 +114,12 @@ class ByteCodeParser {
             ds.Functions += this.createFunction(it, node)
         }
 
-        ds.Extend = Type.getObjectType(node.superName).className
+        val superName = node.superName
+        ds.Extend = classNameFromType(superName, true)
+
         ds.Implements = node.interfaces?.map {
-            Type.getObjectType(it).className
+            val className = classNameFromType(it, true)
+            className
         }?.toTypedArray() ?: arrayOf()
 
         ds.Annotations = node.visibleAnnotations?.map {
@@ -232,7 +234,16 @@ class ByteCodeParser {
             }
         }
 
-        return Type.getObjectType(ownerName).className
+        return classNameFromType(ownerName!!, true)
+    }
+
+    private fun classNameFromType(superName: String, isImport: Boolean): String {
+        val className = Type.getObjectType(superName).className
+        if (isImport) {
+            importCollector.processClassName(className)
+        }
+
+        return className
     }
 
     private fun getParamsFromDesc(desc: String, parameters: MutableList<ParameterNode>): Array<CodeProperty> {
@@ -256,4 +267,4 @@ class ByteCodeParser {
             )
         }.toTypedArray()
     }
-}
+}
diff --git a/scan_bytecode/src/main/kotlin/org/archguard/scanner/bytecode/ImportCollector.kt b/scan_bytecode/src/main/kotlin/org/archguard/scanner/bytecode/ImportCollector.kt
@@ -1,7 +1,8 @@
 package org.archguard.scanner.bytecode
 
 class ImportCollector {
-    private val JAVA_LANG_PACKAGE = "java.lang"
+    private val JAVA_LANG_PACKAGE = "java."
+    private val KOTLIN_LANG_PACKAGE = "kotlin."
 
     private var mapSimpleNames: MutableMap<String, String> = mutableMapOf()
     private val setNotImportedNames: Set<String> = HashSet()
@@ -38,7 +39,7 @@ class ImportCollector {
     }
 
     fun processClassName(className: String) {
-        if (className.startsWith(JAVA_LANG_PACKAGE)) {
+        if (className.startsWith(JAVA_LANG_PACKAGE) || className.startsWith(KOTLIN_LANG_PACKAGE)) {
             return
         }
 

diff --git a/scan_bytecode/src/test/kotlin/org/archguard/scanner/bytecode/ByteCodeParserTest.kt b/scan_bytecode/src/test/kotlin/org/archguard/scanner/bytecode/ByteCodeParserTest.kt
@@ -119,6 +119,17 @@ internal class ByteCodeParserTest {
         val path = Paths.get(resource.toURI()).toFile()
         val ds = ByteCodeParser().parseClassFile(path)
 
+
+        assertEquals(5, ds.Imports.size)
+
+        assertEquals("org.springframework.web.client.RestTemplate," +
+                "com.thoughtworks.archguard.report.infrastructure.QualityGateClientImpl," +
+                "com.thoughtworks.archguard.report.domain.qualitygate.CouplingQualityGate," +
+                "com.thoughtworks.archguard.report.domain.qualitygate.QualityGateClient," +
+                "org.springframework.stereotype.Component",
+
+            ds.Imports.joinToString(",") { it.Source })
+
         var hasRestTemplateCall = false
         ds.Functions.forEach { function ->
             function.FunctionCalls.forEach {