From 3bfea4d590a4761300df82a7b135e08efcbe91b2 Mon Sep 17 00:00:00 2001
From: Ewoud Kohl van Wijngaarden <ewoud@kohlvanwijngaarden.nl>
Date: Thu, 13 Oct 2022 15:28:18 +0200
Subject: [PATCH] Support safe call

This copies the process_call_code method from ruby2ruby. That also
implements process_safe_call which calls process_call with a parameter.
---
 lib/safemode/parser.rb       | 47 ++++++++++++++++++++++++------------
 test/test_safemode_eval.rb   |  4 +++
 test/test_safemode_parser.rb | 10 ++++++++
 3 files changed, 46 insertions(+), 15 deletions(-)

diff --git a/lib/safemode/parser.rb b/lib/safemode/parser.rb
index e2f5776..ce96a06 100644
--- a/lib/safemode/parser.rb
+++ b/lib/safemode/parser.rb
@@ -26,19 +26,22 @@ def parser=(parser)
       end
     end
 
-    def jail(str, parentheses = false)
-      str = parentheses ? "(#{str})." : "#{str}." if str
+    def jail(str, parentheses = false, safe_call: false)
+      str = if str
+              dot = safe_call ? "&." : "."
+              parentheses ? "(#{str})&#{dot}" : "#{str}#{dot}"
+            end
       "#{str}to_jail"
     end
 
     # split up #process_call. see below ...
-    def process_call(exp)
+    def process_call(exp, safe_call = false)
       exp.shift # remove ":call" symbol
-      receiver = jail process_call_receiver(exp)
+      receiver = jail(process_call_receiver(exp), safe_call: safe_call)
       name = exp.shift
       args = process_call_args(exp)
 
-      process_call_code(receiver, name, args)
+      process_call_code(receiver, name, args, safe_call)
     end
 
     def process_fcall(exp)
@@ -159,25 +162,39 @@ def process_call_args(exp)
         end
         args << processed unless (processed.nil? or processed.empty?)
       end
-      args.empty? ? nil : args.join(", ")
+      args
     end
 
-    def process_call_code(receiver, name, args)
+    def process_call_code(receiver, name, args, safe_call)
       case name
-      when :<=>, :==, "!=".to_sym, :<, :>, :<=, :>=, :-, :+, :*, :/, :%, :<<, :>>, :** then
-        "(#{receiver} #{name} #{args})"
+      when *BINARY then
+        if safe_call
+          "#{receiver}&.#{name}(#{args.join(", ")})"
+        elsif args.length > 1
+          "#{receiver}.#{name}(#{args.join(", ")})"
+        else
+          "(#{receiver} #{name} #{args.join(", ")})"
+        end
       when :[] then
-        "#{receiver}[#{args}]"
+        receiver ||= "self"
+        "#{receiver}[#{args.join(", ")}]"
+      when :[]= then
+        receiver ||= "self"
+        rhs = args.pop
+        "#{receiver}[#{args.join(", ")}] = #{rhs}"
+      when :"!" then
+        "(not #{receiver})"
       when :"-@" then
         "-#{receiver}"
       when :"+@" then
         "+#{receiver}"
       else
-        unless receiver.nil? then
-          "#{receiver}.#{name}#{args ? "(#{args})" : args}"
-        else
-          "#{name}#{args ? "(#{args})" : args}"
-        end
+        args     = nil                    if args.empty?
+        args     = "(#{args.join(", ")})" if args
+        receiver = "#{receiver}."         if receiver and not safe_call
+        receiver = "#{receiver}&."        if receiver and safe_call
+
+        "#{receiver}#{name}#{args}"
       end
     end
 
diff --git a/test/test_safemode_eval.rb b/test/test_safemode_eval.rb
index 051360b..2159f35 100644
--- a/test/test_safemode_eval.rb
+++ b/test/test_safemode_eval.rb
@@ -17,6 +17,10 @@ def test_some_stuff_that_should_work
     end
   end
 
+   def test_safe_navigation_operator
+     assert_equal "1", @box.eval('x = 1; x&.to_s')
+   end
+
   def test_unary_operators_on_instances_of_boolean_vars
     assert @box.eval('not false')
     assert @box.eval('!false')
diff --git a/test/test_safemode_parser.rb b/test/test_safemode_parser.rb
index ff95d6f..d011ffb 100644
--- a/test/test_safemode_parser.rb
+++ b/test/test_safemode_parser.rb
@@ -25,6 +25,16 @@ def test_ternary_should_be_usable_for_erb
     assert_jailed "if true then\n 1\n else\n2\nend", "true ? 1 : 2"
   end
 
+  def test_safe_call_simple
+    assert_jailed '@article&.to_jail&.method', '@article&.method'
+  end
+
+  def test_safe_call_with_complex_args
+    unsafe = "@article&.method_with_kwargs('positional')"
+    jailed = "@article&.to_jail&.method_with_kwargs(\"positional\")"
+    assert_jailed jailed, unsafe
+  end
+
   def test_output_buffer_should_be_assignable
     assert_nothing_raised do
       jail('@output_buffer = ""')