rubysrc2cpg: Getting code for nodes, control structure wrapper usage and fixes

joern-cli/frontends/rubysrc2cpg/src/main/scala/io/joern/rubysrc2cpg/astcreation/AstCreator.scala

@@ -15,12 +15,12 @@ import io.shiftleft.codepropertygraph.generated.nodes._
 import org.antlr.v4.runtime.tree.TerminalNode
 import org.antlr.v4.runtime.{CharStreams, CommonTokenStream, Token}
 import org.slf4j.LoggerFactory
-import overflowdb.{BatchedUpdate, Node}
+import overflowdb.{BatchedUpdate}
 
 import java.util
-import scala.collection.immutable.Seq
 import scala.collection.mutable
 import scala.collection.mutable.ListBuffer
+import scala.io.Source
 import scala.jdk.CollectionConverters._
 
 class AstCreator(filename: String, global: Global)
@@ -54,6 +54,8 @@ class AstCreator(filename: String, global: Global)
 
   private val methodAliases = mutable.HashMap[String, String]()
 
+  private val fileContent = mutable.HashMap[Int, String]()
+
   class ScopeIdentifiers {
     val varToIdentiferMap             = mutable.HashMap[String, NewIdentifier]()
     var parentScope: ScopeIdentifiers = null
@@ -83,18 +85,28 @@ class AstCreator(filename: String, global: Global)
     scopeStack.top.varToIdentiferMap.contains(name)
   }
 
+  private def getCodeForTerminalNode(node: TerminalNode): String = {
+    fileContent.getOrElse(node.getSymbol.getLine, "")
+  }
   private def createIdentiferWithScope(
     node: TerminalNode,
     name: String,
-    code: String,
     typeFullName: String,
     dynamicTypeHints: Seq[String]
   ): NewIdentifier = {
-    val newNode = identifierNode(node, name, code, typeFullName, dynamicTypeHints)
+    val newNode = identifierNode(node, name, getCodeForTerminalNode(node), typeFullName, dynamicTypeHints)
     setIdentiferInScope(newNode)
     newNode
   }
 
+  private def readFile(): Unit = {
+    var lineNum = 1
+    for (line <- Source.fromFile(filename).getLines()) {
+      val code = line.stripLeading().stripTrailing()
+      fileContent.addOne(lineNum, code)
+      lineNum += 1
+    }
+  }
   private def getActualMethodName(name: String): String = {
     methodAliases.getOrElse(name, name)
   }
@@ -105,6 +117,8 @@ class AstCreator(filename: String, global: Global)
     val parser      = new RubyParser(tokenStream)
     val programCtx  = parser.program()
 
+    readFile()
+
     val statementCtx = programCtx.compoundStatement().statements()
     pushScope()
     val statementAsts = if (statementCtx != null) {
@@ -162,7 +176,7 @@ class AstCreator(filename: String, global: Global)
     val terminalNode = ctx.children.asScala.map(_.asInstanceOf[TerminalNode]).head
     val token        = terminalNode.getSymbol
     val variableName = token.getText
-    val node         = createIdentiferWithScope(terminalNode, variableName, variableName, Defines.Any, List[String]())
+    val node         = createIdentiferWithScope(terminalNode, variableName, Defines.Any, List[String]())
     setIdentiferInScope(node)
     Seq(Ast(node))
   }
@@ -175,7 +189,7 @@ class AstCreator(filename: String, global: Global)
       val argsAsts    = astForArgumentsContext(ctx.arguments())
       val callNode = NewCall()
         .name(Operators.indexAccess)
-        .code(Operators.indexAccess)
+        .code(getCodeForTerminalNode(ctx.LBRACK()))
         .methodFullName(Operators.indexAccess)
         .signature("")
         .dispatchType(DispatchTypes.STATIC_DISPATCH)
@@ -196,12 +210,12 @@ class AstCreator(filename: String, global: Global)
       }
       val varSymbol = localVar.getSymbol()
       val node =
