Fix if and unless statements in <%= %> tags (#28)

* fixed if and unless stmts in escape tags. Fixed formatting in erb transformer file

* fixed if and unless stmts in escape tags. Fixed formatting in erb transformer file

Author: AndreiDreyer

lib/ruby_ast_gen/erb_to_ruby_transformer.rb

@@ -30,6 +30,7 @@ def transform(input)
   end
 
   private
+
   def visit(node)
     return "" unless node.is_a?(Array)
     case node.first
@@ -59,29 +60,62 @@ def visit(node)
         if ast.is_a?(::Parser::AST::Node)
           case ast.type
           when :if
-            if code.include?(" if ")
-              call, if_cond = code.split(" if ")
+            if code.strip.start_with?("if") || code.strip.start_with?("unless")
+              template_call = if escape_enabled
+                                "joern__template_out_raw"
+                              else
+                                "joern__template_out_escape"
+                              end
+
+              if_cond = extract_code_snippet(ast.children[0].location, code)
+              if_body = extract_code_snippet(ast.children[1].location, code) if ast.children[1]
+              else_body = extract_code_snippet(ast.children[2].location, code) if ast.children[2]
+
               @output << "if #{if_cond}"
+              @output << "#{template_call}(#{if_body})" if if_body
+              @output << "else" if else_body
+              @output << "#{@output_tmp_var} << #{template_call}(#{else_body})" if else_body
             else
-              call, if_cond = code.split(" unless ")
-              @output << "unless #{if_cond}"
+              if code.include?(" if ")
+                call, if_cond = code.split(" if ")
+                @output << "if #{if_cond}"
+              else
+                call, if_cond = code.split(" unless ")
+                @output << "unless #{if_cond}"
+              end
+              template_call = if escape_enabled
+                                "joern__template_out_raw"
+                              else
+                                "joern__template_out_escape"
+                              end
+              @output << "#{@output_tmp_var} << #{template_call}(#{call})"
             end
-            template_call = if escape_enabled then "joern__template_out_raw" else "joern__template_out_escape" end
-            @output << "#{@output_tmp_var} << #{template_call}(#{call})"
             @output << "end"
           else
-            template_call = if escape_enabled then "joern__template_out_raw" else "joern__template_out_escape" end
+            template_call = if escape_enabled then
+                              "joern__template_out_raw"
+                            else
+                              "joern__template_out_escape"
+                            end
             @output << "#{@inner_buffer} << #{template_call}(#{code})"
           end
         end
       elsif is_do_block(code)
         # Do block with variable found, lower
         lower_do_block(code)
       elsif @in_do_block
-        template_call = if escape_enabled then "joern__template_out_raw" else "joern__template_out_escape" end
+        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
+        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
@@ -92,7 +126,7 @@ def visit(node)
       if is_control_struct_start(stripped_code)
         @in_control_block = true
         @output << stripped_code
-      elsif stripped_code.start_with?("end")
+      elsif stripped_code == "end"
         if @in_do_block
           @in_do_block = false
           @output << "#{@inner_buffer}"
@@ -102,10 +136,10 @@ def visit(node)
           @in_control_block = false
           @output << "end"
         end
+      elsif is_do_block(code)
+        lower_do_block(code)
       else
-        if is_do_block(code)
-          lower_do_block(code)
-        end
+        @output << code
       end
     when :newline
     else
@@ -124,12 +158,16 @@ def lambda_incrementor()
   end
 
   def current_lambda()
-    "rails_lambda_#{@current_counter-1}"
+    "rails_lambda_#{@current_counter - 1}"
   end
 
   def flush_static_block()
     unless @static_buff.empty?
-      buffer_to_use = if @in_do_block then @inner_buffer else @output_tmp_var end
+      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
@@ -162,6 +200,14 @@ def lower_do_block(code)
   def is_do_block(code)
     code.match(/do\s*(?:\|([^|]*)\|)?/)
   end
+
+  def extract_code_snippet(location, source_code)
+    return nil unless location
+    range = location.expression || location
+    return nil unless range.is_a?(Parser::Source::Range)
+    snippet = source_code[range.begin_pos...range.end_pos]
+    snippet.strip
+  end
 end