[ruby] capture self variable

Author: AndreiDreyer

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

@@ -100,6 +100,12 @@ trait AstForFunctionsCreator(implicit withSchemaValidation: ValidationMode) { th
 
     // Consider which variables are captured from the outer scope
     val stmtBlockAst = if (isClosure || isSingletonObjectMethod) {
+      // create closure local used for capturing
+      createClosureBindingInformation(scope.lookupSelfInOuterScope.toSet)
+        .collect { case (_, name, _, Some(closureBindingId)) =>
+          val capturingLocal = localNode(node.body, name, name, Defines.Any, closureBindingId = Option(closureBindingId))
+          scope.addToScope(capturingLocal.name, capturingLocal)
+        }
       val baseStmtBlockAst = astForMethodBody(node.body, optionalStatementList)
       transformAsClosureBody(node.body, refs, baseStmtBlockAst)
     } else {
@@ -214,22 +220,25 @@ trait AstForFunctionsCreator(implicit withSchemaValidation: ValidationMode) { th
   private def transformAsClosureBody(originNode: RubyExpression, refs: List[Ast], baseStmtBlockAst: Ast) = {
     // Determine which locals are captured
     val capturedLocalNodes = baseStmtBlockAst.nodes
-      .collect { case x: NewIdentifier => x }
+      .collect { case x: NewIdentifier if x.name != Defines.Self => x }
       .distinctBy(_.name)
       .map {
-        case i if i.name == "self" => scope.lookupSelfInOuterScope // we only need to bind to the closest self in scope
-        case i                     => scope.lookupVariableInOuterScope(i.name)
+        case i if i.name == Defines.Self => scope.lookupSelfInOuterScope
+        case i => scope.lookupVariableInOuterScope(i.name)
       }
       .filter(_.iterator.nonEmpty)
       .flatten
       .toSet
 
+    val selfLocal = scope.lookupSelfInOuterScope.toSet
+    val capturedNodes = capturedLocalNodes ++ selfLocal
+
     val capturedIdentifiers = baseStmtBlockAst.nodes.collect {
-      case i: NewIdentifier if capturedLocalNodes.map(_.name).contains(i.name) => i
+      case i: NewIdentifier if capturedNodes.map(_.name).contains(i.name) => i
     }
     // Copy AST block detaching the REF nodes between parent locals/params and identifiers, with the closures' one
     val capturedBlockAst = baseStmtBlockAst.copy(refEdges = baseStmtBlockAst.refEdges.filterNot {
-      case AstEdge(_: NewIdentifier, dst: DeclarationNew) => capturedLocalNodes.contains(dst)
+      case AstEdge(_: NewIdentifier, dst: DeclarationNew) => capturedNodes.contains(dst)
       case _                                              => false
     })
 
@@ -238,15 +247,8 @@ trait AstForFunctionsCreator(implicit withSchemaValidation: ValidationMode) { th
     val astChildren  = mutable.Buffer.empty[NewNode]
     val refEdges     = mutable.Buffer.empty[(NewNode, NewNode)]
     val captureEdges = mutable.Buffer.empty[(NewNode, NewNode)]
-    capturedLocalNodes
-      .collect {
-        case local: NewLocal =>
-          val closureBindingId = scope.variableScopeFullName(local.name).map(x => s"$x.${local.name}")
-          (local, local.name, local.code, closureBindingId)
-        case param: NewMethodParameterIn =>
-          val closureBindingId = scope.variableScopeFullName(param.name).map(x => s"$x.${param.name}")
-          (param, param.name, param.code, closureBindingId)
-      }
+
+    createClosureBindingInformation(capturedNodes)
       .collect { case (capturedLocal, name, code, Some(closureBindingId)) =>
         val capturingLocal =
           localNode(originNode, name, name, Defines.Any, closureBindingId = Option(closureBindingId))
@@ -624,4 +626,16 @@ trait AstForFunctionsCreator(implicit withSchemaValidation: ValidationMode) { th
       case _                      => false
     }
   }
+
+  private def createClosureBindingInformation(capturedNodes: Set[DeclarationNew]): Set[(DeclarationNew, String, String, Option[String])] = {
+    capturedNodes
+      .collect {
+        case local: NewLocal =>
+          val closureBindingId = scope.variableScopeFullName(local.name).map(x => s"$x.${local.name}")
+          (local, local.name, local.code, closureBindingId)
+        case param: NewMethodParameterIn =>
+          val closureBindingId = scope.variableScopeFullName(param.name).map(x => s"$x.${param.name}")
+          (param, param.name, param.code, closureBindingId)
+      }
+  }
 }

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

@@ -68,7 +68,7 @@ class DoBlockTests extends RubyCode2CpgFixture {
 
     "have the return node under the closure (returning the literal)" in {
       inside(cpg.method("<lambda>0").block.astChildren.l) {
-        case ret :: Nil =>
+        case ret :: _ :: Nil =>
           ret.code shouldBe "\"world!\""
         case xs => fail(s"Expected the closure to have a single call, instead got [${xs.code.mkString(", ")}]")
       }
@@ -229,9 +229,9 @@ class DoBlockTests extends RubyCode2CpgFixture {
       }
     }
 
-    "annotate the nodes via CAPTURE bindings" in {
+    "nnotate the nodes via CAPTURE bindings" in {
       cpg.closureBinding.l match {
-        case myValue :: Nil =>
+        case myValue :: _ :: Nil =>
           inside(myValue._localViaRefOut) {
             case Some(local) =>
               local.name shouldBe "myValue"
@@ -400,17 +400,19 @@ class DoBlockTests extends RubyCode2CpgFixture {
                      |""".stripMargin)
 
     inside(cpg.local.nameNot(".*<tmp-\\d>").l) {
-      case jfsOutsideLocal :: schedules :: hashInsideLocal :: jfsCapturedLocal :: Nil =>
+      case jfsOutsideLocal :: _ :: hashInsideLocal :: jfsCapturedLocal :: selfCapturedLocal :: Nil =>
         jfsOutsideLocal.closureBindingId shouldBe None
         hashInsideLocal.closureBindingId shouldBe None
         jfsCapturedLocal.closureBindingId shouldBe Some("Test0.rb:<main>.get_pto_schedule.jfs")
+        selfCapturedLocal.closureBindingId shouldBe Some("Test0.rb:<main>.get_pto_schedule.<lambda>0.self")
       case xs => fail(s"Expected 6 locals, got ${xs.code.mkString(",")}")
     }
 
     inside(cpg.method.isLambda.local.l) {
-      case hashLocal :: _ :: jfsLocal :: Nil =>
+      case hashLocal :: _ :: jfsLocal :: selfLocal :: Nil =>
         hashLocal.closureBindingId shouldBe None
         jfsLocal.closureBindingId shouldBe Some("Test0.rb:<main>.get_pto_schedule.jfs")
+        selfLocal.closureBindingId shouldBe Some("Test0.rb:<main>.get_pto_schedule.<lambda>0.self")
       case xs => fail(s"Expected 3 locals in lambda, got ${xs.code.mkString(",")}")
     }
   }
@@ -459,7 +461,7 @@ class DoBlockTests extends RubyCode2CpgFixture {
           gSplatLocal.code shouldBe "g"
 
           selfLocal.code shouldBe "self"
-          selfLocal.closureBindingId shouldBe Some("Test0.rb:<global>.self")
+          selfLocal.closureBindingId shouldBe Some("Test0.rb:<main>.<lambda>0.self")
         case xs => fail(s"Expected 4 locals, got [${xs.name.mkString(", ")}]")
       }
     }

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

@@ -247,7 +247,7 @@ class ErbTests extends RubyCode2CpgFixture {
           lambdaParamForm.code shouldBe "form"
 
           inside(lambdaMethod.body.astChildren.l) {
-            case _ :: _ :: (appendCallTemplate: Call) :: _ :: Nil =>
+            case _ :: _ :: (appendCallTemplate: Call) :: _ :: _ :: Nil =>
               appendCallTemplate.code shouldBe "joern__inner_buffer << <%= form.text_field :name %>"
               val List(appendCallArgOne, appendCallArgTwo: Call) = appendCallTemplate.argument.l: @unchecked
 
@@ -390,15 +390,15 @@ class ErbTests extends RubyCode2CpgFixture {
     inside(cpg.method.name("<main>").parameter.name("self").l) {
       case selfMain :: Nil =>
         selfMain._closureBindingViaRefIn.map(_.closureBindingId).toList shouldBe List(
-          Some("index.html.erb:<global>.self")
+          Some("index.html.erb:<main>.<lambda>0.self")
         )
         selfMain.closureBindingId shouldBe None
       case xs => fail(s"Expected one local in global, got ${xs.name.mkString("[", ",", "]")}")
     }
 
     inside(cpg.method.isLambda.local.name("self").l) {
       case selfLocal :: Nil =>
-        selfLocal.closureBindingId shouldBe Some("index.html.erb:<global>.self")
+        selfLocal.closureBindingId shouldBe Some("index.html.erb:<main>.<lambda>0.self")
       case xs => fail(s"Expected one self local, got ${xs.name.mkString("[", ",", "]")}")
     }
   }
@@ -452,7 +452,14 @@ class ErbTests extends RubyCode2CpgFixture {
       cpg.method.fullNameExact("Test0.rb:<main>.Admin.UsersController.show.<lambda>0.<lambda>0").local.name("self").l
     ) {
       case selfLocal :: Nil =>
-        selfLocal.closureBindingId shouldBe Some("Test0.rb:<main>.Admin.UsersController.show.self")
+        selfLocal.closureBindingId shouldBe Some("Test0.rb:<main>.Admin.UsersController.show.<lambda>0.<lambda>0.self")
+      case xs => fail(s"Expected one local for self, got ${xs.name.mkString("[", ",", "]")}")
+    }
+
+    inside(
+      cpg.method.fullNameExact("Test0.rb:<main>.Admin.UsersController.show.<lambda>0").local.name("self").l) {
+      case selfLocal :: Nil =>
+        selfLocal.closureBindingId shouldBe Some("Test0.rb:<main>.Admin.UsersController.show.<lambda>0.self")
       case xs => fail(s"Expected one local for self, got ${xs.name.mkString("[", ",", "]")}")
     }
   }

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

@@ -412,7 +412,7 @@ class MethodReturnTests extends RubyCode2CpgFixture {
 
     "have the return node under the closure (returning the literal)" in {
       inside(cpg.method("<lambda>0").block.astChildren.l) {
-        case ret :: Nil =>
+        case ret :: _ :: Nil =>
           ret.code shouldBe "\"hello\""
         case xs => fail(s"Expected the closure to have a single call, instead got [${xs.code.mkString(", ")}]")
       }