From ba541afae7c5d6f0e2a1da39a5cf60e8212b9ebb Mon Sep 17 00:00:00 2001
From: Justin Collins <justin@presidentbeef.com>
Date: Tue, 13 Jul 2021 12:17:13 -0700
Subject: [PATCH] Use simple method checks from MethInfo

---
 lib/brakeman/tracker/collection.rb  | 27 +++++++--------------------
 lib/brakeman/tracker/method_info.rb |  2 ++
 2 files changed, 9 insertions(+), 20 deletions(-)

diff --git a/lib/brakeman/tracker/collection.rb b/lib/brakeman/tracker/collection.rb
index d6153181e6..0255b1aee1 100644
--- a/lib/brakeman/tracker/collection.rb
+++ b/lib/brakeman/tracker/collection.rb
@@ -50,7 +50,7 @@ def add_option name, exp
 
     def add_method visibility, name, src, file_name
       meth_info = Brakeman::MethodInfo.new(name, src, self, file_name)
-      add_simple_method_maybe meth_info 
+      add_simple_method_maybe meth_info
 
       if src.node_type == :defs
         @class_methods[name] = meth_info
@@ -121,28 +121,15 @@ def get_simple_method_return_value type, name
 
     private
 
-    def add_simple_method_maybe meth_info 
-      src = meth_info.src
-
-      # Simple methods have one (simple) expression in the body and
-      # no arguments
-      if src.formal_args.length == 1 # no args
-        body = src.body
-        if body.length == 1 # single expression in body
-          value = body.first
-
-          if simple_literal? value or
-              all_literals? value, :array or
-              all_literals? value, :hash
-
-            add_simple_method meth_info, value
-          end
-        end
+    def add_simple_method_maybe meth_info
+      if meth_info.very_simple_method?
+        add_simple_method meth_info
       end
     end
 
-    def add_simple_method meth_info, value
-      name = meth_info.name 
+    def add_simple_method meth_info
+      name = meth_info.name
+      value = meth_info.return_value
 
       case meth_info.src.node_type
       when :defn
diff --git a/lib/brakeman/tracker/method_info.rb b/lib/brakeman/tracker/method_info.rb
index e423f40497..31a2b6b4f4 100644
--- a/lib/brakeman/tracker/method_info.rb
+++ b/lib/brakeman/tracker/method_info.rb
@@ -19,6 +19,8 @@ def initialize name, src, owner, file
               else
                 raise "Expected sexp type: #{src.node_type}"
               end
+
+      @simple_method = nil
     end
 
     # To support legacy code that expected a Hash