Ruby: add flow summary for Enumerable#pick

Author: aibaars

ruby/ql/lib/codeql/ruby/frameworks/ActiveSupport.qll

@@ -308,7 +308,65 @@ module ActiveSupport {
           preservesValue = true
         }
       }
-      // TODO: pick, pluck (they require Hash dataflow)
+
+      private string getKeyArgument(MethodCall mc, int i) {
+        mc.getMethodName() = ["pick", "pluck"] and
+        result = DataFlow::Content::getKnownElementIndex(mc.getArgument(i)).serialize()
+      }
+
+      private class PickSingleSummary extends SummarizedCallable {
+        private MethodCall mc;
+        private string key;
+
+        PickSingleSummary() {
+          key = getKeyArgument(mc, 0) and
+          this = "Enumerable.pick(" + key + ")" and
+          mc.getMethodName() = "pick" and
+          mc.getNumberOfArguments() = 1
+        }
+
+        override MethodCall getACall() { result = mc }
+
+        override predicate propagatesFlowExt(string input, string output, boolean preservesValue) {
+          input = "Argument[self].Element[0].Element[" + key + "!]" and
+          output = "ReturnValue" and
+          preservesValue = true
+        }
+      }
+
+      private class PickMultipleSummary extends SummarizedCallable {
+        private MethodCall mc;
+
+        PickMultipleSummary() {
+          mc.getMethodName() = "pick" and
+          mc.getNumberOfArguments() > 1 and
+          exists(int maxKey |
+            maxKey = max(int j | exists(getKeyArgument(mc, j))) and
+            this =
+              "Enumerable.pick(" +
+                concat(int i, string key |
+                  key = getKeyArgument(mc, i)
+                  or
+                  key = "_" and
+                  not exists(getKeyArgument(mc, i)) and
+                  i in [0 .. maxKey]
+                |
+                  key, "," order by i
+                ) + ")"
+          )
+        }
+
+        override MethodCall getACall() { result = mc }
+
+        override predicate propagatesFlowExt(string input, string output, boolean preservesValue) {
+          exists(string s, int i |
+            s = getKeyArgument(mc, i) and
+            input = "Argument[self].Element[0].Element[" + s + "!]" and
+            output = "ReturnValue.Element[" + i + "!]"
+          ) and
+          preservesValue = true
+        }
+      }
     }
   }
 

ruby/ql/test/library-tests/frameworks/active_support/ActiveSupportDataFlow.expected

@@ -320,6 +320,38 @@ edges
 | hash_extensions.rb:91:10:91:10 | j [element] :  | hash_extensions.rb:91:10:91:16 | ...[...] |
 | hash_extensions.rb:92:10:92:10 | j [element] :  | hash_extensions.rb:92:10:92:16 | ...[...] |
 | hash_extensions.rb:92:10:92:10 | j [element] :  | hash_extensions.rb:92:10:92:16 | ...[...] |
