From aa7692b474eaf555a1bede5760a9cd9f7db0b84e Mon Sep 17 00:00:00 2001
From: Konrad Weiss <konrad.weiss@aisec.fraunhofer.de>
Date: Wed, 24 Apr 2024 15:45:23 +0200
Subject: [PATCH] Fix Type resolution for array subscriptions that use a
 mapping

---
 .../frontends/solidity/ExpressionHandler.kt   |  1 +
 .../frontends/solidity/GraphExtensionsPass.kt | 44 ++++++++++++++-----
 2 files changed, 34 insertions(+), 11 deletions(-)

diff --git a/cpg-solidity/src/main/kotlin/de/fraunhofer/aisec/cpg/frontends/solidity/ExpressionHandler.kt b/cpg-solidity/src/main/kotlin/de/fraunhofer/aisec/cpg/frontends/solidity/ExpressionHandler.kt
index f291192af6..e6b4b34be9 100644
--- a/cpg-solidity/src/main/kotlin/de/fraunhofer/aisec/cpg/frontends/solidity/ExpressionHandler.kt
+++ b/cpg-solidity/src/main/kotlin/de/fraunhofer/aisec/cpg/frontends/solidity/ExpressionHandler.kt
@@ -125,6 +125,7 @@ class ExpressionHandler(lang: SolidityLanguageFrontend) : Handler<Expression, Pa
             val arraySub = newArraySubscriptionExpression(frontend.getCodeFromRawNode(ctx), ctx)
             this.handle(expressions[0])?.let { arraySub.arrayExpression = it }
             this.handle(expressions[1])?.let { arraySub.subscriptExpression = it }
+
             return arraySub
         }
 
diff --git a/cpg-solidity/src/main/kotlin/de/fraunhofer/aisec/cpg/frontends/solidity/GraphExtensionsPass.kt b/cpg-solidity/src/main/kotlin/de/fraunhofer/aisec/cpg/frontends/solidity/GraphExtensionsPass.kt
index ab30c14e89..34752bcbde 100644
--- a/cpg-solidity/src/main/kotlin/de/fraunhofer/aisec/cpg/frontends/solidity/GraphExtensionsPass.kt
+++ b/cpg-solidity/src/main/kotlin/de/fraunhofer/aisec/cpg/frontends/solidity/GraphExtensionsPass.kt
@@ -9,7 +9,10 @@ import de.fraunhofer.aisec.cpg.graph.declarations.FunctionDeclaration
 import de.fraunhofer.aisec.cpg.graph.edge.Properties
 import de.fraunhofer.aisec.cpg.graph.edge.PropertyEdge
 import de.fraunhofer.aisec.cpg.graph.statements.ReturnStatement
+import de.fraunhofer.aisec.cpg.graph.statements.expressions.ArraySubscriptionExpression
 import de.fraunhofer.aisec.cpg.graph.statements.expressions.CallExpression
+import de.fraunhofer.aisec.cpg.graph.types.ObjectType
+import de.fraunhofer.aisec.cpg.graph.types.UnknownType
 import de.fraunhofer.aisec.cpg.helpers.SubgraphWalker
 import de.fraunhofer.aisec.cpg.passes.EvaluationOrderGraphPass
 import de.fraunhofer.aisec.cpg.passes.Pass
@@ -23,9 +26,12 @@ class GraphExtensionsPass(ctx: TranslationContext): TranslationResultPass(ctx) {
 
     override fun accept(result: TranslationResult) {
         result?.let {
-            addAdditionalEOGEdges(it)
-            addAdditionalDFGEdges(it)
-            addReturnNodes(it)
+            val allNodes = SubgraphWalker.flattenAST(it)
+            var functions = allNodes.filterIsInstance<FunctionDeclaration>()
+            addSubscriptionType(allNodes)
+            addAdditionalEOGEdges(functions)
+            addAdditionalDFGEdges(functions)
+            addReturnNodes(allNodes)
             result.components
             result.components.map { it.language?.frontend }.filterIsInstance<SolidityLanguageFrontend>().forEach {
                 it.pragmas.forEach { result += it }
@@ -37,9 +43,26 @@ class GraphExtensionsPass(ctx: TranslationContext): TranslationResultPass(ctx) {
     override fun cleanup() {
     }
 
-    private fun addAdditionalEOGEdges(result: TranslationResult) {
-        var funcitons = SubgraphWalker.flattenAST(result).filterIsInstance<FunctionDeclaration>()
-        funcitons.forEach {
+    private fun addSubscriptionType(allNodes: List<Node>) {
+        var subscriptions = allNodes.filterIsInstance<ArraySubscriptionExpression>().toMutableList()
+        while(subscriptions.isNotEmpty()){
+            // Because a subscription expression can have another unresolved subscription expression as its arrayExpression
+            val toProcess = subscriptions.filter { !(it.arrayExpression.type is UnknownType) }
+            if(toProcess.isEmpty())
+                break // No more processeable expressions
+            toProcess.forEach {
+                // If the type of the subscript expression is a mapping access we have to set the type here as it is a custom type
+                val sType = it.arrayExpression.type
+                if(sType is ObjectType && sType.name.toString() == "mapping" && sType.generics.size >= 2){
+                    it.type = sType.generics[1]
+                }
+                subscriptions.remove(it)
+            }
+        }
+    }
+
+    private fun addAdditionalEOGEdges(functions: List<FunctionDeclaration>) {
+        functions.forEach {
             val function = it
             var body = it.body
             body?.let {
@@ -57,9 +80,8 @@ class GraphExtensionsPass(ctx: TranslationContext): TranslationResultPass(ctx) {
             }
         }
     }
-    private fun addAdditionalDFGEdges(result: TranslationResult) {
-        var funcitons = SubgraphWalker.flattenAST(result).filterIsInstance<FunctionDeclaration>()
-        funcitons.forEach {
+    private fun addAdditionalDFGEdges(functions: List<FunctionDeclaration>) {
+        functions.forEach {
             val function = it
             var body = it.body
             body?.let {
@@ -78,9 +100,9 @@ class GraphExtensionsPass(ctx: TranslationContext): TranslationResultPass(ctx) {
         }
     }
 
-    private fun addReturnNodes(result: TranslationResult){
+    private fun addReturnNodes(allNodes: List<Node>){
 
-        var calls = SubgraphWalker.flattenAST(result).filterIsInstance<CallExpression>()
+        var calls = allNodes.filterIsInstance<CallExpression>()
         calls.forEach {
             val call = it
             call.invokes.forEach {