From cb65768745432f37de46aabd677b00677046288d Mon Sep 17 00:00:00 2001
From: Tim Sharpe <tim@bombasticmonkey.com>
Date: Wed, 25 Sep 2019 13:11:32 +1000
Subject: [PATCH] Include hash/array references when enclosing variables

---
 .../check_strings/variables_not_enclosed.rb   | 55 +++++++++++++++++++
 .../variables_not_enclosed_spec.rb            | 16 ++++++
 2 files changed, 71 insertions(+)

diff --git a/lib/puppet-lint/plugins/check_strings/variables_not_enclosed.rb b/lib/puppet-lint/plugins/check_strings/variables_not_enclosed.rb
index 1f852abd..4be7d0e5 100644
--- a/lib/puppet-lint/plugins/check_strings/variables_not_enclosed.rb
+++ b/lib/puppet-lint/plugins/check_strings/variables_not_enclosed.rb
@@ -1,9 +1,19 @@
+require 'set'
+require 'strscan'
+
 # Public: Check the manifest tokens for any variables in a string that have
 # not been enclosed by braces ({}) and record a warning for each instance
 # found.
 #
 # https://puppet.com/docs/puppet/latest/style_guide.html#quoting
 PuppetLint.new_check(:variables_not_enclosed) do
+  STRING_TOKEN_TYPES = Set[
+    :DQMID,
+    :DQPOST,
+    :HEREDOC_MID,
+    :HEREDOC_POST,
+  ]
+
   def check
     tokens.select { |r|
       r.type == :UNENC_VARIABLE
@@ -18,7 +28,52 @@ def check
     end
   end
 
+  def hash_or_array_ref?(token)
+    token.next_token &&
+      STRING_TOKEN_TYPES.include?(token.next_token.type) &&
+      token.next_token.value.start_with?('[')
+  end
+
+  def extract_hash_or_array_ref(token)
+    scanner = StringScanner.new(token.value)
+
+    brack_depth = 0
+    result = { :ref => '' }
+
+    until scanner.eos?
+      result[:ref] += scanner.getch
+
+      # Pass a length of 1 when slicing the last character from the string
+      # to prevent Ruby 1.8 returning a Fixnum instead of a String.
+      case result[:ref][-1, 1]
+      when '['
+        brack_depth += 1
+      when ']'
+        brack_depth -= 1
+      end
+
+      break if brack_depth.zero? && scanner.peek(1) != '['
+    end
+
+    result[:remainder] = scanner.rest
+    result
+  end
+
   def fix(problem)
     problem[:token].type = :VARIABLE
+
+    return unless hash_or_array_ref?(problem[:token])
+
+    string_token = problem[:token].next_token
+    tokens_index = tokens.index(string_token)
+
+    hash_or_array_ref = extract_hash_or_array_ref(string_token)
+
+    ref_tokens = PuppetLint::Lexer.new.tokenise(hash_or_array_ref[:ref])
+    ref_tokens.each_with_index do |token, i|
+      add_token(tokens_index + i, token)
+    end
+
+    string_token.value = hash_or_array_ref[:remainder]
   end
 end
diff --git a/spec/puppet-lint/plugins/check_strings/variables_not_enclosed_spec.rb b/spec/puppet-lint/plugins/check_strings/variables_not_enclosed_spec.rb
index 57c4f547..633b2864 100644
--- a/spec/puppet-lint/plugins/check_strings/variables_not_enclosed_spec.rb
+++ b/spec/puppet-lint/plugins/check_strings/variables_not_enclosed_spec.rb
@@ -86,5 +86,21 @@
         expect(manifest).to eq('"${foo}-${bar}"')
       end
     end
+
+    context 'variable with a hash or array reference not enclosed' do
+      let(:code) { %("$foo['bar'][2]something") }
+
+      it 'should only detect a single problem' do
+        expect(problems).to have(1).problem
+      end
+
+      it 'should fix the manifest' do
+        expect(problems).to contain_fixed(msg).on_line(1).in_column(2)
+      end
+
+      it 'should enclose the variable with the references' do
+        expect(manifest).to eq(%("${foo['bar'][2]}something"))
+      end
+    end
   end
 end