Merge pull request #8784 from luchihoratiu/PUP-11204

(PUP-11204) Allow `puppet lookup --facts` to receive any filename

Author: GabrielNagy

lib/puppet/application/lookup.rb

@@ -54,11 +54,7 @@ class Puppet::Application::Lookup < Puppet::Application
   end
 
   option('--facts FACT_FILE') do |arg|
-    if %w{.yaml .yml .json}.include?(arg.match(/\.[^.]*$/)[0])
-      options[:fact_file] = arg
-    else
-      raise _("The --fact file only accepts yaml and json files.\n%{run_help}") % { run_help: RUN_HELP }
-    end
+    options[:fact_file] = arg
   end
 
   def app_defaults
@@ -356,14 +352,17 @@ def generate_scope
     fact_file = options[:fact_file]
 
     if fact_file
-      if fact_file.end_with?("json")
-        given_facts = Puppet::Util::Json.load(Puppet::FileSystem.read(fact_file, :encoding => 'utf-8'))
-      else
+      if fact_file.end_with?('.json')
+        given_facts = Puppet::Util::Json.load_file(fact_file)
+      elsif fact_file.end_with?('.yml', '.yaml')
         given_facts = Puppet::Util::Yaml.safe_load_file(fact_file)
+      else
+        given_facts = Puppet::Util::Json.load_file_if_valid(fact_file)
+        given_facts = Puppet::Util::Yaml.safe_load_file_if_valid(fact_file) unless given_facts
       end
 
       unless given_facts.instance_of?(Hash)
-        raise _("Incorrect formatted data in %{fact_file} given via the --facts flag") % { fact_file: fact_file }
+        raise _("Incorrectly formatted data in %{fact_file} given via the --facts flag (only accepts yaml and json files)") % { fact_file: fact_file }
       end
       node.add_extra_facts(given_facts)
     end

lib/puppet/util/json.rb

@@ -26,6 +26,23 @@ def self.build(original_exception, data)
       require 'json'
     end
 
+    # Load the content from a file as JSON if
+    # contents are in valid format. This method does not
+    # raise error but returns `nil` when invalid file is
+    # given.
+    def self.load_file_if_valid(filename, options = {})
+      load_file(filename, options)
+    rescue Puppet::Util::Json::ParseError, ArgumentError, Errno::ENOENT => detail
+      Puppet.debug("Could not retrieve JSON content from '#{filename}': #{detail.message}")
+      nil
+    end
+
+    # Load the content from a file as JSON.
+    def self.load_file(filename, options = {})
+      json = Puppet::FileSystem.read(filename, :encoding => 'utf-8')
+      load(json, options)
+    end
+
     # These methods do similar processing to the fallback implemented by MultiJson
     # when using the built-in JSON backend, to ensure consistent behavior
     # whether or not MultiJson can be loaded.

lib/puppet/util/yaml.rb

@@ -42,6 +42,17 @@ def self.safe_load_file(filename, allowed_classes = [])
     safe_load(yaml, allowed_classes, filename)
   end
 
+  # Safely load the content from a file as YAML if
+  # contents are in valid format. This method does not
+  # raise error but returns `nil` when invalid file is
+  # given.
+  def self.safe_load_file_if_valid(filename, allowed_classes = [])
+    safe_load_file(filename, allowed_classes)
+  rescue YamlLoadError, ArgumentError, Errno::ENOENT => detail
+    Puppet.debug("Could not retrieve YAML content from '#{filename}': #{detail.message}")
+    nil
+  end
+
   # @deprecated Use {#safe_load_file} instead.
   def self.load_file(filename, default_value = false, strip_classes = false)
     Puppet.deprecation_warning(_("Puppet::Util::Yaml.load_file is deprecated. Use safe_load_file instead."))

spec/unit/application/lookup_spec.rb

@@ -518,34 +518,120 @@ def run_lookup(lookup)
       end
 
       it 'receives extra facts in top scope' do
-        file_path = tmpdir('lookup_spec')
-        filename = File.join(file_path, "facts.yaml")
-        File.open(filename, "w+") { |f| f.write(<<-YAML.unindent) }
+        file_path = file_containing('facts.yaml', <<~CONTENT)
           ---
           cx: ' C from facts'
-          YAML
+        CONTENT
 
         lookup.options[:node] = node
