From f2ff7835879e4f6c21804930a91c2b1f956ee8ef 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 | 40 +++++++++++++++++++++++++++++++----
 1 file changed, 36 insertions(+), 4 deletions(-)

diff --git a/lib/train-winrm/connection.rb b/lib/train-winrm/connection.rb
index 7c9dcaf..dcbc0fa 100644
--- a/lib/train-winrm/connection.rb
+++ b/lib/train-winrm/connection.rb
@@ -107,15 +107,47 @@ 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
+
+        # Run the command in a thread, to support timing out the command
+        thr = Thread.new do
+          # Surface any exceptions in this thread back to this method
+          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
+            # If this command hits timeout, calling close with the command currently running causes
+            # a RuntimeError error in WinRM's cleanup code. This specific error can be ignored.
+            # The command will be terminated and further commands can be sent on the connection.
+            raise e unless timeout && e.to_s == "opts[:shell_id] is required"
+          rescue WinRM::WinRMHTTPTransportError => e
+            # If this command hits timeout, there is also a potential race in the HTTP transport
+            # where decryption is attempted on an empty message.
+            raise e unless timeout && e.to_s == "Could not decrypt NTLM message. ()."
+          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)