Rework types exported through JSContext

Better model the JSExport protocol logic

Author: atorralba

swift/ql/lib/codeql/swift/frameworks/StandardLibrary/WebView.qll

@@ -91,39 +91,43 @@ private class JsValueSummaries extends SummaryModelCsv {
   }
 }
 
-/** The class `JSContext`. */
-private class JsContextDecl extends ClassDecl {
-  JsContextDecl() { this.getName() = "JSContext" }
+/** The `JSExport` protocol. */
+private class JsExport extends ProtocolDecl {
+  JsExport() { this.getName() = "JSExport" }
 }
 
-/** The type of an object exposed to JavaScript through `JSContext.setObject`. */
-private class TypeExposedThroughJsContext extends Type {
-  TypeExposedThroughJsContext() {
-    exists(ApplyExpr c, FuncDecl f |
-      c.getStaticTarget() = f and
-      f.getEnclosingDecl() instanceof JsContextDecl and
-      f.getName() = "setObject(_:forKeyedSubscript)"
-    |
-      c.getArgument(0).getExpr().getType() = this
-    )
-  }
+/** A protocol inheriting `JSExport`. */
+private class JsExportedProto extends ProtocolDecl {
+  JsExportedProto() { this.getABaseTypeDecl+() instanceof JsExport }
+}
+
+/** A type that adopts a `JSExport`-inherited protocol. */
+private class JsExportedType extends ClassOrStructDecl {
+  JsExportedType() { this.getABaseTypeDecl*() instanceof JsExportedProto }
 }
 
 /**
- * The members (fields and parameters of functions) of a class or struct
- * an instance of which is exposed to JavaScript through `JSContext.setObject`.
+ * A flow source that models properties and methods defined in a `JSExport`-inherited protocol
+ * and implemented in a type adopting that protcol. These members are accessible from JavaScript
+ * when the object is assigned to a `JSContext`.
  */
-private class ExposedThroughJsContextSource extends RemoteFlowSource {
-  ExposedThroughJsContextSource() {
-    exists(ParamDecl p | this.(DataFlow::ParameterNode).getParameter() = p |
-      p.getDeclaringFunction().getEnclosingDecl().(ClassOrStructDecl).getType() instanceof
-        TypeExposedThroughJsContext
+private class JsExportedSource extends RemoteFlowSource {
+  JsExportedSource() {
+    exists(MethodDecl adopter, MethodDecl base |
+      base.getEnclosingDecl() instanceof JsExportedProto and
+      adopter.getEnclosingDecl() instanceof JsExportedType
+    |
+      this.(DataFlow::ParameterNode).getParameter().getDeclaringFunction() = adopter and
+      adopter.getName() = base.getName()
     )
     or
-    exists(FieldDecl f | this.asDefinition().getSourceVariable() = f |
-      f.getEnclosingDecl().(ClassOrStructDecl).getType() instanceof TypeExposedThroughJsContext
+    exists(FieldDecl adopter, FieldDecl base |
+      base.getEnclosingDecl() instanceof JsExportedProto and
+      adopter.getEnclosingDecl() instanceof JsExportedType
+    |
+      this.asDefinition().getSourceVariable() = adopter and adopter.getName() = base.getName()
     )
   }
 
-  override string getSourceType() { result = "Member of a type exposed through JSContext" }
+  override string getSourceType() { result = "Member of a type exposed through JSExport" }
 }

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

@@ -8,6 +8,9 @@
 | url.swift:53:15:53:19 | .resourceBytes | external |
 | url.swift:60:15:60:19 | .lines | external |
 | url.swift:67:16:67:22 | .lines | external |
-| webview.swift:19:82:19:102 | message | external |
-| webview.swift:24:5:24:13 | .globalObject | external |
-| webview.swift:25:5:25:39 | call to objectForKeyedSubscript(_:) | external |
+| webview.swift:20:82:20:102 | message | external |
+| webview.swift:25:5:25:13 | .globalObject | external |
+| webview.swift:26:5:26:39 | call to objectForKeyedSubscript(_:) | external |
+| webview.swift:38:10:38:10 | self | Member of a type exposed through JSExport |
+| webview.swift:38:18:38:24 | arg1 | Member of a type exposed through JSExport |
+| webview.swift:38:29:38:35 | arg2 | Member of a type exposed through JSExport |

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

@@ -13,6 +13,7 @@ class JSContext {
     func objectForKeyedSubscript(_: Any!) -> JSValue! { return JSValue() } 
     func setObject(_: Any, forKeyedSubscript: (NSCopying & NSObjectProtocol) ) {}
 }
+protocol JSExport {}
 
 // --- tests ---
 class TestMessageHandler: WKScriptMessageHandler {
@@ -23,13 +24,20 @@ class TestMessageHandler: WKScriptMessageHandler {
 func testJsContext(context: JSContext) {
     context.globalObject // SOURCE
     context.objectForKeyedSubscript("") // SOURCE
-    context.setObject(Exposed.self, forKeyedSubscript: "exposed" as! NSCopying & NSObjectProtocol)
 }
 
-class Exposed {
+protocol Exported : JSExport {
+    var tainted: Any { get }
+    func tainted(arg1: Any, arg2: Any)
+}
+class ExportedImpl : Exported {
     var tainted: Any { get { return "" } } // SOURCE
 
+    var notTainted: Any { get { return ""} }
+
     func tainted(arg1: Any, arg2: Any) { // SOURCES
+    }
 
+    func notTainted(arg1: Any, arg2: Any) {
     }
 }