-        createIdentiferWithScope(localVar, varSymbol.getText, varSymbol.getText, Defines.Any, List(Defines.Any))
+        createIdentiferWithScope(localVar, varSymbol.getText, Defines.Any, List(Defines.Any))
       val yAst = Ast(node)
 
       val callNode = NewCall()
         .name(Operators.fieldAccess)
-        .code(Operators.fieldAccess)
+        .code(getCodeForTerminalNode(localVar))
         .methodFullName(Operators.fieldAccess)
         .signature("")
         .dispatchType(DispatchTypes.STATIC_DISPATCH)
@@ -734,8 +748,8 @@ class AstCreator(filename: String, global: Global)
     val primaryAst = astForPrimaryContext(ctx.primary())
     val localVar   = ctx.CONSTANT_IDENTIFIER()
     val varSymbol  = localVar.getSymbol()
-    val node = createIdentiferWithScope(localVar, varSymbol.getText, varSymbol.getText, Defines.Any, List(Defines.Any))
-    val constAst = Ast(node)
+    val node       = createIdentiferWithScope(localVar, varSymbol.getText, Defines.Any, List(Defines.Any))
+    val constAst   = Ast(node)
 
     val callNode = NewCall()
       .name(ctx.COLON2().getText)
@@ -802,11 +816,25 @@ class AstCreator(filename: String, global: Global)
       val classOrModuleRefAst =
         astForClassOrModuleReferenceContext(ctx.classDefinition().classOrModuleReference(), baseClassName)
       val bodyAst = astForBodyStatementContext(ctx.classDefinition().bodyStatement())
+      val bodyAstSansModifiers = bodyAst
+        .filterNot(ast => {
+          val nodes = ast.nodes
+            .filter(_.isInstanceOf[NewIdentifier])
+
+          if (nodes.size == 1) {
+            val varName = nodes
+              .map(_.asInstanceOf[NewIdentifier].name)
+              .head
+            varName == "public" || varName == "protected" || varName == "private"
+          } else {
+            false
+          }
+        })
 
       if (classStack.size > 0) {
         classStack.pop()
       }
