diff --git a/lib/train-winrm/connection.rb b/lib/train-winrm/connection.rb index 2be966f..24747c2 100644 --- a/lib/train-winrm/connection.rb +++ b/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) diff --git a/test/unit/connection_test.rb b/test/unit/connection_test.rb index afeb431..8e11ced 100644 --- a/test/unit/connection_test.rb +++ b/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 diff --git a/test/unit/transport_test.rb b/test/unit/transport_test.rb index 71c3710..eb52d39 100644 --- a/test/unit/transport_test.rb +++ b/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,13 +73,13 @@ 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 @@ -87,7 +87,7 @@ 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