-        lookup.options[:fact_file] = filename
+        lookup.options[:fact_file] = file_path
         lookup.options[:render_as] = :s
         allow(lookup.command_line).to receive(:args).and_return(['c'])
         expect(run_lookup(lookup)).to eql("This is C from facts")
       end
 
       it 'receives extra facts in the facts hash' do
-        file_path = tmpdir('lookup_spec')
-        filename = File.join(file_path, "facts.yaml")
-        File.open(filename, "w+") { |f| f.write(<<-YAML.unindent) }
+        file_path = file_containing('facts.yaml', <<~CONTENT)
           ---
           cx: ' G from facts'
-        YAML
+        CONTENT
 
         lookup.options[:node] = node
-        lookup.options[:fact_file] = filename
+        lookup.options[:fact_file] = file_path
         lookup.options[:render_as] = :s
         allow(lookup.command_line).to receive(:args).and_return(['g'])
         expect(run_lookup(lookup)).to eql("This is G from facts in facts hash")
       end
+
+      describe 'when retrieving given facts' do
+        before do
+          lookup.options[:node] = node
+          allow(lookup.command_line).to receive(:args).and_return(['g'])
+        end
+
+        it 'loads files with yaml extension as yaml on first try' do
+          file_path = file_containing('facts.yaml', <<~CONTENT)
+            ---
+            cx: ' G from facts'
+          CONTENT
+          lookup.options[:fact_file] = file_path
+
+          expect(Puppet::Util::Yaml).to receive(:safe_load_file).with(file_path, []).and_call_original
+          run_lookup(lookup)
+        end
+
+        it 'loads files with yml extension as yaml on first try' do
+          file_path = file_containing('facts.yml', <<~CONTENT)
+            ---
+            cx: ' G from facts'
+          CONTENT
+          lookup.options[:fact_file] = file_path
+
+          expect(Puppet::Util::Yaml).to receive(:safe_load_file).with(file_path, []).and_call_original
+          run_lookup(lookup)
+        end
+
+        it 'loads files with json extension as json on first try' do
+          file_path = file_containing('facts.json', <<~CONTENT)
+            {
+              "cx": " G from facts"
+            }
+          CONTENT
+          lookup.options[:fact_file] = file_path
+
+          expect(Puppet::Util::Json).to receive(:load_file).with(file_path, {}).and_call_original
+          run_lookup(lookup)
+        end
+
+        it 'detects file format accordingly even with random file extension' do
+          file_path = file_containing('facts.txt', <<~CONTENT)
+            {
+              "cx": " G from facts"
+            }
+          CONTENT
+          lookup.options[:fact_file] = file_path
+
+          expect(Puppet::Util::Json).to receive(:load_file_if_valid).with(file_path).and_call_original
+          expect(Puppet::Util::Yaml).not_to receive(:safe_load_file_if_valid).with(file_path)
+          run_lookup(lookup)
+        end
+
+        it 'detects file without extension as json due to valid contents' do
+          file_path = file_containing('facts', <<~CONTENT)
+            {
+              "cx": " G from facts"
+            }
+          CONTENT
+          lookup.options[:fact_file] = file_path
+
+          expect(Puppet::Util::Json).to receive(:load_file_if_valid).with(file_path).and_call_original
+          expect(Puppet::Util::Yaml).not_to receive(:safe_load_file_if_valid).with(file_path)
+          run_lookup(lookup)
+        end
+
+        it 'detects file without extension as yaml due to valid contents' do
+          file_path = file_containing('facts', <<~CONTENT)
+            ---
+            cx: ' G from facts'
+          CONTENT
+          lookup.options[:fact_file] = file_path
+
+          expect(Puppet::Util::Json.load_file_if_valid(file_path)).to be_nil
+          expect(Puppet::Util::Yaml).to receive(:safe_load_file_if_valid).with(file_path).and_call_original
+          run_lookup(lookup)
+        end
+
+        it 'raises error due to invalid contents' do
+          file_path = file_containing('facts.yml', <<~CONTENT)
+            INVALID CONTENT
+          CONTENT
+          lookup.options[:fact_file] = file_path
+
+          expect {
+            lookup.run_command
+          }.to raise_error(/Incorrectly formatted data in .+ given via the --facts flag \(only accepts yaml and json files\)/)
+        end
+      end
     end
 
     context 'using a puppet function as data provider' do

spec/unit/util/json_spec.rb

