From ef18b457869f1dad5d629a6edb14008b2aa3adbe Mon Sep 17 00:00:00 2001 From: Rui Lopes Date: Mon, 31 May 2021 20:20:12 +0100 Subject: [PATCH 1/6] ensure that the shell provisioner only emits complete output lines to the ui partial lines are buffered until they are complete (have a line ending) see https://github.com/hashicorp/vagrant/issues/11047 --- lib/vagrant/util/line_buffer.rb | 35 ++++++++++ plugins/provisioners/shell/provisioner.rb | 78 ++++++++++++++++------- 2 files changed, 91 insertions(+), 22 deletions(-) create mode 100644 lib/vagrant/util/line_buffer.rb diff --git a/lib/vagrant/util/line_buffer.rb b/lib/vagrant/util/line_buffer.rb new file mode 100644 index 000000000..250c4749a --- /dev/null +++ b/lib/vagrant/util/line_buffer.rb @@ -0,0 +1,35 @@ +require 'stringio' + +module Vagrant + module Util + class LineBuffer + def initialize + @buffer = StringIO.new + end + + def lines(data, &block) + if data == nil + return + end + remaining_buffer = StringIO.new + @buffer << data + @buffer.string.each_line do |line| + if line.end_with? "\n" + block.call(line.rstrip) + else + remaining_buffer << line + break + end + end + @buffer = remaining_buffer + end + + def remaining(&block) + if @buffer.length > 0 + block.call(@buffer.string.rstrip) + @buffer = StringIO.new + end + end + end + end +end diff --git a/plugins/provisioners/shell/provisioner.rb b/plugins/provisioners/shell/provisioner.rb index 53800e254..51aa662a4 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,15 +66,42 @@ module VagrantPlugins protected - # This handles outputting the communication data 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. - color = type == :stdout ? :green : :red + # This buffers a communicator output data (which can have partial + # lines) into full lines and emits them back to the UI. + def comm_line_buffer(&block) + stdout = Vagrant::Util::LineBuffer.new + stderr = Vagrant::Util::LineBuffer.new - # Clear out the newline since we add one - data = data.chomp - return if data.empty? + line_buffer_handle_comm = Proc.new do |type, data| + case type + when :stdout + stdout.lines data do |line| + handle_comm(type, line) + end + when :stderr + stderr.lines data do |line| + handle_comm(type, line) + end + end + end + + # execute the block and emit complete lines to the UI. + block.call(line_buffer_handle_comm) + + # emit the remaining as incomplete/partial lines to the UI. + stdout.remaining do |line| + handle_comm(:stdout, line) + end + stderr.remaining do |line| + handle_comm(:stderr, line) + end + end + + # This handles outputting the communication line back to the UI + def handle_comm(type, line) + if [:stderr, :stdout].include?(type) + # Output the line with the proper color based on the stream. + color = type == :stdout ? :green : :red options = {} options[:color] = color if !config.keep_color @@ -121,12 +149,14 @@ 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) + comm_line_buffer do |handle_comm| + comm.execute( + command, + sudo: config.privileged, + error_key: :ssh_bad_exit_status_muted + ) do |type, data| + handle_comm.call(type, data) + end end end end @@ -176,12 +206,14 @@ 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) + comm_line_buffer do |handle_comm| + comm.execute( + command, + shell: :powershell, + error_key: :ssh_bad_exit_status_muted + ) do |type, data| + handle_comm.call(type, data) + end end end end @@ -245,8 +277,10 @@ 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) + comm_line_buffer do |handle_comm| + comm.sudo(command, { elevated: config.privileged, interactive: config.powershell_elevated_interactive }) do |type, data| + handle_comm.call(type, data) + end end end end From 35de0d724e5ccf240c1d06de399884a77a09a69c Mon Sep 17 00:00:00 2001 From: Chris Roberts Date: Tue, 29 Jun 2021 09:57:18 -0700 Subject: [PATCH 2/6] Remove stringio usage from line buffer util, add max buffer length --- lib/vagrant/util/line_buffer.rb | 67 ++++++++++++++++++++++----------- 1 file changed, 46 insertions(+), 21 deletions(-) diff --git a/lib/vagrant/util/line_buffer.rb b/lib/vagrant/util/line_buffer.rb index 250c4749a..14271b98d 100644 --- a/lib/vagrant/util/line_buffer.rb +++ b/lib/vagrant/util/line_buffer.rb @@ -1,34 +1,59 @@ -require 'stringio' - module Vagrant module Util class LineBuffer - def initialize - @buffer = StringIO.new + + # 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 - def lines(data, &block) - if data == nil - return - end - remaining_buffer = StringIO.new - @buffer << data - @buffer.string.each_line do |line| - if line.end_with? "\n" - block.call(line.rstrip) - else - remaining_buffer << line - break + # 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 - @buffer = remaining_buffer + self end - def remaining(&block) - if @buffer.length > 0 - block.call(@buffer.string.rstrip) - @buffer = StringIO.new + # 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 From 6cc6b89dcfd29f5cb8c306cf080e3fae4c973314 Mon Sep 17 00:00:00 2001 From: Chris Roberts Date: Tue, 29 Jun 2021 09:57:56 -0700 Subject: [PATCH 3/6] Include line buffer util in autoloads --- lib/vagrant/util.rb | 1 + 1 file changed, 1 insertion(+) 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' From b1ae4d88fb95f960ef317128eadda618ccd8e99e Mon Sep 17 00:00:00 2001 From: Chris Roberts Date: Tue, 29 Jun 2021 09:58:11 -0700 Subject: [PATCH 4/6] Update shell provisioner usage of line buffer --- plugins/provisioners/shell/provisioner.rb | 77 ++++++++++------------- 1 file changed, 33 insertions(+), 44 deletions(-) diff --git a/plugins/provisioners/shell/provisioner.rb b/plugins/provisioners/shell/provisioner.rb index 51aa662a4..477cf81d9 100644 --- a/plugins/provisioners/shell/provisioner.rb +++ b/plugins/provisioners/shell/provisioner.rb @@ -66,39 +66,19 @@ module VagrantPlugins protected - # This buffers a communicator output data (which can have partial - # lines) into full lines and emits them back to the UI. - def comm_line_buffer(&block) - stdout = Vagrant::Util::LineBuffer.new - stderr = Vagrant::Util::LineBuffer.new - - line_buffer_handle_comm = Proc.new do |type, data| - case type - when :stdout - stdout.lines data do |line| - handle_comm(type, line) - end - when :stderr - stderr.lines data do |line| - handle_comm(type, line) - end - end - end - - # execute the block and emit complete lines to the UI. - block.call(line_buffer_handle_comm) - - # emit the remaining as incomplete/partial lines to the UI. - stdout.remaining do |line| - handle_comm(:stdout, line) - end - stderr.remaining do |line| - handle_comm(:stderr, line) - end + 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, line) + def handle_comm(type, data) if [:stderr, :stdout].include?(type) # Output the line with the proper color based on the stream. color = type == :stdout ? :green : :red @@ -149,14 +129,16 @@ module VagrantPlugins end # Execute it with sudo - comm_line_buffer do |handle_comm| + outputs, handler = build_outputs + begin comm.execute( command, sudo: config.privileged, - error_key: :ssh_bad_exit_status_muted - ) do |type, data| - handle_comm.call(type, data) - end + error_key: :ssh_bad_exit_status_muted, + &handler + ) + ensure + outputs.values.map(&:close) end end end @@ -206,14 +188,16 @@ module VagrantPlugins end # Execute it with sudo - comm_line_buffer do |handle_comm| + begin + outputs, handler = build_outputs comm.execute( command, shell: :powershell, - error_key: :ssh_bad_exit_status_muted - ) do |type, data| - handle_comm.call(type, data) - end + error_key: :ssh_bad_exit_status_muted, + &handler + ) + ensure + outputs.values.map(&:close) end end end @@ -277,10 +261,15 @@ module VagrantPlugins end # Execute it with sudo - comm_line_buffer do |handle_comm| - comm.sudo(command, { elevated: config.privileged, interactive: config.powershell_elevated_interactive }) do |type, data| - handle_comm.call(type, data) - end + 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 From aee58785b29fb1930be7ee197810325ef502d36b Mon Sep 17 00:00:00 2001 From: Chris Roberts Date: Tue, 29 Jun 2021 09:58:55 -0700 Subject: [PATCH 5/6] Add tests on provisioner to validate expected output behavior --- .../provisioners/shell/provisioner_test.rb | 41 ++++++++++++++++++- 1 file changed, 39 insertions(+), 2 deletions(-) diff --git a/test/unit/plugins/provisioners/shell/provisioner_test.rb b/test/unit/plugins/provisioners/shell/provisioner_test.rb index fc9f0f586..87e797369 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 From 9b232d0ff39659f989c3344f3648f7b5108c9b82 Mon Sep 17 00:00:00 2001 From: Chris Roberts Date: Tue, 29 Jun 2021 09:59:17 -0700 Subject: [PATCH 6/6] Add test coverage on line buffer util --- test/unit/vagrant/util/line_buffer_test.rb | 60 ++++++++++++++++++++++ 1 file changed, 60 insertions(+) create mode 100644 test/unit/vagrant/util/line_buffer_test.rb 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