kotlin2cpg: add visibility modifiers to fns (#2799)

joernio · Jun 2, 2023 · cfafc11 · cfafc11
1 parent dbd35f5
commit cfafc11
Show file tree

Hide file tree

Showing 5 changed files with 50 additions and 3 deletions.
diff --git a/joern-cli/frontends/kotlin2cpg/src/main/scala/io/joern/kotlin2cpg/ast/KtPsiToAst.scala b/joern-cli/frontends/kotlin2cpg/src/main/scala/io/joern/kotlin2cpg/ast/KtPsiToAst.scala
@@ -19,6 +19,8 @@ import io.joern.x2cpg.utils.NodeBuilders.{
   newLocalNode,
   newMethodReturnNode
 }
+import org.jetbrains.kotlin.descriptors.DescriptorVisibilities
+import org.jetbrains.kotlin.descriptors.DescriptorVisibility
 
 import java.util.UUID.randomUUID
 import org.jetbrains.kotlin.psi._
@@ -401,18 +403,41 @@ trait KtPsiToAst {
     val typeFullName      = registerType(typeInfoProvider.returnType(ktFn, explicitTypeName))
     val _methodReturnNode = newMethodReturnNode(typeFullName, None, line(ktFn), column(ktFn))
 
+    val visibility = typeInfoProvider.visibility(ktFn)
+    val visibilityModifierType =
+      modifierTypeForVisibility(visibility.getOrElse(DescriptorVisibilities.UNKNOWN))
+    val visibilityModifier = modifierNode(visibilityModifierType)
+
     val modifierNodes =
       if (withVirtualModifier) Seq(modifierNode(ModifierTypes.VIRTUAL))
       else Seq()
 
     val annotationEntries = ktFn.getAnnotationEntries.asScala.map(astForAnnotationEntry).toSeq
     Seq(
-      methodAst(_methodNode, thisParameterAsts ++ methodParametersAsts, bodyAst, _methodReturnNode, modifierNodes)
+      methodAst(
+        _methodNode,
+        thisParameterAsts ++ methodParametersAsts,
+        bodyAst,
+        _methodReturnNode,
+        List(visibilityModifier) ++ modifierNodes
+      )
         .withChildren(otherBodyAsts)
         .withChildren(annotationEntries)
     )
   }
 
+  private def modifierTypeForVisibility(visibility: DescriptorVisibility): String = {
+    if (visibility.toString == DescriptorVisibilities.PUBLIC.toString)
+      ModifierTypes.PUBLIC
+    else if (visibility.toString == DescriptorVisibilities.PRIVATE.toString)
+      ModifierTypes.PRIVATE
+    else if (visibility.toString == DescriptorVisibilities.PROTECTED.toString)
+      ModifierTypes.PROTECTED
+    else if (visibility.toString == DescriptorVisibilities.INTERNAL.toString)
+      ModifierTypes.INTERNAL
+    else "UNKNOWN"
+  }
+
   def astsForBlock(
     expr: KtBlockExpression,
     argIdx: Option[Int],

diff --git a/...ontends/kotlin2cpg/src/main/scala/io/joern/kotlin2cpg/types/DefaultTypeInfoProvider.scala b/...ontends/kotlin2cpg/src/main/scala/io/joern/kotlin2cpg/types/DefaultTypeInfoProvider.scala
@@ -11,7 +11,13 @@ import org.jetbrains.kotlin.cli.jvm.compiler.{
   NoScopeRecordCliBindingTrace
 }
 import org.jetbrains.kotlin.com.intellij.util.keyFMap.KeyFMap
-import org.jetbrains.kotlin.descriptors.{DeclarationDescriptor, FunctionDescriptor, ValueDescriptor}
+import org.jetbrains.kotlin.descriptors.{
+  DeclarationDescriptor,
+  DescriptorVisibilities,
+  DescriptorVisibility,
+  FunctionDescriptor,
+  ValueDescriptor
+}
 import org.jetbrains.kotlin.descriptors.impl.{
   ClassConstructorDescriptorImpl,
   EnumEntrySyntheticClassDescriptor,
@@ -104,6 +110,12 @@ class DefaultTypeInfoProvider(environment: KotlinCoreEnvironment) extends TypeIn
       .getOrElse(defaultValue)
   }
 
+  def visibility(fn: KtNamedFunction): Option[DescriptorVisibility] = {
+    val mapForEntity = bindingsForEntity(bindingContext, fn)
+    Option(mapForEntity.get(BindingContext.FUNCTION.getKey))
+      .map(_.getVisibility)
+  }
+
   def containingTypeDeclFullName(ktFn: KtNamedFunction, defaultValue: String): String = {
     val mapForEntity = bindingsForEntity(bindingContext, ktFn)
     Option(mapForEntity.get(BindingContext.FUNCTION.getKey))

diff --git a/...-cli/frontends/kotlin2cpg/src/main/scala/io/joern/kotlin2cpg/types/TypeInfoProvider.scala b/...-cli/frontends/kotlin2cpg/src/main/scala/io/joern/kotlin2cpg/types/TypeInfoProvider.scala
@@ -1,6 +1,7 @@
 package io.joern.kotlin2cpg.types
 
 import io.shiftleft.passes.KeyPool
+import org.jetbrains.kotlin.descriptors.DescriptorVisibility
 import org.jetbrains.kotlin.psi.{
   KtAnnotationEntry,
   KtBinaryExpression,
@@ -32,6 +33,8 @@ trait TypeInfoProvider {
 
   def isStaticMethodCall(expr: KtQualifiedExpression): Boolean
 
+  def visibility(fn: KtNamedFunction): Option[DescriptorVisibility]
+
   def returnType(elem: KtNamedFunction, defaultValue: String): String
 
   def containingDeclFullName(expr: KtCallExpression): Option[String]

diff --git a/.../frontends/kotlin2cpg/src/test/scala/io/joern/kotlin2cpg/querying/MethodReturnTests.scala b/.../frontends/kotlin2cpg/src/test/scala/io/joern/kotlin2cpg/querying/MethodReturnTests.scala
@@ -17,7 +17,6 @@ class MethodReturnTests extends KotlinCode2CpgFixture(withOssDataflow = false) {
       val List(x) = cpg.method.name("foo").methodReturn.l
       x.code shouldBe "RET"
       x.evaluationStrategy shouldBe EvaluationStrategies.BY_VALUE
-      x.order shouldBe 4
       x.lineNumber shouldBe Some(2)
       x.columnNumber shouldBe Some(4)
     }

diff --git a/joern-cli/frontends/kotlin2cpg/src/test/scala/io/joern/kotlin2cpg/querying/MethodTests.scala b/joern-cli/frontends/kotlin2cpg/src/test/scala/io/joern/kotlin2cpg/querying/MethodTests.scala
@@ -40,6 +40,14 @@ class MethodTests extends KotlinCode2CpgFixture(withOssDataflow = false) {
       y.filename.endsWith(".kt") shouldBe true
     }
 
+    "should contain MODIFIER nodes attached to the METHOD nodes" in {
+      val List(mod1) = cpg.method.nameExact("double").modifier.l
+      mod1.modifierType shouldBe "PUBLIC"
+
+      val List(mod2) = cpg.method.nameExact("main").modifier.l
+      mod2.modifierType shouldBe "PUBLIC"
+    }
+
     "should allow traversing to parameters" in {
       cpg.method.name("double").isExternal(false).parameter.name.toSet shouldBe Set("x")
       cpg.method.name("main").isExternal(false).parameter.name.toSet shouldBe Set("args")