beckermj/vaguerent
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #12437 from chrisroberts/line-buf-shell-util

Browse Source 
Line buffer utility and shell provisioner usage
This commit is contained in:
Chris Roberts 2021-06-29 13:11:19 -07:00 committed by GitHub
commit bbd1404c1e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 203 additions and 22 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
lib/vagrant/util.rb
View File

@ -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'

60
lib/vagrant/util/line_buffer.rb Normal file
View File

@ -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

63
plugins/provisioners/shell/provisioner.rb
View File

@ -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

41
test/unit/plugins/provisioners/shell/provisioner_test.rb
View File

@ -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

60
test/unit/vagrant/util/line_buffer_test.rb Normal file
View File

@ -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