Allow timeout option for WinRM commands

Allows end users (e.g. InSpec test coders) to specify a timeout for a potentially long running command.
If the timeout is reached, the command is expected to be terminated on the host and an exception is raised (a subsequent change to InSpec will handle this exception).
This complements recent changes to the base connection and ssh connection in train: inspec/train#625

Signed-off-by: James Stocks <[email protected]>

lib/train-winrm/connection.rb

@@ -107,15 +107,41 @@ def file_via_connection(path)
         Train::File::Remote::Windows.new(self, path)
       end
 
-      def run_command_via_connection(command, &data_handler)
+      def run_command_via_connection(command, opts = {}, &data_handler)
         return if command.nil?
 
         logger.debug("[WinRM] #{self} (#{command})")
         out = ""
+        response = nil
+        timeout = opts[:timeout]&.to_i
+
+        thr = Thread.new do
+          # Run the command in a thread, surfacing any exceptions to the main thread
+          Thread.current.report_on_exception = false
+          Thread.current.abort_on_exception = true
+          begin
+            response = session.run(command) do |stdout, _|
+              yield(stdout) if data_handler && stdout
+              out << stdout if stdout
+            end
+          rescue RuntimeError => e
+            # Ref: https://github.com/WinRb/WinRM/issues/315
+            # Calling close on a session with a command currently running causes a RuntimeError
+            # in WinRM's cleanup code. This specific error can be ignored.
+            raise e unless timeout && e.to_s == "opts[:shell_id] is required"
+          end
+        end
 
-        response = session.run(command) do |stdout, _|
-          yield(stdout) if data_handler && stdout
-          out << stdout if stdout
+        if timeout
+          res = thr.join(timeout)
+          unless res
+            msg = "PowerShell command '(#{command})' reached timeout (#{timeout}s)"
+            logger.info("[WinRM] #{msg}")
+            close
+            raise Train::CommandTimeoutReached.new msg
+          end
+        else
+          thr.join
         end
 
         CommandResult.new(out, response.stderr, response.exitcode)