Merge pull request #41 from SpringKill-team/feature/enhance-callgraph

Feature/enhance callgraph

Author: springkill

.run/Run Plugin.run.xml

@@ -14,7 +14,7 @@
           <option value="runIde" />
         </list>
       </option>
-      <option name="vmOptions" value="-Xmx10240m" />
+      <option name="vmOptions" value="-Xmx15240m" />
     </ExternalSystemSettings>
     <ExternalSystemDebugServerProcess>true</ExternalSystemDebugServerProcess>
     <ExternalSystemReattachDebugProcess>true</ExternalSystemReattachDebugProcess>

src/main/kotlin/org/skgroup/securityinspector/actions/BuildCallGraphAction.kt

@@ -30,8 +30,6 @@ class BuildCallGraphAction : AnAction() {
     }
 
     override fun update(e: AnActionEvent) {
-//        val project = e.project
-//        val psiElement = e.getData(CommonDataKeys.PSI_ELEMENT)
 
         CoroutineScope(Dispatchers.Default).launch {
             val isMethod = readAction {

src/main/kotlin/org/skgroup/securityinspector/actions/SearchAsSinkAction.kt

@@ -30,8 +30,6 @@ class SearchAsSinkAction : AnAction() {
     }
 
     override fun update(e: AnActionEvent) {
-//        val project = e.project
-//        val psiElement = e.getData(CommonDataKeys.PSI_ELEMENT)
 
         CoroutineScope(Dispatchers.Default).launch {
 

src/main/kotlin/org/skgroup/securityinspector/analysis/graphs/callgraph/CallGraph.kt

@@ -17,17 +17,6 @@ data class CallGraph(
     val edgeTypes: MutableMap<Pair<MethodNode, MethodNode>, EdgeType> = mutableMapOf(),
     val nodeTypes: MutableMap<MethodNode, NodeType> = mutableMapOf()
 ) {
-    fun addEdge(from: MethodNode, to: MethodNode, type: EdgeType) {
-        edges.getOrPut(from) { mutableSetOf() }.add(to)
-        edgeTypes[Pair(from, to)] = type
-        nodes.add(from)
-        nodes.add(to)
-    }
-
-    fun addSpecialNode(node: MethodNode, type: NodeType) {
-        nodeTypes[node] = type
-    }
-
     fun merge(other: CallGraph) {
         // 合并节点
         this.nodes.addAll(other.nodes)
@@ -42,4 +31,8 @@ data class CallGraph(
         this.nodeTypes.putAll(other.nodeTypes)
     }
 
+    fun findSinks() = nodes.filter { edges[it].isNullOrEmpty() }
+
+    fun findRoots() = nodes.filter { it !in edges.values.flatten().toSet() }
+
 }

src/main/kotlin/org/skgroup/securityinspector/analysis/graphs/callgraph/CallGraphBuilder.kt

@@ -4,14 +4,14 @@ import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.project.Project
 import com.intellij.psi.*
 import com.intellij.psi.search.GlobalSearchScope
-import com.intellij.psi.search.searches.ClassInheritorsSearch
 import com.intellij.psi.search.searches.ReferencesSearch
 import com.intellij.psi.util.PsiTreeUtil
 import org.skgroup.securityinspector.analysis.ast.nodes.MethodNode
 import org.skgroup.securityinspector.analysis.di.DIProcessor
 import org.skgroup.securityinspector.enums.RefMode
 import org.skgroup.securityinspector.utils.GraphUtils
-import java.util.ArrayDeque
+import org.skgroup.securityinspector.utils.PsiUtil.findConcreteImplementations
+import java.util.*
 
 /**
  * Call graph builder 是一个用于构建调用图的类
@@ -32,8 +32,6 @@ class CallGraphBuilder : JavaRecursiveElementVisitor() {
     // 用于记录当前访问的方法调用栈
     private val currentMethodStack = ArrayDeque<MethodNode>()
 
-//    private val diProcessor = DIProcessor(callGraph)
-
     /**
      * Visit method 方法用于访问一个Java方法并将其加入调用图
      * 用栈处理递归方法调用
@@ -46,7 +44,6 @@ class CallGraphBuilder : JavaRecursiveElementVisitor() {
         var calleeStackNode = GraphUtils.getMethodNode(method)
 
         // 查找对该方法的所有引用，建立反向调用关系
-        // TODO 这里的逻辑我自己也混乱了，可能会有问题，后面再看
         val projectScope = GlobalSearchScope.projectScope(method.project)
         ReferencesSearch.search(method, projectScope).forEach { reference ->
             val callerMethod = PsiTreeUtil.getParentOfType(reference.element, PsiMethod::class.java)
@@ -57,6 +54,14 @@ class CallGraphBuilder : JavaRecursiveElementVisitor() {
             callGraph.nodes.add(calleeMethodNode)
 
             val callerMethodNode = GraphUtils.getMethodNode(callerMethod, reference)
+
+            if (method.hasModifierProperty(PsiModifier.ABSTRACT) || method.containingClass?.isInterface == true) {
+                findConcreteImplementations(method).forEach { implMethod ->
+                    val implNode = GraphUtils.getMethodNode(implMethod)
+                    callGraph.edges.getOrPut(callerMethodNode) { mutableSetOf() }.add(implNode)
+                }
+            }
+
             // 建立反向调用关系
             callGraph.nodes.add(callerMethodNode)
             callGraph.edges.getOrPut(callerMethodNode) { mutableSetOf() }.add(calleeMethodNode)
@@ -79,7 +84,6 @@ class CallGraphBuilder : JavaRecursiveElementVisitor() {
      */
     override fun visitMethodCallExpression(expression: PsiMethodCallExpression) {
         if (currentMethodStack.isEmpty()) return
-        expression.methodExpression.qualifierExpression?.accept(this)
 
         val caller = currentMethodStack.peek()
         // 解析被调用的方法
@@ -106,16 +110,6 @@ class CallGraphBuilder : JavaRecursiveElementVisitor() {
                 .add(calleeMethodNode)
             handleIoCContainerCall(expression, caller, calleeMethodNode)
         } ?: run {
-//            if (clazz != null) {
-//                if (clazz.isInterface) {
-//                    println("$clazz is an interface")
-//                    println("and the method call expression $expression can not resolve reference.")
-//                }
-//                if (clazz.hasModifierProperty(PsiModifier.ABSTRACT)) {
-//                    println("$clazz is an abstract class")
-//                    println("and the method call expression $expression can not resolve reference.")
-//                }
-//            }
             if (PsiTreeUtil.getParentOfType(expression, PsiLambdaExpression::class.java) != null) {
                 println("in lambda method call expression $expression can not resolve reference.")
             } else {
@@ -139,7 +133,6 @@ class CallGraphBuilder : JavaRecursiveElementVisitor() {
      */
     override fun visitNewExpression(expression: PsiNewExpression) {
         if (currentMethodStack.isEmpty()) return
-//        expression.methodNewExpression.qualifierExpression?.accept(this)
 
         val callerMethodNode = currentMethodStack.peek()
         // 解析构造方法
@@ -160,16 +153,6 @@ class CallGraphBuilder : JavaRecursiveElementVisitor() {
                 .add(calleeMethodNode)
             println("expression $expression resolve reference success.")
         } ?: run {
-//            if (clazz != null) {
-//                if (clazz.isInterface) {
-//                    println("$clazz is an interface")
-//                    println("and the method call expression $expression can not resolve reference.")
-//                }
-//                if (clazz.hasModifierProperty(PsiModifier.ABSTRACT)) {
-//                    println("$clazz is an abstract class")
-//                    println("and the method call expression $expression can not resolve referencee.")
-//                }
-//            }
             if (PsiTreeUtil.getParentOfType(expression, PsiLambdaExpression::class.java) != null) {
                 println("in lambda method call expression $expression can not resolve reference.")
             } else {
@@ -329,26 +312,4 @@ class CallGraphBuilder : JavaRecursiveElementVisitor() {
         return callGraph
     }
 
-    private fun findConcreteImplementations(abstractMethod: PsiMethod): List<PsiMethod> {
-        val implementations = mutableListOf<PsiMethod>()
-        val containingClass = abstractMethod.containingClass ?: return emptyList()
-
-        // 查找所有子类或实现类
-        val inheritors = if (containingClass.isInterface) {
-            ClassInheritorsSearch.search(containingClass, abstractMethod.useScope, true).toList()
-        } else {
-            ClassInheritorsSearch.search(containingClass).toList()
-        }
-
-        inheritors.forEach { implClass ->
-            implClass.findMethodBySignature(abstractMethod, true)?.let {
-                if (!it.hasModifierProperty(PsiModifier.ABSTRACT)) {
-                    implementations.add(it)
-                }
-            }
-        }
-
-        return implementations
-    }
-
 }

src/main/kotlin/org/skgroup/securityinspector/analysis/graphs/processor/DIProcessor.kt

@@ -104,7 +104,7 @@ class DIProcessor(
 
                         // 将 cls -> field.type 关系记录到调用图
                         // 为了和现有的 method-based callGraph 统一，我们可以构造一个伪“类级方法节点”
-                        val callerNode = GraphUtils.getMethodNode(cls)   // 伪方法节点：类本身
+                        val callerNode = GraphUtils.getClassNode(cls)   // 伪方法节点：类本身
                         val calleeNode = findOrCreateMethodNodeForType(field.type)
 
                         // 入图
@@ -164,8 +164,8 @@ class DIProcessor(
 
                 if (implCls.isInheritor(serviceClass, true)) {
                     // 建立关联 service -> impl
-                    val serviceNode = GraphUtils.getMethodNode(serviceClass)
-                    val implNode = GraphUtils.getMethodNode(implCls)
+                    val serviceNode = GraphUtils.getClassNode(serviceClass)
+                    val implNode = GraphUtils.getClassNode(implCls)
 
                     callGraph.nodes.add(serviceNode)
                     callGraph.nodes.add(implNode)
@@ -214,7 +214,7 @@ class DIProcessor(
         val psiClass = PsiUtil.resolveClassInType(type)
         return if (psiClass != null) {
             // 直接用 psiClass 构造 MethodNode
-            GraphUtils.getMethodNode(psiClass)
+            GraphUtils.getClassNode(psiClass)
         } else {
             // 找不到对应类，做一个兜底处理
             MethodNode(type.canonicalText ?: "UnknownType",type.canonicalText ?: "UnknownType", "DI", emptyList(), emptyList(),RefMode.CALL)

src/main/kotlin/org/skgroup/securityinspector/enums/AnalysisScope.kt

@@ -10,9 +10,10 @@ enum class AnalysisScope(private val displayName: String) {
 
     ENTIRE_PROJECT("Entire Project"),
     MODULE("Selected Module");
-//    OPEN_FILE("Open Files");
+    // OPEN_FILE("Open Files");
 
     override fun toString(): String {
         return displayName
     }
+
 }

src/main/kotlin/org/skgroup/securityinspector/ui/component/CallGraphUIComponents.kt

@@ -34,8 +34,7 @@ class CallGraphUIComponents(val project: Project) {
     val rootListModel = DefaultListModel<MethodNode>()
     val rootList = JBList(rootListModel)
 
-    //    val infoArea = JTextArea().apply { isEditable = false }
-    val searchComboBox = ComboBox<AnalysisScope>(AnalysisScope.values())
+    val searchComboBox = ComboBox(AnalysisScope.values())
 
     val searchResultRootNode = DefaultMutableTreeNode("Search Results")
     val searchResultTreeModel = DefaultTreeModel(searchResultRootNode)
@@ -223,8 +222,6 @@ class CallGraphUIComponents(val project: Project) {
         // 初始化列表渲染器
         val methodRenderer = MethodListRenderer.createMethodRenderer()
         rootList.cellRenderer = methodRenderer
-//        sinkList.cellRenderer = methodRenderer
-//        searchResultList.cellRenderer = methodRenderer
 
         // 创建各个功能面板
         val rootPanel = createTitledPanel("Root Methods", JBScrollPane(rootList))

src/main/kotlin/org/skgroup/securityinspector/ui/service/CallGraphGenerator.kt

@@ -98,7 +98,10 @@ object CallGraphGenerator {
                     ApplicationManager.getApplication().runReadAction {
                         if (psiFile is PsiJavaFile) {
                             if (processedFiles.add(psiFile)) {
-                                psiFile.accept(builder)
+                                PsiTreeUtil.findChildrenOfType(psiFile, PsiMethod::class.java).forEach{
+                                    it.accept(builder)
+                                }
+//                                psiFile.accept(builder)
                             }
                         }
                     }