From aa8615ce61ccd1f8ab22817a5dbc3a29bb7f3055 Mon Sep 17 00:00:00 2001
From: Christian Banse <christian.banse@aisec.fraunhofer.de>
Date: Fri, 24 Nov 2023 08:50:34 +0100
Subject: [PATCH] Added testcase for unwrapReference

---
 .../fraunhofer/aisec/cpg/graph/Extensions.kt  | 22 +++++++++++++++++++
 .../de/fraunhofer/aisec/cpg/helpers/Util.kt   | 16 --------------
 .../aisec/cpg/graph/ShortcutsTest.kt          | 15 +++++++++++++
 3 files changed, 37 insertions(+), 16 deletions(-)

diff --git a/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/Extensions.kt b/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/Extensions.kt
index 2da58dba07..2502a45947 100644
--- a/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/Extensions.kt
+++ b/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/graph/Extensions.kt
@@ -662,3 +662,25 @@ private fun Node.eogDistanceTo(to: Node): Int {
 
     return i
 }
+
+/**
+ * This is a small utility function to "unwrap" a [Reference] that it is wrapped in (multiple)
+ * [Expression] nodes. This will only work on expression that only have one "argument" (such as a
+ * unary operator), in order to avoid ambiguous results. This can be useful for data-flow analysis,
+ * if you want to quickly retrieve the reference that is affected by an operation. For example in
+ * C++ it is common to take an address of a variable and cast it into an appropriate type:
+ * ```cpp
+ * int64_t addr = (int64_t) &a;
+ * ```
+ *
+ * When called on the right-hand side of this assignment, this function will return `a`.
+ */
+fun Expression?.unwrapReference(): Reference? {
+    return when {
+        this is Reference -> this
+        this is UnaryOperator && (this.operatorCode == "*" || this.operatorCode == "&") ->
+            this.input.unwrapReference()
+        this is CastExpression -> this.expression.unwrapReference()
+        else -> null
+    }
+}
diff --git a/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/helpers/Util.kt b/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/helpers/Util.kt
index aed1750bb3..49edb26de3 100644
--- a/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/helpers/Util.kt
+++ b/cpg-core/src/main/kotlin/de/fraunhofer/aisec/cpg/helpers/Util.kt
@@ -405,20 +405,4 @@ object Util {
         ENTRIES,
         EXITS
     }
-
-    fun unwrapReference(node: Node?): Reference? {
-        return if (node is Reference) node
-        else if (node is UnaryOperator && (node.operatorCode == "*" || node.operatorCode == "&"))
-            unwrapReference(node.input)
-        else if (node is CastExpression) unwrapReference(node.expression)
-        else if (
-            node is BinaryOperator &&
-                (node.operatorCode == "&" || node.operatorCode == "*") &&
-                node.lhs is CastExpression &&
-                (node.lhs as CastExpression).expression is ProblemExpression
-        )
-        // TODO: This is only a hotfix for a bug in the language frontend.
-        unwrapReference(node.rhs)
-        else null
-    }
 }
diff --git a/cpg-core/src/test/kotlin/de/fraunhofer/aisec/cpg/graph/ShortcutsTest.kt b/cpg-core/src/test/kotlin/de/fraunhofer/aisec/cpg/graph/ShortcutsTest.kt
index 0663cc55eb..45bb92eea3 100644
--- a/cpg-core/src/test/kotlin/de/fraunhofer/aisec/cpg/graph/ShortcutsTest.kt
+++ b/cpg-core/src/test/kotlin/de/fraunhofer/aisec/cpg/graph/ShortcutsTest.kt
@@ -27,6 +27,7 @@ package de.fraunhofer.aisec.cpg.graph
 
 import de.fraunhofer.aisec.cpg.*
 import de.fraunhofer.aisec.cpg.frontends.TestLanguage
+import de.fraunhofer.aisec.cpg.frontends.TestLanguageFrontend
 import de.fraunhofer.aisec.cpg.graph.declarations.FunctionDeclaration
 import de.fraunhofer.aisec.cpg.graph.declarations.MethodDeclaration
 import de.fraunhofer.aisec.cpg.graph.declarations.RecordDeclaration
@@ -348,6 +349,20 @@ class ShortcutsTest {
         assertEquals(3, (paramPassed.last() as? Literal<*>)?.value)
     }
 
+    @Test
+    fun testUnwrapReference() {
+        with(TestLanguageFrontend()) {
+            val a = newReference("a")
+            val op = newUnaryOperator("&", prefix = true, postfix = false)
+            op.input = a
+            val cast = newCastExpression()
+            cast.castType = objectType("int64")
+            cast.expression = op
+
+            assertEquals(a, cast.unwrapReference())
+        }
+    }
+
     private lateinit var shortcutClassResult: TranslationResult
 
     @BeforeAll