Ruby: Use a newtype instead of DataFlow::FlowState for code-injection

ruby/ql/lib/codeql/ruby/security/CodeInjectionCustomizations.qll

@@ -13,27 +13,74 @@ private import codeql.ruby.dataflow.BarrierGuards
 module CodeInjection {
   /** Flow states used to distinguish whether an attacker controls the entire string. */
   module FlowState {
-    /** Flow state used for normal tainted data, where an attacker might only control a substring. */
-    DataFlow::FlowState substring() { result = "substring" }
+    /**
+     * Flow state used for normal tainted data, where an attacker might only control a substring.
+     * DEPRECATED: Use `Full()`
+     */
+    deprecated DataFlow::FlowState substring() { result = "substring" }
 
-    /** Flow state used for data that is entirely controlled by the attacker. */
-    DataFlow::FlowState full() { result = "full" }
+    /**
+     * Flow state used for data that is entirely controlled by the attacker.
+     * DEPRECATED: Use `Full()`
+     */
+    deprecated DataFlow::FlowState full() { result = "full" }
+
+    private newtype TState =
+      TFull() or
+      TSubString()
+
+    /** Flow states used to distinguish whether an attacker controls the entire string. */
+    class State extends TState {
+      string toString() {
+        this = TSubString() and result = "substring"
+        or
+        this = TFull() and result = "full"
+      }
+    }
+
+    /**
+     * Flow state used for normal tainted data, where an attacker might only control a substring.
+     */
+    class SubString extends State, TSubString { }
+
+    /**
+     * Flow state used for data that is entirely controlled by the attacker.
+     */
+    class Full extends State, TFull { }
   }
 
   /**
    * A data flow source for "Code injection" vulnerabilities.
    */
   abstract class Source extends DataFlow::Node {
+    /**
+     * Gets a flow state for which this is a source.
+     * DEPRECATED: Use `getAState()`
+     */
+    deprecated DataFlow::FlowState getAFlowState() {
+      result = [FlowState::substring(), FlowState::full()]
+    }
+
     /** Gets a flow state for which this is a source. */
-    DataFlow::FlowState getAFlowState() { result = [FlowState::substring(), FlowState::full()] }
+    FlowState::State getAState() {
+      result instanceof FlowState::SubString or result instanceof FlowState::Full
+    }
   }
 
   /**
    * A data flow sink for "Code injection" vulnerabilities.
    */
   abstract class Sink extends DataFlow::Node {
+    /**
+     * Holds if this sink is safe for an attacker that only controls a substring.
+     * DEPRECATED: Use `getAState()`
+     */
+    deprecated DataFlow::FlowState getAFlowState() {
+      result = [FlowState::substring(), FlowState::full()]
+    }
+
     /** Holds if this sink is safe for an attacker that only controls a substring. */
-    DataFlow::FlowState getAFlowState() { result = [FlowState::substring(), FlowState::full()] }
+    FlowState::State getAState() { any() }
   }
 
   /**
@@ -43,8 +90,15 @@ module CodeInjection {
     /**
      * Gets a flow state for which this is a sanitizer.
      * Sanitizes all states if the result is empty.
+     * DEPRECATED: Use `getAState()`
      */
-    DataFlow::FlowState getAFlowState() { none() }
+    deprecated DataFlow::FlowState getAFlowState() { none() }
+
+    /**
+     * Gets a flow state for which this is a sanitizer.
+     * Sanitizes all states if the result is empty.
+     */
+    FlowState::State getAState() { none() }
   }
 
   /**
@@ -67,12 +121,17 @@ module CodeInjection {
 
     CodeExecutionAsSink() { this = c.getCode() }
 
-    /** Gets a flow state for which this is a sink. */
-    override DataFlow::FlowState getAFlowState() {
+    deprecated override DataFlow::FlowState getAFlowState() {
       if c.runsArbitraryCode()
       then result = [FlowState::substring(), FlowState::full()] // If it runs arbitrary code then it's always vulnerable.
       else result = FlowState::full() // If it "just" loads something, then it's only vulnerable if the attacker controls the entire string.
     }
+
+    override FlowState::State getAState() {
+      if c.runsArbitraryCode()
+      then any() // If it runs arbitrary code then it's always vulnerable.
+      else result instanceof FlowState::Full // If it "just" loads something, then it's only vulnerable if the attacker controls the entire string.
+    }
   }
 
   private import codeql.ruby.AST as Ast
@@ -92,6 +151,8 @@ module CodeInjection {
       )
     }
 
-    override DataFlow::FlowState getAFlowState() { result = FlowState::full() }
+    deprecated override DataFlow::FlowState getAFlowState() { result = FlowState::full() }
+
+    override FlowState::State getAState() { result instanceof FlowState::Full }
   }
 }

ruby/ql/lib/codeql/ruby/security/CodeInjectionQuery.qll

@@ -44,25 +44,21 @@ deprecated class Configuration extends TaintTracking::Configuration {
 }
 
 private module Config implements DataFlow::StateConfigSig {
-  class FlowState = DataFlow::FlowState;
+  class FlowState = FlowState::State;
 
-  predicate isSource(DataFlow::Node source, FlowState state) {
-    state = source.(Source).getAFlowState()
-  }
+  predicate isSource(DataFlow::Node source, FlowState state) { state = source.(Source).getAState() }
 
-  predicate isSink(DataFlow::Node sink, FlowState state) { state = sink.(Sink).getAFlowState() }
+  predicate isSink(DataFlow::Node sink, FlowState state) { state = sink.(Sink).getAState() }
 
   predicate isBarrier(DataFlow::Node node) {
-    node instanceof Sanitizer and not exists(node.(Sanitizer).getAFlowState())
+    node instanceof Sanitizer and not exists(node.(Sanitizer).getAState())
     or
     node instanceof StringConstCompareBarrier
     or
     node instanceof StringConstArrayInclusionCallBarrier
   }
 
-  predicate isBarrier(DataFlow::Node node, DataFlow::FlowState state) {
-    node.(Sanitizer).getAFlowState() = state
-  }
+  predicate isBarrier(DataFlow::Node node, FlowState state) { node.(Sanitizer).getAState() = state }
 }
 
 /**

ruby/ql/src/queries/security/cwe-094/CodeInjection.ql

@@ -29,7 +29,7 @@ where
         otherSink) and
       otherSink.getNode() = sink.getNode()
     |
-      otherSink order by otherSink.getState()
+      otherSink order by otherSink.getState().toString()
     )
 select sink.getNode(), source, sink, "This code execution depends on a $@.", sourceNode,
   "user-provided value"