diff --git a/lib/vagrant/util.rb b/lib/vagrant/util.rb index 3eca4daa3..c135dacc2 100644 --- a/lib/vagrant/util.rb +++ b/lib/vagrant/util.rb @@ -26,6 +26,7 @@ module Vagrant autoload :IPV4Interfaces, 'vagrant/util/ipv4_interfaces' autoload :IsPortOpen, 'vagrant/util/is_port_open' autoload :KeyPair, 'vagrant/util/key_pair' + autoload :LineBuffer, 'vagrant/util/line_buffer' autoload :LineEndingHelpers, 'vagrant/util/line_ending_helpers' autoload :LoggingFormatter, 'vagrant/util/logging_formatter' autoload :MapCommandOptions, 'vagrant/util/map_command_options' diff --git a/lib/vagrant/util/line_buffer.rb b/lib/vagrant/util/line_buffer.rb new file mode 100644 index 000000000..14271b98d --- /dev/null +++ b/lib/vagrant/util/line_buffer.rb @@ -0,0 +1,60 @@ +module Vagrant + module Util + class LineBuffer + + # Maximum number of characters to buffer before sending + # to callback without detecting a new line + MAX_LINE_LENGTH = 5000.freeze + + # Create a new line buffer. The registered block + # will be called when a new line is encountered on + # provided input, or the max line length is reached + def initialize(&callback) + raise ArgumentError, + "Expected callback but received none" if callback.nil? + @mu = Mutex.new + @callback = callback + @buffer = "" + end + + # Add string data to output + # + # @param [String] str String of data to output + # @return [self] + def <<(str) + @mu.synchronize do + while i = str.index("\n") + @callback.call((@buffer + str[0, i+1]).rstrip) + @buffer.clear + str = str[i+1, str.length].to_s + end + + @buffer << str.to_s + + if @buffer.length > MAX_LINE_LENGTH + @callback.call(@buffer.dup) + @buffer.clear + end + end + self + end + + # Closes the buffer. Any remaining data that has + # been buffered will be given to the callback. + # Once closed the instance will no longer be usable. + # + # @return [self] + def close + @mu.synchronize do + # Send any remaining output on the buffer + @callback.call(@buffer.dup) if !@buffer.empty? + # Disable this buffer instance + @callback = nil + @buffer.clear + @buffer.freeze + end + self + end + end + end +end diff --git a/plugins/provisioners/shell/provisioner.rb b/plugins/provisioners/shell/provisioner.rb index 53800e254..477cf81d9 100644 --- a/plugins/provisioners/shell/provisioner.rb +++ b/plugins/provisioners/shell/provisioner.rb @@ -2,6 +2,7 @@ require "pathname" require "tempfile" require "vagrant/util/downloader" +require "vagrant/util/line_buffer" require "vagrant/util/retryable" module VagrantPlugins @@ -65,16 +66,23 @@ module VagrantPlugins protected - # This handles outputting the communication data back to the UI + def build_outputs + outputs = { + stdout: Vagrant::Util::LineBuffer.new { |line| handle_comm(:stdout, line) }, + stderr: Vagrant::Util::LineBuffer.new { |line| handle_comm(:stderr, line) }, + } + block = proc { |type, data| + outputs[type] << data if outputs[type] + } + return outputs, block + end + + # This handles outputting the communication line back to the UI def handle_comm(type, data) if [:stderr, :stdout].include?(type) - # Output the data with the proper color based on the stream. + # Output the line with the proper color based on the stream. color = type == :stdout ? :green : :red - # Clear out the newline since we add one - data = data.chomp - return if data.empty? - options = {} options[:color] = color if !config.keep_color @@ -121,12 +129,16 @@ module VagrantPlugins end # Execute it with sudo - comm.execute( - command, - sudo: config.privileged, - error_key: :ssh_bad_exit_status_muted - ) do |type, data| - handle_comm(type, data) + outputs, handler = build_outputs + begin + comm.execute( + command, + sudo: config.privileged, + error_key: :ssh_bad_exit_status_muted, + &handler + ) + ensure + outputs.values.map(&:close) end end end @@ -176,12 +188,16 @@ module VagrantPlugins end # Execute it with sudo - comm.execute( - command, - shell: :powershell, - error_key: :ssh_bad_exit_status_muted - ) do |type, data| - handle_comm(type, data) + begin + outputs, handler = build_outputs + comm.execute( + command, + shell: :powershell, + error_key: :ssh_bad_exit_status_muted, + &handler + ) + ensure + outputs.values.map(&:close) end end end @@ -245,8 +261,15 @@ module VagrantPlugins end # Execute it with sudo - comm.sudo(command, { elevated: config.privileged, interactive: config.powershell_elevated_interactive }) do |type, data| - handle_comm(type, data) + begin + outputs, handler = build_outputs + comm.sudo(command, + elevated: config.privileged, + interactive: config.powershell_elevated_interactive, + &handler + ) + ensure + outputs.values.map(&:close) end end end diff --git a/test/unit/plugins/provisioners/shell/provisioner_test.rb b/test/unit/plugins/provisioners/shell/provisioner_test.rb index 691a20fe8..7b8e07214 100644 --- a/test/unit/plugins/provisioners/shell/provisioner_test.rb +++ b/test/unit/plugins/provisioners/shell/provisioner_test.rb @@ -438,7 +438,8 @@ describe "Vagrant::Shell::Provisioner" do :powershell_args => "", :name => nil, :privileged => false, - :powershell_elevated_interactive => false + :powershell_elevated_interactive => false, + :keep_color => true, ) } @@ -456,8 +457,26 @@ describe "Vagrant::Shell::Provisioner" do allow(machine).to receive(:communicate).and_return(communicator) allow(machine).to receive(:guest).and_return(guest) allow(machine).to receive(:ui).and_return(ui) + allow(vsp).to receive(:with_script_file).and_yield(config.path) + allow(communicator).to receive(:upload).with(config.path, /arbitrary.ps1$/) } + it "should output all received output" do + stdout = ["two lines\n", "from stdout\n"] + stderr = ["one line\n", "and partial from stderr"] + expect(communicator).to receive(:sudo). + and_yield(:stdout, stdout.first). + and_yield(:stderr, stderr.first). + and_yield(:stderr, stderr.last). + and_yield(:stdout, stdout.last) + allow(ui).to receive(:detail) + expect(ui).to receive(:detail).with("two lines", any_args) + expect(ui).to receive(:detail).with("from stdout", any_args) + expect(ui).to receive(:detail).with("one line", any_args) + expect(ui).to receive(:detail).with("and partial from stderr", any_args) + vsp.send(:provision_winrm, "") + end + it "ensures that files are uploaded with an extension" do allow(vsp).to receive(:with_script_file).and_yield(config.path) expect(communicator).to receive(:upload).with(config.path, /arbitrary.ps1$/) @@ -533,7 +552,8 @@ describe "Vagrant::Shell::Provisioner" do :powershell_args => "", :name => nil, :privileged => false, - :powershell_elevated_interactive => false + :powershell_elevated_interactive => false, + :keep_color => true, ) } @@ -548,6 +568,7 @@ describe "Vagrant::Shell::Provisioner" do before { allow(guest).to receive(:capability?).with(:wait_for_reboot).and_return(false) allow(communicator).to receive(:sudo) + allow(communicator).to receive(:upload) allow(communicator).to receive_message_chain(:machine_config_ssh, :shell) allow(machine).to receive(:communicate).and_return(communicator) allow(machine).to receive(:guest).and_return(guest) @@ -567,6 +588,22 @@ describe "Vagrant::Shell::Provisioner" do expect(communicator).to receive(:execute).with(/powershell.*arbitrary.ps1/, anything) vsp.send(:provision_winssh, "") end + + it "should output all received output" do + stdout = ["two lines\n", "from stdout\n"] + stderr = ["one line\n", "and partial from stderr"] + expect(communicator).to receive(:execute). + and_yield(:stdout, stdout.first). + and_yield(:stderr, stderr.first). + and_yield(:stderr, stderr.last). + and_yield(:stdout, stdout.last) + allow(ui).to receive(:detail) + expect(ui).to receive(:detail).with("two lines", any_args) + expect(ui).to receive(:detail).with("from stdout", any_args) + expect(ui).to receive(:detail).with("one line", any_args) + expect(ui).to receive(:detail).with("and partial from stderr", any_args) + vsp.send(:provision_winssh, "") + end end context "bat file being uploaded" do diff --git a/test/unit/vagrant/util/line_buffer_test.rb b/test/unit/vagrant/util/line_buffer_test.rb new file mode 100644 index 000000000..df16830f1 --- /dev/null +++ b/test/unit/vagrant/util/line_buffer_test.rb @@ -0,0 +1,60 @@ +require File.expand_path("../../../base", __FILE__) +require "vagrant/util/line_buffer" + +describe Vagrant::Util::LineBuffer do + it "should raise error when no callback is provided" do + expect { subject }.to raise_error(ArgumentError) + end + + context "with block defined" do + let(:block) { proc{ |l| output << l } } + let(:output) { [] } + let(:partial) { "this is part of a line. " } + let(:line) { "this is a full line\n" } + + subject { described_class.new(&block) } + + it "should not raise an error when callback is provided" do + expect { subject }.not_to raise_error + end + + describe "#<<" do + it "should add line to the output" do + subject << line + expect(output).to eq([line.rstrip]) + end + + it "should not add partial line to output" do + subject << partial + expect(output).to be_empty + end + + it "should add partial line to output once full line is given" do + subject << partial + expect(output).to be_empty + subject << line + expect(output).to eq([partial + line.rstrip]) + end + + it "should add line once it has surpassed max line length" do + overflow = "a" * (described_class.const_get(:MAX_LINE_LENGTH) + 1) + subject << overflow + expect(output).to eq([overflow]) + end + end + + describe "#close" do + it "should output any partial data left in buffer" do + subject << partial + expect(output).to be_empty + subject.close + expect(output).to eq([partial]) + end + + it "should not be writable after closing" do + subject.close + expect { subject << partial }.to raise_error(FrozenError) + end + end + end +end