Allow timeout option for WinRM commands

lib/train-winrm/connection.rb

@@ -33,6 +33,8 @@
 
 require "train"
 require "train/plugins"
+# This module may need to directly require WinRM to reference its exception classes
+require "winrm" unless defined?(WinRM)
 
 module TrainPlugins
   module WinRM
@@ -108,15 +110,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 ::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. ()."
+          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"
+          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)

test/unit/connection_test.rb

@@ -37,9 +37,9 @@
       # We need to test run_command b/c run_command_via_connection is private.
       winrm.run_command("test") do |data|
         called = true
-        data.must_equal "testdata"
+        _(data).must_equal "testdata"
       end
-      called.must_equal true
+      _(called).must_equal true
     end
   end
 

test/unit/transport_test.rb

@@ -24,47 +24,47 @@
     let(:winrm) { cls.new({ host: "dummy", logger: Logger.new(STDERR, level: :info) }) }
 
     it "can be instantiated (with valid config)" do
-      winrm.wont_be_nil
+      _(winrm).wont_be_nil
     end
 
     it "configures the host" do
-      winrm.options[:host].must_equal "dummy"
+      _(winrm.options[:host]).must_equal "dummy"
     end
 
     it "has default endpoint" do
-      winrm.options[:endpoint].must_be_nil
+      _(winrm.options[:endpoint]).must_be_nil
     end
 
     it "has default path set" do
-      winrm.options[:path].must_equal "/wsman"
+      _(winrm.options[:path]).must_equal "/wsman"
     end
 
     it "has default ssl set" do
-      winrm.options[:ssl].must_equal false
+      _(winrm.options[:ssl]).must_equal false
     end
 
     it "has default self_signed set" do
-      winrm.options[:self_signed].must_equal false
+      _(winrm.options[:self_signed]).must_equal false
     end
 
     it "has default rdp_port set" do
-      winrm.options[:rdp_port].must_equal 3389
+      _(winrm.options[:rdp_port]).must_equal 3389
     end
 
     it "has default winrm_transport set" do
-      winrm.options[:winrm_transport].must_equal :negotiate
+      _(winrm.options[:winrm_transport]).must_equal :negotiate
     end
 
     it "has default winrm_disable_sspi set" do
-      winrm.options[:winrm_disable_sspi].must_equal false
+      _(winrm.options[:winrm_disable_sspi]).must_equal false
     end
 
     it "has default winrm_basic_auth_only set" do
-      winrm.options[:winrm_basic_auth_only].must_equal false
+      _(winrm.options[:winrm_basic_auth_only]).must_equal false
     end
 
     it "has default user" do
-      winrm.options[:user].must_equal "administrator"
+      _(winrm.options[:user]).must_equal "administrator"
     end
   end
 
@@ -73,21 +73,21 @@
     let(:connection) { winrm.connection }
     it "without ssl genrates uri" do
       conf[:host] = "dummy_host"
-      connection.uri.must_equal "winrm://administrator@http://dummy_host:5985/wsman:3389"
+      _(connection.uri).must_equal "winrm://administrator@http://dummy_host:5985/wsman:3389"
     end
 
     it "without ssl genrates uri" do
       conf[:ssl] = true
       conf[:host] = "dummy_host_ssl"
-      connection.uri.must_equal "winrm://administrator@https://dummy_host_ssl:5986/wsman:3389"
+      _(connection.uri).must_equal "winrm://administrator@https://dummy_host_ssl:5986/wsman:3389"
     end
   end
 
   describe "options validation" do
     let(:winrm) { cls.new(conf) }
     it "raises an error when a non-supported winrm_transport is specificed" do
       conf[:winrm_transport] = "invalid"
-      proc { winrm.connection }.must_raise Train::ClientError
+      _(proc { winrm.connection }).must_raise Train::ClientError
     end
   end
 end