ERB transformation rework (#20)

* bug fixes

* Reworked ERB transformation to perform a series of appending calls to a buffer'

* rename buffers to prefix with joern__ to distinguish variables

* Remove prefix and postfix tag calls. Add call to original function before lambda lowering

* code cleanup

* add unit test showing lowered ERB code

Author: andreimid

lib/ruby_ast_gen.rb

@@ -124,7 +124,6 @@ def self.parse_file(file_path, relative_input_path)
         file_content = File.read(file_path)
         get_erb_content(file_content)
       end
-    puts code
     buffer = Parser::Source::Buffer.new(file_path)
     buffer.source = code
     parser = Parser::CurrentRuby.new

lib/ruby_ast_gen/erb_to_ruby_transformer.rb

@@ -4,25 +4,27 @@
 class ErbToRubyTransformer
   def initialize
     @parser = Temple::ERB::Parser.new
-    @indent_level = 0
-    @current_line = []
     @in_control_block = false
-    @output_tmp_var = "tmp0"
-    @is_first_output = true
+    @output_tmp_var = "joern__buffer"
+    @in_do_block = false
+    @inner_buffer = "joern__inner_buffer"
+    @current_counter = 0
+    @current_lambda_vars = ""
     @output = []
-    @no_control_struct = true
-    @open_heredoc = false
+    @static_buff = []
   end
 
   def transform(input)
     ast = @parser.call(input)
-    content = "#{@output_tmp_var} = \"\" \n#{visit(ast)}"
-    if @in_control_block
+    @output << "#{@output_tmp_var} = \"\""
+    visit(ast)
+    @output << "return #{@output_tmp_var}"
+
+    if @in_control_block || @in_do_block
       raise ::StandardError, "Invalid ERB Syntax"
     end
     <<~RUBY
-      #{content}
-      return #{@output_tmp_var}
+      #{@output.join("\n")}
     RUBY
   end
 
@@ -32,77 +34,104 @@ def visit(node)
     case node.first
     when :multi
       node[1..-1].each do |child|
-        transformed = visit(child)
-        unless transformed.strip.empty?
-          if @is_first_output
-            @open_heredoc = true
-            @current_line << "#{@output_tmp_var} += <<-HEREDOC\n"
-            @is_first_output = false
-          end
-          @current_line << transformed
-        end
-      end
-
-      if @open_heredoc
-        @current_line << "\nHEREDOC\n"
-        @open_heredoc = false
+        visit(child)
       end
-
-      flush_current_line(@output) unless @current_line.empty?
-      @output.join("\n")
     when :static
-      "#{node[1].to_s}"
+      unless node[1].to_s != nil && node[1].to_s.strip.empty?
+        @static_buff << "\"#{node[1].to_s.gsub('"', '\"').strip}\""
+      end
     when :dynamic
-      "#{node[1].to_s}"
+      unless node[1].to_s != nil && node[1].to_s.strip.empty?
+        @output << "\"#{node[1].to_s.gsub('"', '\"')}\""
+      end
     when :escape
+      unless @static_buff.empty?
+        buffer_to_use = if @in_do_block then "#{@inner_buffer}" else "#{@output_tmp_var}" end
+        @output << "#{buffer_to_use} << \"#{@static_buff.join('\n').gsub(/(?<!\\)"/, '')}\""
+        @static_buff = [] # clear static buffer
+      end
+
       escape_enabled = node[1]
       inner_node = node[2]
       code = inner_node[1].to_s.strip
-      template_call = if escape_enabled then "joern__template_out_raw" else "joern__template_out_escape" end
-      "\#{#{template_call}(#{code})}"
+
+      # Do block with variable found, lower
+      if is_do_block(code)
+        lower_do_block(code)
+      elsif @in_do_block
+        template_call = if escape_enabled then "joern__template_out_raw" else "joern__template_out_escape" end
+        @output << "#{@inner_buffer} << #{template_call}(#{code})"
+      else
+        template_call = if escape_enabled then "joern__template_out_raw" else "joern__template_out_escape" end
+        @output << "#{@output_tmp_var} << #{template_call}(#{code})"
+      end
     when :code
+      unless @static_buff.empty?
+        buffer_to_use = if @in_do_block then "#{@inner_buffer}" else "buffer" end
+        @output << "#{buffer_to_use} << \"#{@static_buff.join('\n').gsub(/(?<!\\)"/, '')}\""
+        @static_buff = [] # clear static buffer
+      end
       code = node[1].to_s.strip
       # Using this to determine if we should throw a StandardError for "invalid" ERB
       if is_control_struct_start(code)
         @in_control_block = true
