From fdec904a7a69dee901b2240f71d331306eb8475b Mon Sep 17 00:00:00 2001
From: tombolton <tombolton@github.com>
Date: Wed, 9 Mar 2022 14:01:30 +0000
Subject: [PATCH] add new security queries to extraction query

---
 .../extraction/ExtractEndpointData.qll          | 15 +++++++++++++++
 .../modelbuilding/extraction/Queries.qll        | 17 ++++++++++++++++-
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/javascript/ql/experimental/adaptivethreatmodeling/modelbuilding/extraction/ExtractEndpointData.qll b/javascript/ql/experimental/adaptivethreatmodeling/modelbuilding/extraction/ExtractEndpointData.qll
index bea47423797..9231d66ad97 100644
--- a/javascript/ql/experimental/adaptivethreatmodeling/modelbuilding/extraction/ExtractEndpointData.qll
+++ b/javascript/ql/experimental/adaptivethreatmodeling/modelbuilding/extraction/ExtractEndpointData.qll
@@ -18,6 +18,9 @@ import experimental.adaptivethreatmodeling.NosqlInjectionATM as NosqlInjectionAT
 import experimental.adaptivethreatmodeling.SqlInjectionATM as SqlInjectionATM
 import experimental.adaptivethreatmodeling.TaintedPathATM as TaintedPathATM
 import experimental.adaptivethreatmodeling.XssATM as XssATM
+import experimental.adaptivethreatmodeling.StoredXssATM as StoredXssATM
+import experimental.adaptivethreatmodeling.XssThroughDomATM as XssThroughDomATM
+import experimental.adaptivethreatmodeling.CodeInjectionATM as CodeInjectionATM
 import Labels
 import NoFeaturizationRestrictionsConfig
 import Queries
@@ -32,6 +35,12 @@ ATMConfig getATMCfg(Query query) {
   query instanceof TaintedPathQuery and result instanceof TaintedPathATM::TaintedPathATMConfig
   or
   query instanceof XssQuery and result instanceof XssATM::DomBasedXssATMConfig
+  or
+  query instanceof StoredXssQuery and result instanceof StoredXssATM::StoredXssATMConfig
+  or
+  query instanceof XssThroughDomQuery and result instanceof XssThroughDomATM::XssThroughDOMATMConfig
+  or
+  query instanceof CodeInjectionQuery and result instanceof CodeInjectionATM::CodeInjectionATMConfig
 }
 
 /** Gets the ATM data flow configuration for the specified query. */
@@ -43,6 +52,12 @@ DataFlow::Configuration getDataFlowCfg(Query query) {
   query instanceof TaintedPathQuery and result instanceof TaintedPathATM::Configuration
   or
   query instanceof XssQuery and result instanceof XssATM::Configuration
+  or
+  query instanceof StoredXssQuery and result instanceof StoredXssATM::Configuration
+  or
+  query instanceof XssThroughDomQuery and result instanceof XssThroughDomATM::Configuration
+  or
+  query instanceof CodeInjectionQuery and result instanceof CodeInjectionATM::Configuration
 }
 
 /** Gets a known sink for the specified query. */
diff --git a/javascript/ql/experimental/adaptivethreatmodeling/modelbuilding/extraction/Queries.qll b/javascript/ql/experimental/adaptivethreatmodeling/modelbuilding/extraction/Queries.qll
index 51dd3ffec84..1fd80c9a85f 100644
--- a/javascript/ql/experimental/adaptivethreatmodeling/modelbuilding/extraction/Queries.qll
+++ b/javascript/ql/experimental/adaptivethreatmodeling/modelbuilding/extraction/Queries.qll
@@ -8,7 +8,10 @@ newtype TQuery =
   TNosqlInjectionQuery() or
   TSqlInjectionQuery() or
   TTaintedPathQuery() or
-  TXssQuery()
+  TXssQuery() or
+  TStoredXssQuery() or
+  TXssThroughDomQuery() or
+  TCodeInjectionQuery()
 
 abstract class Query extends TQuery {
   abstract string getName();
@@ -31,3 +34,15 @@ class TaintedPathQuery extends Query, TTaintedPathQuery {
 class XssQuery extends Query, TXssQuery {
   override string getName() { result = "Xss" }
 }
+
+class StoredXssQuery extends Query, TStoredXssQuery {
+  override string getName() { result = "StoredXss" }
+}
+
+class XssThroughDomQuery extends Query, TXssThroughDomQuery {
+  override string getName() { result = "XssThroughDom" }
+}
+
+class CodeInjectionQuery extends Query, TCodeInjectionQuery {
+  override string getName() { result = "CodeInjection" }
+}