+| hash_extensions.rb:98:21:98:31 | call to source :  | hash_extensions.rb:99:10:99:15 | values [element 0, element :id] :  |
+| hash_extensions.rb:98:21:98:31 | call to source :  | hash_extensions.rb:99:10:99:15 | values [element 0, element :id] :  |
+| hash_extensions.rb:98:21:98:31 | call to source :  | hash_extensions.rb:101:10:101:15 | values [element 0, element :id] :  |
+| hash_extensions.rb:98:21:98:31 | call to source :  | hash_extensions.rb:101:10:101:15 | values [element 0, element :id] :  |
+| hash_extensions.rb:98:21:98:31 | call to source :  | hash_extensions.rb:104:10:104:15 | values [element 0, element :id] :  |
+| hash_extensions.rb:98:21:98:31 | call to source :  | hash_extensions.rb:104:10:104:15 | values [element 0, element :id] :  |
+| hash_extensions.rb:98:40:98:54 | call to source :  | hash_extensions.rb:100:10:100:15 | values [element 0, element :name] :  |
+| hash_extensions.rb:98:40:98:54 | call to source :  | hash_extensions.rb:100:10:100:15 | values [element 0, element :name] :  |
+| hash_extensions.rb:98:40:98:54 | call to source :  | hash_extensions.rb:102:10:102:15 | values [element 0, element :name] :  |
+| hash_extensions.rb:98:40:98:54 | call to source :  | hash_extensions.rb:102:10:102:15 | values [element 0, element :name] :  |
+| hash_extensions.rb:98:40:98:54 | call to source :  | hash_extensions.rb:103:10:103:15 | values [element 0, element :name] :  |
+| hash_extensions.rb:98:40:98:54 | call to source :  | hash_extensions.rb:103:10:103:15 | values [element 0, element :name] :  |
+| hash_extensions.rb:99:10:99:15 | values [element 0, element :id] :  | hash_extensions.rb:99:10:99:25 | call to pick |
+| hash_extensions.rb:99:10:99:15 | values [element 0, element :id] :  | hash_extensions.rb:99:10:99:25 | call to pick |
+| hash_extensions.rb:100:10:100:15 | values [element 0, element :name] :  | hash_extensions.rb:100:10:100:27 | call to pick |
+| hash_extensions.rb:100:10:100:15 | values [element 0, element :name] :  | hash_extensions.rb:100:10:100:27 | call to pick |
+| hash_extensions.rb:101:10:101:15 | values [element 0, element :id] :  | hash_extensions.rb:101:10:101:32 | call to pick [element 0] :  |
+| hash_extensions.rb:101:10:101:15 | values [element 0, element :id] :  | hash_extensions.rb:101:10:101:32 | call to pick [element 0] :  |
+| hash_extensions.rb:101:10:101:32 | call to pick [element 0] :  | hash_extensions.rb:101:10:101:35 | ...[...] |
+| hash_extensions.rb:101:10:101:32 | call to pick [element 0] :  | hash_extensions.rb:101:10:101:35 | ...[...] |
+| hash_extensions.rb:102:10:102:15 | values [element 0, element :name] :  | hash_extensions.rb:102:10:102:32 | call to pick [element 1] :  |
+| hash_extensions.rb:102:10:102:15 | values [element 0, element :name] :  | hash_extensions.rb:102:10:102:32 | call to pick [element 1] :  |
+| hash_extensions.rb:102:10:102:32 | call to pick [element 1] :  | hash_extensions.rb:102:10:102:35 | ...[...] |
+| hash_extensions.rb:102:10:102:32 | call to pick [element 1] :  | hash_extensions.rb:102:10:102:35 | ...[...] |
+| hash_extensions.rb:103:10:103:15 | values [element 0, element :name] :  | hash_extensions.rb:103:10:103:32 | call to pick [element 0] :  |
+| hash_extensions.rb:103:10:103:15 | values [element 0, element :name] :  | hash_extensions.rb:103:10:103:32 | call to pick [element 0] :  |
+| hash_extensions.rb:103:10:103:32 | call to pick [element 0] :  | hash_extensions.rb:103:10:103:35 | ...[...] |
+| hash_extensions.rb:103:10:103:32 | call to pick [element 0] :  | hash_extensions.rb:103:10:103:35 | ...[...] |
+| hash_extensions.rb:104:10:104:15 | values [element 0, element :id] :  | hash_extensions.rb:104:10:104:32 | call to pick [element 1] :  |
+| hash_extensions.rb:104:10:104:15 | values [element 0, element :id] :  | hash_extensions.rb:104:10:104:32 | call to pick [element 1] :  |
+| hash_extensions.rb:104:10:104:32 | call to pick [element 1] :  | hash_extensions.rb:104:10:104:35 | ...[...] |
+| hash_extensions.rb:104:10:104:32 | call to pick [element 1] :  | hash_extensions.rb:104:10:104:35 | ...[...] |
 nodes
 | active_support.rb:10:9:10:18 | call to source :  | semmle.label | call to source :  |
 | active_support.rb:11:10:11:10 | x :  | semmle.label | x :  |