+        @output << code
       elsif code.start_with?("end")
-        @in_control_block = false
-      end
-
-      if @open_heredoc
-        @open_heredoc = false
-        @current_line << "\nHEREDOC"
+        if @in_do_block
+          @in_do_block = false
+          @output << "#{@inner_buffer}"
+          @output << "end"
+          @output << "#{@output_tmp_var} << #{current_lambda}.call(#{@current_lambda_vars})"
+        else
+          @in_control_block = false
+          @output << "end"
+        end
+      else
+        if is_do_block(code)
+          lower_do_block(code)
+        end
       end
-
-      @current_line << "\n#{node[1].to_s.strip}\n"
-      @is_first_output = true
-      ""
     when :newline
-      ""
     else
       RubyAstGen::Logger::debug("Invalid node type: #{node}")
-      ""
     end
   end
 
-  def indent
-    "  " * @indent_level
+
+  def is_control_struct_start(line)
+    line.start_with?('if', 'unless', 'elsif', 'else', /@?\w+\.each\sdo/)
   end
 
-  def flush_current_line(output)
-    unless @current_line.empty?
-      line = @current_line.join.rstrip
-      output << line unless line.empty?
-      @current_line.clear
-    end
+  def lambda_incrementor()
+    new_lambda = "rails_lambda_#{@current_counter}"
+    @current_counter += 1
+    new_lambda
   end
 
-  def is_control_struct_start(line)
-    line.start_with?('if', 'unless', 'elsif', 'else', /@?\w+\.each\sdo/)
+  def current_lambda()
+    "rails_lambda_#{@current_counter-1}"
+  end
+
+  def lower_do_block(code)
+    if (code_match = code.match(/do\s*(?:\|([^|]*)\|)?/))
+      @current_lambda_vars = code_match[1]
+      before_do, _ = code.split(/\bdo\b/)
+      unless before_do.nil?
+        method_call = before_do.strip
+        call_name, rest = method_call.split(' ', 2)
+        if rest != nil && !rest.start_with?('(') && !rest.end_with?(')')
+          method_call = "#{call_name}(#{rest})"
+        end
+        @output << "#{@output_tmp_var} << #{method_call}"
+      end
+      @in_do_block = true
+      @output << "#{lambda_incrementor} = lambda do |#{@current_lambda_vars}|"
+      @output << "#{@inner_buffer} = \"\""
+    end
   end
 
-  def is_control_struct(line)
-    is_control_struct_start(line) || line.start_with?('end')
+  def is_do_block(code)
+    code.match(/do\s*(?:\|([^|]*)\|)?/)
   end
 end
 

spec/ruby_ast_gen_spec.rb

@@ -103,7 +103,7 @@ def foo(a, bar: "default")
   end
 
   it "should parse ERB structure with no ruby expressions" do
-    code(<<-CODE)
+    erb_code(<<-CODE)
 app_name: <%= ENV['APP_NAME'] %>
 version: <%= ENV['APP_VERSION'] %>
 
@@ -116,7 +116,7 @@ def foo(a, bar: "default")
   end
 
   it "should parse ERB structure with expressions" do
-    code(<<-CODE)
+    erb_code(<<-CODE)
 app_name: <%= ENV['APP_NAME'] %>
 version: <%= ENV['APP_VERSION'] %>
 
@@ -135,7 +135,7 @@ def foo(a, bar: "default")
   end
 
   it "should still return some AST even if the ERB is invalid" do
-    code(<<-CODE)
+    erb_code(<<-CODE)
 app_name: <%= ENV['APP_NAME'] %>
 version: <%= ENV['APP_VERSION'] %>
 
@@ -151,4 +151,27 @@ def foo(a, bar: "default")
     ast = RubyAstGen::parse_file(temp_erb_file.path, temp_name)
     expect(ast).not_to be_nil
   end
+
+  it "should lower ERB code" do
+    erb_code(<<-CODE)
+      <%= form_with url: some_url do |form| %>
+        <%= form.text_field :name %>
+      <% end %>
+    CODE
+
+    file_content = File.read(temp_erb_file.path)
+    code = RubyAstGen::get_erb_content(file_content)
+    expected = <<-HEREDOC
+joern__buffer = ""
+joern__buffer << form_with(url: some_url)
+rails_lambda_0 = lambda do |form|
+joern__inner_buffer = ""
+joern__inner_buffer << joern__template_out_escape(form.text_field :name) 
+joern_inner_buffer
+end
+joern__buffer << rails_lambda_0.call(form)
+return joern__buffer
+    HEREDOC
+    expect(code).equal?(expected)
+  end
 end