Skip to content

Commit

Permalink
Merge pull request #671 from inspec/cw/add-timeout-to-mixlib-based-lo…
Browse files Browse the repository at this point in the history
…cal-runners
  • Loading branch information
clintoncwolfe authored Mar 24, 2021
2 parents 6893b8d + 11f6c11 commit badc6bd
Show file tree
Hide file tree
Showing 2 changed files with 21 additions and 9 deletions.
25 changes: 18 additions & 7 deletions lib/train/transports/local.rb
Original file line number Diff line number Diff line change
Expand Up @@ -86,9 +86,9 @@ def force_runner(command_runner)
end
end

def run_command_via_connection(cmd, &_data_handler)
def run_command_via_connection(cmd, opts, &_data_handler)
# Use the runner if it is available
return @runner.run_command(cmd) if defined?(@runner)
return @runner.run_command(cmd, opts) if defined?(@runner)

# If we don't have a runner, such as at the beginning of setting up the
# transport and performing the first few steps of OS detection, fall
Expand All @@ -115,13 +115,18 @@ def initialize(connection, options)
@cmd_wrapper = Local::CommandWrapper.load(connection, options)
end

def run_command(cmd)
def run_command(cmd, opts = {})
if defined?(@cmd_wrapper) && !@cmd_wrapper.nil?
cmd = @cmd_wrapper.run(cmd)
end

res = Mixlib::ShellOut.new(cmd)
res.run_command
res.timeout = opts[:timeout]
begin
res.run_command
rescue Mixlib::ShellOut::CommandTimeout
raise Train::CommandTimeoutReached
end
Local::CommandResult.new(res.stdout, res.stderr, res.exitstatus)
end

Expand All @@ -138,7 +143,7 @@ def initialize(powershell_cmd = "powershell")
@powershell_cmd = powershell_cmd
end

def run_command(script)
def run_command(script, opts)
# Prevent progress stream from leaking into stderr
script = "$ProgressPreference='SilentlyContinue';" + script

Expand All @@ -149,7 +154,12 @@ def run_command(script)
cmd = "#{@powershell_cmd} -NoProfile -EncodedCommand #{base64_script}"

res = Mixlib::ShellOut.new(cmd)
res.run_command
res.timeout = opts[:timeout]
begin
res.run_command
rescue Mixlib::ShellOut::CommandTimeout
raise Train::CommandTimeoutReached
end
Local::CommandResult.new(res.stdout, res.stderr, res.exitstatus)
end

Expand All @@ -176,9 +186,10 @@ def initialize(powershell_cmd = "powershell")
# A command that succeeds without setting an exit code will have exitstatus 0
# A command that exits with an exit code will have that value as exitstatus
# A command that fails (e.g. throws exception) before setting an exit code will have exitstatus 1
def run_command(cmd)
def run_command(cmd, _opts)
script = "$ProgressPreference='SilentlyContinue';" + cmd
encoded_script = Base64.strict_encode64(script)
# TODO: no way to safely implement timeouts here.
@pipe.puts(encoded_script)
@pipe.flush
res = OpenStruct.new(JSON.parse(Base64.decode64(@pipe.readline)))
Expand Down
5 changes: 3 additions & 2 deletions test/unit/transports/local_test.rb
Original file line number Diff line number Diff line change
Expand Up @@ -113,6 +113,7 @@ def initialize(user_opts = {})

def mock_run_cmd(cmd, &block)
cmd_runner.expect :run_command, nil
cmd_runner.expect :timeout=, nil, [nil]
Mixlib::ShellOut.stub :new, cmd_runner do |*args|
yield
end
Expand Down Expand Up @@ -164,7 +165,7 @@ def mock_run_cmd(cmd, &block)
.expects(:new)
.never

runner.expects(:run_command).with("not actually executed")
runner.expects(:run_command).with("not actually executed", {})
connection.run_command("not actually executed")
end

Expand All @@ -177,7 +178,7 @@ def mock_run_cmd(cmd, &block)
.expects(:new)
.returns(runner)

runner.expects(:run_command).with("not actually executed")
runner.expects(:run_command).with("not actually executed", {})
connection.run_command("not actually executed")
end
end
Expand Down

0 comments on commit badc6bd

Please sign in to comment.