-      Seq(classOrModuleRefAst.head.withChildren(bodyAst))
+      Seq(classOrModuleRefAst.head.withChildren(bodyAstSansModifiers))
     } else {
       // TODO test for this is pending due to lack of understanding to generate an example
       val astExprOfCommand = astForExpressionOrCommandContext(ctx.classDefinition().expressionOrCommand())
@@ -1193,7 +1221,7 @@ class AstCreator(filename: String, global: Global)
       .signature(localIdentifier.getText())
       .typeFullName(MethodFullNames.UnknownFullName)
       .dispatchType(DispatchTypes.STATIC_DISPATCH)
-      .code(localIdentifier.getText())
+      .code(getCodeForTerminalNode(localIdentifier))
       .lineNumber(line)
       .columnNumber(column)
     Seq(callAst(callNode))
@@ -1217,7 +1245,7 @@ class AstCreator(filename: String, global: Global)
       val varSymbol = localVar.getSymbol()
       if (lookupIdentiferInScope(varSymbol.getText)) {
         val node =
-          createIdentiferWithScope(localVar, varSymbol.getText, varSymbol.getText, Defines.Any, List(Defines.Any))
+          createIdentiferWithScope(localVar, varSymbol.getText, Defines.Any, List(Defines.Any))
         Seq(Ast(node))
       } else {
         astForCallNode(localVar)
@@ -1227,7 +1255,7 @@ class AstCreator(filename: String, global: Global)
       val varSymbol = localVar.getSymbol()
       if (lookupIdentiferInScope(varSymbol.getText)) {
         val node =
-          createIdentiferWithScope(localVar, varSymbol.getText, varSymbol.getText, Defines.Any, List(Defines.Any))
+          createIdentiferWithScope(localVar, varSymbol.getText, Defines.Any, List(Defines.Any))
         Seq(Ast(node))
       } else {
         astForCallNode(localVar)
@@ -1287,13 +1315,13 @@ class AstCreator(filename: String, global: Global)
       val localVar  = ctx.LOCAL_VARIABLE_IDENTIFIER()
       val varSymbol = localVar.getSymbol()
       val node =
-        createIdentiferWithScope(localVar, varSymbol.getText, varSymbol.getText, Defines.Any, List(Defines.Any))
+        createIdentiferWithScope(localVar, varSymbol.getText, Defines.Any, List(Defines.Any))
       Seq(Ast(node))
     } else if (ctx.CONSTANT_IDENTIFIER() != null) {
       val localVar  = ctx.CONSTANT_IDENTIFIER()
       val varSymbol = localVar.getSymbol()
       val node =
-        createIdentiferWithScope(localVar, varSymbol.getText, varSymbol.getText, Defines.Any, List(Defines.Any))
+        createIdentiferWithScope(localVar, varSymbol.getText, Defines.Any, List(Defines.Any))
       Seq(Ast(node))
     } else {
       Seq(Ast())
@@ -1368,7 +1396,7 @@ class AstCreator(filename: String, global: Global)
     localVarList
       .map(localVar => {
         val varSymbol = localVar.getSymbol()
-        createIdentiferWithScope(localVar, varSymbol.getText, varSymbol.getText, Defines.Any, Seq[String](Defines.Any))
+        createIdentiferWithScope(localVar, varSymbol.getText, Defines.Any, Seq[String](Defines.Any))
         val param = NewMethodParameterIn()
           .name(varSymbol.getText)
           .code(varSymbol.getText)
@@ -1547,7 +1575,7 @@ class AstCreator(filename: String, global: Global)
   def astForSimpleScopedConstantReferencePrimaryContext(ctx: SimpleScopedConstantReferencePrimaryContext): Seq[Ast] = {
     val localVar  = ctx.CONSTANT_IDENTIFIER()
     val varSymbol = localVar.getSymbol()
-    val node = createIdentiferWithScope(localVar, varSymbol.getText, varSymbol.getText, Defines.Any, List(Defines.Any))
+    val node      = createIdentiferWithScope(localVar, varSymbol.getText, Defines.Any, List(Defines.Any))
 
     val callNode = NewCall()
       .name(ctx.COLON2().getText)
@@ -1761,21 +1789,17 @@ class AstCreator(filename: String, global: Global)
 
   def astForUntilExpressionContext(ctx: UntilExpressionContext): Seq[Ast] = {
     // until will be modelled as a while
-    val untilNode = NewControlStructure()
-      .controlStructureType(ControlStructureTypes.WHILE)
-      .code(ctx.UNTIL().getText)
-      .lineNumber(ctx.UNTIL().getSymbol.getLine)
-      .columnNumber(ctx.UNTIL().getSymbol.getCharPositionInLine)
-
-    val untilCondAsts = astForExpressionOrCommandContext(ctx.expressionOrCommand())
-    val doClauseAst   = astForDoClauseContext(ctx.doClause())
-
-    Seq(
-      Ast(untilNode)
-        .withChildren(untilCondAsts)
-        .withConditionEdge(untilNode, untilCondAsts.head.nodes.head)
-        .withChildren(doClauseAst)
+    val untilCondAst = astForExpressionOrCommandContext(ctx.expressionOrCommand()).headOption
+    val doClauseAsts = astForDoClauseContext(ctx.doClause())
+
+    val ast = whileAst(
+      untilCondAst,
+      doClauseAsts,
+      Some(getCodeForTerminalNode(ctx.UNTIL())),
+      Some(ctx.UNTIL().getSymbol.getLine),
+      Some(ctx.UNTIL().getSymbol.getCharPositionInLine)
     )
+    Seq(ast)
   }
 
   def astForPseudoVariableIdentifierContext(ctx: PseudoVariableIdentifierContext): Seq[Ast] = {
@@ -1790,7 +1814,7 @@ class AstCreator(filename: String, global: Global)
       else return Seq(Ast())
     }
 
-    val astNode = createIdentiferWithScope(node, ctx.getText, ctx.getText, Defines.Any, List(Defines.Any))
+    val astNode = createIdentiferWithScope(node, ctx.getText, Defines.Any, List(Defines.Any))
     Seq(Ast(astNode))
   }
 
@@ -1811,21 +1835,17 @@ class AstCreator(filename: String, global: Global)
   }
 
   def astForWhileExpressionContext(ctx: WhileExpressionContext): Seq[Ast] = {
-    val whileNode = NewControlStructure()
-      .controlStructureType(ControlStructureTypes.WHILE)
-      .code(ctx.getText)
-      .lineNumber(ctx.WHILE().getSymbol.getLine)
-      .columnNumber(ctx.WHILE().getSymbol.getCharPositionInLine)
-
-    val whileCondAsts = astForExpressionOrCommandContext(ctx.expressionOrCommand())
-    val doClauseAsts  = astForDoClauseContext(ctx.doClause())
-
-    Seq(
-      Ast(whileNode)
-        .withChildren(whileCondAsts)
-        .withConditionEdge(whileNode, whileCondAsts.head.nodes.head)
-        .withChildren(doClauseAsts)
+    val whileCondAst = astForExpressionOrCommandContext(ctx.expressionOrCommand()).headOption
+    val doClauseAsts = astForDoClauseContext(ctx.doClause())
+
+    val ast = whileAst(
+      whileCondAst,
+      doClauseAsts,
+      Some(getCodeForTerminalNode(ctx.WHILE())),
+      Some(ctx.WHILE().getSymbol.getLine),
+      Some(ctx.WHILE().getSymbol.getCharPositionInLine)
     )
+    Seq(ast)
   }
 
   def astForBlockArgumentContext(ctx: BlockArgumentContext): Seq[Ast] = {
@@ -1838,8 +1858,7 @@ class AstCreator(filename: String, global: Global)
   }
 
   def astForBlockSplattingTypeArgumentsContext(ctx: BlockSplattingTypeArgumentsContext): Seq[Ast] = {
-    val blockNode = NewBlock().typeFullName(Defines.Any)
-    val splatAst  = astForSplattingArgumentContext(ctx.splattingArgument())
+    val splatAst = astForSplattingArgumentContext(ctx.splattingArgument())
     if (ctx.blockArgument() != null) {
       val blockArgAst = astForBlockArgumentContext(ctx.blockArgument())
       blockArgAst ++ splatAst

joern-cli/frontends/rubysrc2cpg/src/test/scala/io/joern/rubysrc2cpg/querying/IdentifierTests.scala

@@ -824,12 +824,6 @@ class IdentifierTests extends RubyCode2CpgFixture {
         .l
         .size shouldBe 1
     }
-    "recognise all call nodes" in {
-      cpg.method
-        .name("puts")
-        .l
-        .size shouldBe 1
-    }
 
     "successfully plot ASTs" in {
       cpg.method.name(":program").dotAst.l
@@ -867,9 +861,47 @@ class IdentifierTests extends RubyCode2CpgFixture {
         .size shouldBe 1
     }
     "successfully plot ASTs" in {
-      cpg.method.name(":some_method").dotAst.l
+      cpg.method.name("some_method").dotAst.l
       cpg.method.name(":program").dotAst.l
     }
   }
 
+  "CPG for code with private/protected/public" should {
+    val cpg = code(
+      """
+        |class SomeClass
+        |  private
+        |  def method1
+        |  end
+        |
+        |  protected
+        |  def method2
+        |  end
+        |
+        |  public
+        |  def method3
+        |  end
+        |end
+        |
+        |""".stripMargin)
+
+    "recognise all method nodes" in {
+      cpg.method
+        .name("method1")
+        .l
+        .size shouldBe 1
+      cpg.method
+        .name("method2")
+        .l
+        .size shouldBe 1
+      cpg.method
+        .name("method3")
+        .l
+        .size shouldBe 1
+
+    }
+    "successfully plot ASTs" in {
+      cpg.method.name(":program").dotAst.l
+    }
+  }
 }