@@ -722,6 +754,42 @@ nodes
 | hash_extensions.rb:92:10:92:10 | j [element] :  | semmle.label | j [element] :  |
 | hash_extensions.rb:92:10:92:16 | ...[...] | semmle.label | ...[...] |
 | hash_extensions.rb:92:10:92:16 | ...[...] | semmle.label | ...[...] |
+| hash_extensions.rb:98:21:98:31 | call to source :  | semmle.label | call to source :  |
+| hash_extensions.rb:98:21:98:31 | call to source :  | semmle.label | call to source :  |
+| hash_extensions.rb:98:40:98:54 | call to source :  | semmle.label | call to source :  |
+| hash_extensions.rb:98:40:98:54 | call to source :  | semmle.label | call to source :  |
+| hash_extensions.rb:99:10:99:15 | values [element 0, element :id] :  | semmle.label | values [element 0, element :id] :  |
+| hash_extensions.rb:99:10:99:15 | values [element 0, element :id] :  | semmle.label | values [element 0, element :id] :  |
+| hash_extensions.rb:99:10:99:25 | call to pick | semmle.label | call to pick |
+| hash_extensions.rb:99:10:99:25 | call to pick | semmle.label | call to pick |
+| hash_extensions.rb:100:10:100:15 | values [element 0, element :name] :  | semmle.label | values [element 0, element :name] :  |
+| hash_extensions.rb:100:10:100:15 | values [element 0, element :name] :  | semmle.label | values [element 0, element :name] :  |
+| hash_extensions.rb:100:10:100:27 | call to pick | semmle.label | call to pick |
+| hash_extensions.rb:100:10:100:27 | call to pick | semmle.label | call to pick |
+| hash_extensions.rb:101:10:101:15 | values [element 0, element :id] :  | semmle.label | values [element 0, element :id] :  |
+| hash_extensions.rb:101:10:101:15 | values [element 0, element :id] :  | semmle.label | values [element 0, element :id] :  |
+| hash_extensions.rb:101:10:101:32 | call to pick [element 0] :  | semmle.label | call to pick [element 0] :  |
+| hash_extensions.rb:101:10:101:32 | call to pick [element 0] :  | semmle.label | call to pick [element 0] :  |
+| hash_extensions.rb:101:10:101:35 | ...[...] | semmle.label | ...[...] |
+| hash_extensions.rb:101:10:101:35 | ...[...] | semmle.label | ...[...] |
+| hash_extensions.rb:102:10:102:15 | values [element 0, element :name] :  | semmle.label | values [element 0, element :name] :  |
+| hash_extensions.rb:102:10:102:15 | values [element 0, element :name] :  | semmle.label | values [element 0, element :name] :  |
+| hash_extensions.rb:102:10:102:32 | call to pick [element 1] :  | semmle.label | call to pick [element 1] :  |
+| hash_extensions.rb:102:10:102:32 | call to pick [element 1] :  | semmle.label | call to pick [element 1] :  |
+| hash_extensions.rb:102:10:102:35 | ...[...] | semmle.label | ...[...] |
+| hash_extensions.rb:102:10:102:35 | ...[...] | semmle.label | ...[...] |
+| hash_extensions.rb:103:10:103:15 | values [element 0, element :name] :  | semmle.label | values [element 0, element :name] :  |
+| hash_extensions.rb:103:10:103:15 | values [element 0, element :name] :  | semmle.label | values [element 0, element :name] :  |
+| hash_extensions.rb:103:10:103:32 | call to pick [element 0] :  | semmle.label | call to pick [element 0] :  |
+| hash_extensions.rb:103:10:103:32 | call to pick [element 0] :  | semmle.label | call to pick [element 0] :  |
+| hash_extensions.rb:103:10:103:35 | ...[...] | semmle.label | ...[...] |
+| hash_extensions.rb:103:10:103:35 | ...[...] | semmle.label | ...[...] |
+| hash_extensions.rb:104:10:104:15 | values [element 0, element :id] :  | semmle.label | values [element 0, element :id] :  |
+| hash_extensions.rb:104:10:104:15 | values [element 0, element :id] :  | semmle.label | values [element 0, element :id] :  |
+| hash_extensions.rb:104:10:104:32 | call to pick [element 1] :  | semmle.label | call to pick [element 1] :  |
+| hash_extensions.rb:104:10:104:32 | call to pick [element 1] :  | semmle.label | call to pick [element 1] :  |
+| hash_extensions.rb:104:10:104:35 | ...[...] | semmle.label | ...[...] |
+| hash_extensions.rb:104:10:104:35 | ...[...] | semmle.label | ...[...] |
 subpaths
 #select
 | active_support.rb:182:10:182:13 | ...[...] | active_support.rb:180:10:180:17 | call to source :  | active_support.rb:182:10:182:13 | ...[...] | $@ | active_support.rb:180:10:180:17 | call to source :  | call to source :  |
