From df29596a0122efc3077baddf67c79e2329bbeac9 Mon Sep 17 00:00:00 2001
From: Andrei Dreyer <andrei@whirlylabs.com>
Date: Tue, 12 Nov 2024 11:21:11 +0200
Subject: [PATCH] Accept default handling for MemberAccess and IndexAccess on
 assoc key (#5100)

---
 .../AstForExpressionsCreator.scala            |  2 +-
 .../rubysrc2cpg/querying/CallTests.scala      | 44 +++++++++++++++++++
 2 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/joern-cli/frontends/rubysrc2cpg/src/main/scala/io/joern/rubysrc2cpg/astcreation/AstForExpressionsCreator.scala b/joern-cli/frontends/rubysrc2cpg/src/main/scala/io/joern/rubysrc2cpg/astcreation/AstForExpressionsCreator.scala
index e27eb5c4389d..59bb2266426d 100644
--- a/joern-cli/frontends/rubysrc2cpg/src/main/scala/io/joern/rubysrc2cpg/astcreation/AstForExpressionsCreator.scala
+++ b/joern-cli/frontends/rubysrc2cpg/src/main/scala/io/joern/rubysrc2cpg/astcreation/AstForExpressionsCreator.scala
@@ -1004,7 +1004,7 @@ trait AstForExpressionsCreator(implicit withSchemaValidation: ValidationMode) {
       case keyIdentifier: SimpleIdentifier => setArgumentName(value, keyIdentifier.text)
       case symbol @ StaticLiteral(typ) if typ == getBuiltInType(Defines.Symbol) =>
         setArgumentName(value, symbol.text.stripPrefix(":"))
-      case _: (LiteralExpr | RubyCall | ProcOrLambdaExpr) => astForExpression(assoc)
+      case _: (LiteralExpr | RubyCall | ProcOrLambdaExpr | MemberAccess | IndexAccess) => astForExpression(assoc)
       case x =>
         logger.warn(s"Not explicitly handled argument association key of type ${x.getClass.getSimpleName}")
         astForExpression(assoc)
diff --git a/joern-cli/frontends/rubysrc2cpg/src/test/scala/io/joern/rubysrc2cpg/querying/CallTests.scala b/joern-cli/frontends/rubysrc2cpg/src/test/scala/io/joern/rubysrc2cpg/querying/CallTests.scala
index 38c2ce3ecbc2..46e217b88426 100644
--- a/joern-cli/frontends/rubysrc2cpg/src/test/scala/io/joern/rubysrc2cpg/querying/CallTests.scala
+++ b/joern-cli/frontends/rubysrc2cpg/src/test/scala/io/joern/rubysrc2cpg/querying/CallTests.scala
@@ -523,4 +523,48 @@ class CallTests extends RubyCode2CpgFixture(withPostProcessing = true) {
       case xs => fail(s"Expected 6 parameters for call, got [${xs.code.mkString(", ")}]")
     }
   }
+
+  "Call with association IndexAccess key" in {
+    val cpg = code("""
+        |foo(bar[:baz] => nil)
+        |""".stripMargin)
+
+    inside(cpg.call.name("foo").argument.l) {
+      case _ :: (assocParam: Call) :: Nil =>
+        assocParam.methodFullName shouldBe RubyOperators.association
+        assocParam.code shouldBe "bar[:baz] => nil"
+
+        inside(assocParam.argument.l) {
+          case (lhs: Call) :: (rhs: Literal) :: Nil =>
+            lhs.methodFullName shouldBe Operators.indexAccess
+            lhs.code shouldBe "bar[:baz]"
+
+            rhs.code shouldBe "nil"
+          case xs => fail(s"Expected lhs and rhs for association, got [${xs.code.mkString(",")}]")
+        }
+      case xs => fail(s"Expected two params, got [${xs.code.mkString(",")}]")
+    }
+  }
+
+  "Call with association MemberAccess key" in {
+    val cpg = code("""
+        |foo(bar.baz => nil)
+        |""".stripMargin)
+
+    inside(cpg.call.name("foo").argument.l) {
+      case _ :: (assocParam: Call) :: Nil =>
+        assocParam.methodFullName shouldBe RubyOperators.association
+        assocParam.code shouldBe "bar.baz => nil"
+
+        inside(assocParam.argument.l) {
+          case (lhs: Call) :: (rhs: Literal) :: Nil =>
+            lhs.methodFullName shouldBe Operators.fieldAccess
+            lhs.code shouldBe "bar.baz"
+
+            rhs.code shouldBe "nil"
+          case xs => fail(s"Expected lhs and rhs for association, got [${xs.code.mkString(",")}]")
+        }
+      case xs => fail(s"Expected two params, got [${xs.code.mkString(",")}]")
+    }
+  }
 }