Keep `kill_process_tree` from killing protected system processes.

2 files changed, 68 insertions, 16 deletions

diff --git a/lib/mixlib/shellout/windows.rb b/lib/mixlib/shellout/windows.rb
index fd8039b..6957ea9 100644
--- a/lib/mixlib/shellout/windows.rb
+++ b/lib/mixlib/shellout/windows.rb
@@ -321,31 +321,43 @@ module Mixlib
           File.executable?(path) && !File.directory?(path)
         end
 
+        def self.system_required_processes
+          [
+            'System Idle Process',
+            'System',
+            'spoolsv.exe',
+            'lsass.exe',
+            'csrss.exe',
+            'smss.exe',
+            'svchost.exe'
+          ]
+        end
+
         # recursively kills all child processes of given pid
         # calls itself querying for children child procs until
         # none remain. Important that a single WmiLite instance
         # is passed in since each creates its own WMI rpc process
         def self.kill_process_tree(pid, wmi, logger)
           wmi.query("select * from Win32_Process where ParentProcessID=#{pid}").each do |instance|
+            next if system_required_processes.include? instance.wmi_ole_object.name
             child_pid = instance.wmi_ole_object.processid
             kill_process_tree(child_pid, wmi, logger)
-            begin
-              logger.debug([
-                "killing child process #{child_pid}::",
-                "#{instance.wmi_ole_object.Name} of parent #{pid}"
-                ].join) if logger
-              kill_process(instance)
-            rescue Errno::EIO, SystemCallError
-              logger.debug([
-                "Failed to kill child process #{child_pid}::",
-                "#{instance.wmi_ole_object.Name} of parent #{pid}"
-              ].join) if logger
-            end
+            kill_process(instance, logger)
           end
         end
 
-        def self.kill_process(instance)
+        def self.kill_process(instance, logger)
+          child_pid = instance.wmi_ole_object.processid
+          logger.debug([
+            "killing child process #{child_pid}::",
+            "#{instance.wmi_ole_object.Name} of parent #{pid}"
+            ].join) if logger
           Process.kill(:KILL, instance.wmi_ole_object.processid)
+        rescue Errno::EIO, SystemCallError
+          logger.debug([
+            "Failed to kill child process #{child_pid}::",
+            "#{instance.wmi_ole_object.Name} of parent #{pid}"
+          ].join) if logger
         end
 
         def self.format_process(process, app_name, command_line, timeout)
diff --git a/spec/mixlib/shellout/windows_spec.rb b/spec/mixlib/shellout/windows_spec.rb
index dd8a80b..eb8ee78 100644
--- a/spec/mixlib/shellout/windows_spec.rb
+++ b/spec/mixlib/shellout/windows_spec.rb
@@ -1,7 +1,7 @@
 require 'spec_helper'
 
 describe 'Mixlib::ShellOut::Windows', :windows_only do
- 
+
   describe 'Utils' do
     describe '.should_run_under_cmd?' do
       subject { Mixlib::ShellOut::Windows::Utils.should_run_under_cmd?(command) }
@@ -108,6 +108,44 @@ describe 'Mixlib::ShellOut::Windows', :windows_only do
         end
       end
     end
+
+    describe '.kill_process_tree' do
+      let(:utils) { Mixlib::ShellOut::Windows::Utils }
+      let(:wmi) { Object.new }
+      let(:wmi_ole_object) { Object.new }
+      let(:wmi_process) { Object.new }
+
+      before do
+        allow(wmi).to receive(:query).and_return([wmi_process])
+        allow(wmi_process).to receive(:wmi_ole_object).and_return(wmi_ole_object)
+      end
+
+      context 'with a protected system process in the process tree' do
+        before do
+          allow(wmi_ole_object).to receive(:name).and_return('csrss.exe')
+          allow(wmi_ole_object).to receive(:processid).and_return(100)
+        end
+
+        it 'does not attempt to kill csrss.exe' do
+          expect(utils).to_not receive(:kill_process)
+          utils.kill_process_tree(200, wmi, nil)
+        end
+      end
+
+      context 'with a non-system-critical process in the process tree' do
+        before do
+          allow(wmi_ole_object).to receive(:name).and_return('blah.exe')
+          allow(wmi_ole_object).to receive(:processid).and_return(300)
+        end
+
+        it 'does attempt to kill blah.exe' do
+          expect(utils).to receive(:kill_process).with(wmi_process, nil)
+          expect(utils).to receive(:kill_process_tree).with(200, wmi, nil).and_call_original
+          expect(utils).to receive(:kill_process_tree).with(300, wmi, nil)
+          utils.kill_process_tree(200, wmi, nil)
+        end
+      end
+    end
   end
 
   # Caveat: Private API methods are subject to change without notice.
@@ -170,7 +208,7 @@ describe 'Mixlib::ShellOut::Windows', :windows_only do
           allow(ENV).to receive(:[]).with('COMSPEC').and_return('C:\Windows\system32\cmd.exe')
           allow(File).to receive(:executable?).and_return(false)
           allow(File).to receive(:executable?).with(executable_path).and_return(true)
-          allow(File).to receive(:directory?).and_return(false)            
+          allow(File).to receive(:directory?).and_return(false)
         end
 
         it 'should return with full path with extension' do
@@ -181,7 +219,7 @@ describe 'Mixlib::ShellOut::Windows', :windows_only do
           before do
             # File.executable? returns true for directories
             allow(File).to receive(:executable?).with(cmd).and_return(true)
-            allow(File).to receive(:directory?).with(cmd).and_return(true)            
+            allow(File).to receive(:directory?).with(cmd).and_return(true)
           end
 
           it 'should return with full path with extension' do
@@ -196,6 +234,8 @@ describe 'Mixlib::ShellOut::Windows', :windows_only do
         let(:cmd_exe) { "C:\\Windows\\system32\\cmd.exe" }
         let(:cmd) { "#{executable_path}" }
 
+        before { ENV['ComSpec'] = 'C:\Windows\system32\cmd.exe' }
+
         context 'with .bat file' do
           let(:executable_path) { '"C:\Program Files\Application\Start.bat"' }