add InclusionTest to PostMessageEventSanitizer

javascript/ql/lib/semmle/javascript/dataflow/TaintTracking.qll

@@ -1225,19 +1225,25 @@ module TaintTracking {
    * An equality test on `e.origin` or `e.source` where `e` is a `postMessage` event object,
    * considered as a sanitizer for `e`.
    */
-  private class PostMessageEventSanitizer extends AdditionalSanitizerGuardNode, DataFlow::ValueNode {
+  private class PostMessageEventSanitizer extends AdditionalSanitizerGuardNode {
     VarAccess event;
-    override EqualityTest astNode;
+    boolean polarity;
 
     PostMessageEventSanitizer() {
-      exists(string prop | prop = "origin" or prop = "source" |
-        astNode.getAnOperand().(PropAccess).accesses(event, prop) and
-        event.mayReferToParameter(any(PostMessageEventHandler h).getEventParameter())
+      event.mayReferToParameter(any(PostMessageEventHandler h).getEventParameter()) and
+      exists(DataFlow::PropRead read | read.accesses(event.flow(), ["origin", "source"]) |
+        exists(EqualityTest test | polarity = test.getPolarity() and this.getAstNode() = test |
+          test.getAnOperand().flow() = read
+        )
+        or
+        exists(InclusionTest test | polarity = test.getPolarity() and this = test |
+          test.getContainedNode() = read
+        )
       )
     }
 
     override predicate sanitizes(boolean outcome, Expr e) {
-      outcome = astNode.getPolarity() and
+      outcome = polarity and
       e = event
     }
 

javascript/ql/test/query-tests/Security/CWE-079/DomBasedXss/addEventListener.js

@@ -14,4 +14,17 @@ function test() {
     }
 
     window.addEventListener("message", foo.bind(null, {data: 'items'}));
+
+    window.onmessage = e => {
+        if (e.origin !== "https://foobar.com") {
+            return;
+        }
+        document.write(e.data); // OK - there is an origin check
+    }
+
+    window.onmessage = e => {
+        if (mySet.includes(e.origin)) {
+            document.write(e.data); // OK - there is an origin check
+        }
+    }
 }