From a6323fb2cbe5ce75e4e8d7fec092033fe0c06345 Mon Sep 17 00:00:00 2001
From: James Stocks <jstocks@chef.io>
Date: Wed, 23 Sep 2020 16:11:11 +0100
Subject: [PATCH] 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: https://github.com/inspec/train/pull/625

Signed-off-by: James Stocks <jstocks@chef.io>
---
 lib/train-winrm/connection.rb | 34 ++++++++++++++++++++++++++++++----
 1 file changed, 30 insertions(+), 4 deletions(-)

diff --git a/lib/train-winrm/connection.rb b/lib/train-winrm/connection.rb
index 7c9dcaf..ea84185 100644
--- a/lib/train-winrm/connection.rb
+++ b/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)