@@ -0,0 +1,126 @@
+# coding: utf-8
+require 'spec_helper'
+require 'puppet/util/json'
+
+describe Puppet::Util::Json do
+  include PuppetSpec::Files
+
+  shared_examples_for 'json file loader' do |load_method|
+    it 'reads a JSON file from disk' do
+      file_path = file_containing('input', JSON.dump({ "my" => "data" }))
+
+      expect(load_method.call(file_path)).to eq({ "my" => "data" })
+    end
+
+    it 'reads JSON as UTF-8' do
+      file_path = file_containing('input', JSON.dump({ "my" => "𠜎" }))
+
+      expect(load_method.call(file_path)).to eq({ "my" => "𠜎" })
+    end
+  end
+
+  context "#load" do
+    it 'raises an error if JSON is invalid' do
+      expect {
+        Puppet::Util::Json.load('{ invalid')
+      }.to raise_error(Puppet::Util::Json::ParseError, /unexpected token at '{ invalid'/)
+    end
+
+    it 'raises an error if the content is empty' do
+      expect {
+        Puppet::Util::Json.load('')
+      }.to raise_error(Puppet::Util::Json::ParseError)
+    end
+
+    it 'loads true' do
+      expect(Puppet::Util::Json.load('true')).to eq(true)
+    end
+
+    it 'loads false' do
+      expect(Puppet::Util::Json.load('false')).to eq(false)
+    end
+
+    it 'loads a numeric' do
+      expect(Puppet::Util::Json.load('42')).to eq(42)
+    end
+
+    it 'loads a string' do
+      expect(Puppet::Util::Json.load('"puppet"')).to eq('puppet')
+    end
+
+    it 'loads an array' do
+      expect(Puppet::Util::Json.load(<<~JSON)).to eq([1, 2])
+        [1, 2]
+      JSON
+    end
+
+    it 'loads a hash' do
+      expect(Puppet::Util::Json.load(<<~JSON)).to eq('a' => 1, 'b' => 2)
+        {
+          "a": 1,
+          "b": 2
+        }
+      JSON
+    end
+  end
+
+  context "load_file_if_valid" do
+    before do
+      Puppet[:log_level] = 'debug'
+    end
+
+    it_should_behave_like 'json file loader', Puppet::Util::Json.method(:load_file_if_valid)
+
+    it 'returns nil when the file is invalid JSON and debug logs about it' do
+      file_path = file_containing('input', '{ invalid')
+      expect(Puppet).to receive(:debug)
+        .with(/Could not retrieve JSON content .+: unexpected token at '{ invalid'/).and_call_original
+
+      expect(Puppet::Util::Json.load_file_if_valid(file_path)).to eql(nil)
+    end
+
+    it 'returns nil when the filename is illegal and debug logs about it' do
+      expect(Puppet).to receive(:debug)
+        .with(/Could not retrieve JSON content .+: pathname contains null byte/).and_call_original
+
+      expect(Puppet::Util::Json.load_file_if_valid("not\0allowed")).to eql(nil)
+    end
+
+    it 'returns nil when the file does not exist and debug logs about it' do
+      expect(Puppet).to receive(:debug)
+        .with(/Could not retrieve JSON content .+: No such file or directory/).and_call_original
+
+      expect(Puppet::Util::Json.load_file_if_valid('does/not/exist.json')).to eql(nil)
+    end
+  end
+
+  context '#load_file' do
+    it_should_behave_like 'json file loader', Puppet::Util::Json.method(:load_file)
+
+    it 'raises an error when the file is invalid JSON' do
+      file_path = file_containing('input', '{ invalid')
+
+      expect {
+        Puppet::Util::Json.load_file(file_path)
+      }.to raise_error(Puppet::Util::Json::ParseError, /unexpected token at '{ invalid'/)
+    end
+
+    it 'raises an error when the filename is illegal' do
+      expect {
+        Puppet::Util::Json.load_file("not\0allowed")
+      }.to raise_error(ArgumentError, /null byte/)
+    end
+
+    it 'raises an error when the file does not exist' do
+      expect {
+        Puppet::Util::Json.load_file('does/not/exist.json')
+      }.to raise_error(Errno::ENOENT, /No such file or directory/)
+    end
+
+    it 'writes data formatted as JSON to disk' do
+      file_path = file_containing('input', Puppet::Util::Json.dump({ "my" => "data" }))
+
+      expect(Puppet::Util::Json.load_file(file_path)).to eq({ "my" => "data" })
+    end
+  end
+end