Backport "Make the DMI IDs we whitelist configurable" to Ohai 13

Backport of #970

Signed-off-by: Gavin Reynolds <gavin@chef.io>

4 files changed, 44 insertions, 25 deletions

diff --git a/lib/ohai/common/dmi.rb b/lib/ohai/common/dmi.rb
index 988e1493..3403bd93 100644
--- a/lib/ohai/common/dmi.rb
+++ b/lib/ohai/common/dmi.rb
@@ -72,10 +72,22 @@ module Ohai
         127 => "end_of_table_marker",
       }
 
-      # list of IDs to collect, otherwise we generate pages of hashes about cache chip size and whatnot
-      # See OHAI-260. When we can give the user a choice, this will be a default.
+      # list of IDs to collect from config or default to a sane list that prunes
+      # away some of the less useful IDs
       ID_TO_CAPTURE = [ 0, 1, 2, 3, 4, 6, 11 ]
 
+      # return the list of DMI IDs to capture
+      def whitelisted_ids
+        if Ohai.config[:additional_dmi_ids]
+          if [ Integer, Array ].include?(Ohai.config[:additional_dmi_ids].class)
+            return ID_TO_CAPTURE + Array(Ohai.config[:additional_dmi_ids])
+          else
+            Ohai::Log.warn("The DMI plugin additional_dmi_ids config must be an array of IDs!")
+          end
+        end
+        ID_TO_CAPTURE
+      end
+
       # look up DMI ID
       def id_lookup(id)
         id = id.to_i
@@ -85,7 +97,7 @@ module Ohai
           id = DMI::ID_TO_DESCRIPTION[id]
         else
           Ohai::Log.debug("unrecognized header id; falling back to 'unknown'")
-          id = "unknown"
+          id = "unknown_dmi_id_#{id}"
         end
       rescue
         Ohai::Log.debug("failed to look up id #{id}, returning unchanged")
@@ -122,7 +134,7 @@ module Ohai
         end
       end
 
-      module_function :id_lookup, :convenience_keys
+      module_function :id_lookup, :convenience_keys, :whitelisted_ids
     end
   end
 end
diff --git a/lib/ohai/plugins/dmi.rb b/lib/ohai/plugins/dmi.rb
index ed7ae9f2..e350873c 100644
--- a/lib/ohai/plugins/dmi.rb
+++ b/lib/ohai/plugins/dmi.rb
@@ -74,9 +74,8 @@ Ohai.plugin(:DMI) do
         elsif table_location = table_location_line.match(line)
           dmi[:table_location] = table_location[1]
 
-        elsif handle = handle_line.match(line)
-          # Don't overcapture for now (OHAI-260)
-          unless Ohai::Common::DMI::ID_TO_CAPTURE.include?(handle[2].to_i)
+        elsif ( handle = handle_line.match(line) )
+          unless Ohai::Common::DMI.whitelisted_ids.include?(handle[2].to_i)
             dmi_record = nil
             next
           end
diff --git a/lib/ohai/plugins/solaris2/dmi.rb b/lib/ohai/plugins/solaris2/dmi.rb
index d2aa35a4..750c5f81 100644
--- a/lib/ohai/plugins/solaris2/dmi.rb
+++ b/lib/ohai/plugins/solaris2/dmi.rb
@@ -129,7 +129,7 @@ Ohai.plugin(:DMI) do
           id = smb_to_id[header_information[3]]
 
           # Don't overcapture for now (OHAI-260)
-          unless Ohai::Common::DMI::ID_TO_CAPTURE.include?(id)
+          unless Ohai::Common::DMI.whitelisted_ids.include?(id)
             dmi_record = nil
             next
           end
diff --git a/spec/unit/plugins/dmi_spec.rb b/spec/unit/plugins/dmi_spec.rb
index 9c855a9c..a8074bf0 100644
--- a/spec/unit/plugins/dmi_spec.rb
+++ b/spec/unit/plugins/dmi_spec.rb
@@ -101,15 +101,16 @@ Chassis Information
 EOS
 
 describe Ohai::System, "plugin dmi" do
+  let(:plugin) { get_plugin("dmi") }
+  let(:stdout) { DMI_OUT }
+
   before(:each) do
-    @plugin = get_plugin("dmi")
-    @stdout = DMI_OUT
-    allow(@plugin).to receive(:shell_out).with("dmidecode").and_return(mock_shell_out(0, @stdout, ""))
+    allow(plugin).to receive(:shell_out).with("dmidecode").and_return(mock_shell_out(0, stdout, ""))
   end
 
-  it "should run dmidecode" do
-    expect(@plugin).to receive(:shell_out).with("dmidecode").and_return(mock_shell_out(0, @stdout, ""))
-    @plugin.run
+  it "runs dmidecode" do
+    expect(plugin).to receive(:shell_out).with("dmidecode").and_return(mock_shell_out(0, stdout, ""))
+    plugin.run
   end
 
   # Test some simple sample data
@@ -128,21 +129,28 @@ describe Ohai::System, "plugin dmi" do
     },
   }.each do |id, data|
     data.each do |attribute, value|
-      it "should have [:dmi][:#{id}][:#{attribute}] set" do
-        @plugin.run
-        expect(@plugin[:dmi][id][attribute]).to eql(value)
+      it "attribute [:dmi][:#{id}][:#{attribute}] is set" do
+        plugin.run
+        expect(plugin[:dmi][id][attribute]).to eql(value)
       end
-      it "should have [:dmi][:#{id}][:#{attribute}] set for windows output" do
-        @stdout = convert_windows_output(DMI_OUT)
-        expect(@plugin).to receive(:shell_out).with("dmidecode").and_return(mock_shell_out(0, @stdout, ""))
-        @plugin.run
-        expect(@plugin[:dmi][id][attribute]).to eql(value)
+      it "attribute [:dmi][:#{id}][:#{attribute}] set for windows output" do
+        stdout = convert_windows_output(DMI_OUT)
+        expect(plugin).to receive(:shell_out).with("dmidecode").and_return(mock_shell_out(0, stdout, ""))
+        plugin.run
+        expect(plugin[:dmi][id][attribute]).to eql(value)
       end
     end
   end
 
-  it "should correctly ignore unwanted data" do
-    @plugin.run
-    expect(@plugin[:dmi][:base_board]).not_to have_key(:error_correction_type)
+  it "allows capturing additional DMI data" do
+    Ohai.config[:additional_dmi_ids] = [ 16 ]
+    plugin.run
+    expect(plugin[:dmi]).to have_key(:physical_memory_array)
+  end
+
+  it "correctly ignores data in excluded DMI IDs" do
+    expect(plugin).to receive(:shell_out).with("dmidecode").and_return(mock_shell_out(0, stdout, ""))
+    plugin.run
+    expect(plugin[:dmi]).not_to have_key(:physical_memory_array)
   end
 end