Kotlin: Change return type of Android specific ConcurrentHashMap.keySet

Author: tamasvajk

java/kotlin-extractor/src/main/kotlin/KotlinFileExtractor.kt

@@ -704,16 +704,17 @@ open class KotlinFileExtractor(
 
                 val paramsSignature = allParamTypes.joinToString(separator = ",", prefix = "(", postfix = ")") { it.javaResult.signature!! }
 
-                val substReturnType = typeSubstitution?.let { it(f.returnType, TypeContext.RETURN, pluginContext) } ?: f.returnType
+                val adjustedReturnType = getAdjustedReturnType(f)
+                val substReturnType = typeSubstitution?.let { it(adjustedReturnType, TypeContext.RETURN, pluginContext) } ?: adjustedReturnType
 
                 val locId = locOverride ?: getLocation(f, classTypeArgsIncludingOuterClasses)
 
                 if (f.symbol is IrConstructorSymbol) {
                     val unitType = useType(pluginContext.irBuiltIns.unitType, TypeContext.RETURN)
                     val shortName = when {
-                        f.returnType.isAnonymous -> ""
+                        adjustedReturnType.isAnonymous -> ""
                         typeSubstitution != null -> useType(substReturnType).javaResult.shortName
-                        else -> f.returnType.classFqName?.shortName()?.asString() ?: f.name.asString()
+                        else -> adjustedReturnType.classFqName?.shortName()?.asString() ?: f.name.asString()
                     }
                     val constrId = id.cast<DbConstructor>()
                     tw.writeConstrs(constrId, shortName, "$shortName$paramsSignature", unitType.javaResult.id, parentId, sourceDeclaration.cast<DbConstructor>())

java/kotlin-extractor/src/main/kotlin/KotlinUsesExtractor.kt

@@ -1,6 +1,7 @@
 package com.github.codeql
 
 import com.github.codeql.utils.*
+import com.github.codeql.utils.versions.codeQlWithHasQuestionMark
 import com.github.codeql.utils.versions.isRawType
 import com.semmle.extractor.java.OdasaOutput
 import org.jetbrains.kotlin.backend.common.extensions.IrPluginContext
@@ -861,6 +862,27 @@ open class KotlinUsesExtractor(
         return getFunctionLabel(f, null, classTypeArgsIncludingOuterClasses)
     }
 
+    fun getAdjustedReturnType(f: IrFunction) : IrType {
+        // The return type of `java.util.concurrent.ConcurrentHashMap<K,V>.keySet/0` is defined as `Set<K>` in the stubs inside the Android SDK.
+        // This does not match the Java SDK return type: `ConcurrentHashMap.KeySetView<K,V>`, so it's adjusted here.
+        // This is a deliberate change in the Android SDK: https://github.com/AndroidSDKSources/android-sdk-sources-for-api-level-31/blob/2c56b25f619575bea12f9c5520ed2259620084ac/java/util/concurrent/ConcurrentHashMap.java#L1244-L1249
+        // The annotation on the source is not visible in the android.jar, so we can't make the change based on that.
+        // TODO: there are other instances of `dalvik.annotation.codegen.CovariantReturnType` in the Android SDK, we should handle those too if they cause DB inconsistencies
+        val parentClass = f.parentClassOrNull
+        if (parentClass == null ||
+            parentClass.fqNameWhenAvailable?.asString() != "java.util.concurrent.ConcurrentHashMap" ||
+            getFunctionShortName(f).nameInDB != "keySet" ||
+            f.valueParameters.isNotEmpty() ||
+            f.returnType.classFqName?.asString() != "kotlin.collections.MutableSet") {
+            return f.returnType
+        }
+
+        val otherKeySet = parentClass.declarations.filterIsInstance<IrFunction>().find { it.name.asString() == "keySet" && it.valueParameters.size == 1 }
+            ?: return f.returnType
+
+        return otherKeySet.returnType.codeQlWithHasQuestionMark(false)
+    }
+
     /*
      * There are some pairs of classes (e.g. `kotlin.Throwable` and
      * `java.lang.Throwable`) which are really just 2 different names
@@ -877,7 +899,7 @@ open class KotlinUsesExtractor(
             maybeParentId,
             getFunctionShortName(f).nameInDB,
             f.valueParameters,
-            f.returnType,
+            getAdjustedReturnType(f),
             f.extensionReceiverParameter,
             getFunctionTypeParameters(f),
             classTypeArgsIncludingOuterClasses,

java/ql/test/kotlin/library-tests/android_function_return_types/returnTypes.expected

@@ -0,0 +1,4 @@
+| test.kt:4:5:6:5 | keySet | 0 | java.util.concurrent.ConcurrentHashMap$KeySetView<K,V> |
+| test.kt:8:5:10:5 | keySet | 1 | java.util.concurrent.ConcurrentHashMap$KeySetView<K,V> |
+| test.kt:17:5:19:5 | keySet | 0 | java.util.Set<K> |
+| test.kt:21:5:23:5 | keySet | 1 | java.util.concurrent.OtherConcurrentHashMap$KeySetView<K,V> |

java/ql/test/kotlin/library-tests/android_function_return_types/returnTypes.ql

@@ -0,0 +1,5 @@
+import java
+
+from Method m
+where m.fromSource()
+select m, m.getNumberOfParameters(), m.getReturnType().(RefType).getQualifiedName()

java/ql/test/kotlin/library-tests/android_function_return_types/test.kt

@@ -0,0 +1,27 @@
+package java.util.concurrent
+
+class ConcurrentHashMap<K,V> {
+    fun keySet(): MutableSet<K> {
+        return null!!
+    }
+
+    fun keySet(p: V): KeySetView<K,V>? {
+        return null
+    }
+
+    class KeySetView<K,V> {
+    }
+}
+
+class OtherConcurrentHashMap<K,V> {
+    fun keySet(): MutableSet<K> {
+        return null!!
+    }
+
+    fun keySet(p: V): KeySetView<K,V>? {
+        return null
+    }
+
+    class KeySetView<K,V> {
+    }
+}