Swift: Fix for extensions.

Author: geoffw0

swift/ql/lib/codeql/swift/dataflow/ExternalFlow.qll

@@ -402,6 +402,14 @@ private NominalType getDeclType(IterableDeclContext decl) {
   result = decl.(ProtocolDecl).getType()
 }
 
+private NominalTypeDecl resolveExtensions(IterableDeclContext decl) {
+  // TODO: this should be a method on IterableDeclContext
+  result = decl.(NominalTypeDecl)
+  or
+  result = decl.(ExtensionDecl).getExtendedTypeDecl()
+  // TODO: or a protocol added by the extension
+}
+
 /**
  * Gets the element in module `namespace` that satisfies the following properties:
  * 1. If the element is a member of a class-like type, then the class-like type has name `type`
@@ -453,10 +461,10 @@ private Element interpretElement0(
       result = field
     |
       subtypes = true and
-      decl = nomTypeDecl.getADerivedTypeDecl*()
+      resolveExtensions(decl) = nomTypeDecl.getADerivedTypeDecl*()
       or
       subtypes = false and
-      decl = nomTypeDecl
+      resolveExtensions(decl) = nomTypeDecl
     )
   )
 }

swift/ql/test/library-tests/dataflow/flowsources/FlowSources.expected

@@ -20,16 +20,20 @@
 | generics.swift:55:9:55:17 | .source2 | external |
 | generics.swift:57:9:57:17 | .source4 | external |
 | generics.swift:58:9:58:17 | .source5 | external |
+| generics.swift:60:9:60:17 | .source7 | external |
 | generics.swift:62:9:62:18 | .source1 | external |
 | generics.swift:63:9:63:18 | .source2 | external |
 | generics.swift:65:9:65:18 | .source9 | external |
 | generics.swift:66:9:66:18 | .source10 | external |
 | generics.swift:67:9:67:27 | call to source11() | external |
+| generics.swift:68:9:68:18 | .source12 | external |
 | generics.swift:69:9:69:27 | call to source13() | external |
 | generics.swift:88:9:88:15 | .source1 | external |
 | generics.swift:89:9:89:15 | .source2 | external |
 | generics.swift:90:9:90:14 | .source1 | external |
+| generics.swift:91:9:91:14 | .source2 | external |
 | generics.swift:92:9:92:15 | .source1 | external |
+| generics.swift:93:9:93:15 | .source2 | external |
 | generics.swift:112:9:112:15 | .source1 | external |
 | generics.swift:113:9:113:15 | .source2 | external |
 | nsdata.swift:18:17:18:17 | call to init(contentsOf:) | external |

swift/ql/test/library-tests/dataflow/flowsources/generics.swift

@@ -57,15 +57,15 @@ func useDerived(generic: MyGeneric<Int>, generic2: MyGeneric<Any>, derived: MyDe
     _ = derived.source4 // SOURCE
     _ = derived.source5 // SOURCE
     _ = derived.source6() // SOURCE [NOT DETECTED]
-    _ = derived.source7 // SOURCE [NOT DETECTED]
+    _ = derived.source7 // SOURCE
     _ = derived.source8() // SOURCE [NOT DETECTED]
     _ = derived2.source1 // SOURCE
     _ = derived2.source2 // SOURCE
     _ = derived2.source3() // SOURCE [NOT DETECTED]
     _ = derived2.source9 // SOURCE
     _ = derived2.source10 // SOURCE
     _ = derived2.source11() // SOURCE
-    _ = derived2.source12 // SOURCE [NOT DETECTED]
+    _ = derived2.source12 // SOURCE
     _ = derived2.source13() // SOURCE
 }
 
@@ -88,9 +88,9 @@ func useProtocol(proto: MyProtocol, impl: MyImpl<Int>, impl2: MyImpl<Any>) {
     _ = proto.source1 // SOURCE
     _ = proto.source2 // SOURCE
     _ = impl.source1 // SOURCE
-    _ = impl.source2 // SOURCE [NOT DETECTED]
+    _ = impl.source2 // SOURCE
     _ = impl2.source1 // SOURCE
-    _ = impl2.source2 // SOURCE [NOT DETECTED]
+    _ = impl2.source2 // SOURCE
 }
 
 // ---