@@ -766,3 +834,9 @@ subpaths
 | hash_extensions.rb:87:10:87:16 | ...[...] | hash_extensions.rb:83:9:83:19 | call to source :  | hash_extensions.rb:87:10:87:16 | ...[...] | $@ | hash_extensions.rb:83:9:83:19 | call to source :  | call to source :  |
 | hash_extensions.rb:91:10:91:16 | ...[...] | hash_extensions.rb:89:27:89:37 | call to source :  | hash_extensions.rb:91:10:91:16 | ...[...] | $@ | hash_extensions.rb:89:27:89:37 | call to source :  | call to source :  |
 | hash_extensions.rb:92:10:92:16 | ...[...] | hash_extensions.rb:89:27:89:37 | call to source :  | hash_extensions.rb:92:10:92:16 | ...[...] | $@ | hash_extensions.rb:89:27:89:37 | call to source :  | call to source :  |
+| hash_extensions.rb:99:10:99:25 | call to pick | hash_extensions.rb:98:21:98:31 | call to source :  | hash_extensions.rb:99:10:99:25 | call to pick | $@ | hash_extensions.rb:98:21:98:31 | call to source :  | call to source :  |
+| hash_extensions.rb:100:10:100:27 | call to pick | hash_extensions.rb:98:40:98:54 | call to source :  | hash_extensions.rb:100:10:100:27 | call to pick | $@ | hash_extensions.rb:98:40:98:54 | call to source :  | call to source :  |
+| hash_extensions.rb:101:10:101:35 | ...[...] | hash_extensions.rb:98:21:98:31 | call to source :  | hash_extensions.rb:101:10:101:35 | ...[...] | $@ | hash_extensions.rb:98:21:98:31 | call to source :  | call to source :  |
+| hash_extensions.rb:102:10:102:35 | ...[...] | hash_extensions.rb:98:40:98:54 | call to source :  | hash_extensions.rb:102:10:102:35 | ...[...] | $@ | hash_extensions.rb:98:40:98:54 | call to source :  | call to source :  |
+| hash_extensions.rb:103:10:103:35 | ...[...] | hash_extensions.rb:98:40:98:54 | call to source :  | hash_extensions.rb:103:10:103:35 | ...[...] | $@ | hash_extensions.rb:98:40:98:54 | call to source :  | call to source :  |
+| hash_extensions.rb:104:10:104:35 | ...[...] | hash_extensions.rb:98:21:98:31 | call to source :  | hash_extensions.rb:104:10:104:35 | ...[...] | $@ | hash_extensions.rb:98:21:98:31 | call to source :  | call to source :  |

ruby/ql/test/library-tests/frameworks/active_support/hash_extensions.rb

@@ -93,3 +93,15 @@ def m_index_with
 end
 
 m_index_with()
+
+def m_pick
+    values = [{ id: source("1"), name: source("David") }, { id: source("2"), name: source("Rafael") }]
+    sink(values.pick(:id)) # $ hasValueFlow=1
+    sink(values.pick(:name)) # $ hasValueFlow=David
+    sink(values.pick(:id, :name)[0]) # $ hasValueFlow=1
+    sink(values.pick(:id, :name)[1]) # $ hasValueFlow=David
+    sink(values.pick(:name, :id)[0]) # $ hasValueFlow=David
+    sink(values.pick(:name, :id)[1]) # $ hasValueFlow=1
+end
+
+m_pick()