diff --git a/lib/vagrant/ssh.rb b/lib/vagrant/ssh.rb index 1d7e394bf..87467812e 100644 --- a/lib/vagrant/ssh.rb +++ b/lib/vagrant/ssh.rb @@ -49,7 +49,7 @@ module Vagrant env.config[:ssh][:username], opts.merge( :port => port, :keys => [env.config.ssh.private_key_path])) do |ssh| - yield ssh + yield SSH::Session.new(ssh) end end @@ -58,7 +58,7 @@ module Vagrant # the arguments to `Net::SCP#upload!` so view that for more information. def upload!(from, to) execute do |ssh| - scp = Net::SCP.new(ssh) + scp = Net::SCP.new(ssh.session) scp.upload!(from, to) end end @@ -132,4 +132,52 @@ module Vagrant return env.config.ssh.port end end + + class SSH + # A helper class which wraps around `Net::SSH::Connection::Session` + # in order to provide basic command error checking while still + # providing access to the actual session object. + class Session + attr_reader :session + + def initialize(session) + @session = session + end + + # Executes a given command on the SSH session and blocks until + # the command completes. This is an almost line for line copy of + # the actual `exec!` implementation, except that this + # implementation also reports `:exit_status` to the block if given. + def exec!(command, &block) + block ||= Proc.new do |ch, type, data| + ch[:result] ||= "" + ch[:result] << data if [:stdout, :stderr].include?(type) + end + + metach = session.open_channel do |channel| + channel.exec(command) do |ch, success| + raise "could not execute command: #{command.inspect}" unless success + + # Output stdout data to the block + channel.on_data do |ch2, data| + block.call(ch2, :stdout, data) + end + + # Output stderr data to the block + channel.on_extended_data do |ch2, type, data| + block.call(ch2, :stderr, data) + end + + # Output exit status information to the block + channel.on_request("exit-status") do |ch2, data| + block.call(ch2, :exit_status, data.read_long) + end + end + end + + metach.wait + metach[:result] + end + end + end end diff --git a/test/vagrant/ssh_test.rb b/test/vagrant/ssh_test.rb index 0b3b77cf8..a13b4ed60 100644 --- a/test/vagrant/ssh_test.rb +++ b/test/vagrant/ssh_test.rb @@ -118,6 +118,15 @@ class SshTest < Test::Unit::TestCase Net::SSH.expects(:start).with(@env.config.ssh.host, @env.config.ssh.username, anything).once @ssh.execute end + + should "yield an SSH session object" do + raw = mock("raw") + Net::SSH.expects(:start).yields(raw) + @ssh.execute do |ssh| + assert ssh.is_a?(Vagrant::SSH::Session) + assert_equal raw, ssh.session + end + end end context "SCPing files to the remote host" do @@ -128,8 +137,10 @@ class SshTest < Test::Unit::TestCase should "use Vagrant::SSH execute to setup an SCP connection and upload" do scp = mock("scp") ssh = mock("ssh") + sess = mock("session") + ssh.stubs(:session).returns(sess) scp.expects(:upload!).with("foo", "bar").once - Net::SCP.expects(:new).with(ssh).returns(scp).once + Net::SCP.expects(:new).with(ssh.session).returns(scp).once @ssh.expects(:execute).yields(ssh).once @ssh.upload!("foo", "bar") end diff --git a/vagrant.gemspec b/vagrant.gemspec index 94971275c..83ca2b7b8 100644 --- a/vagrant.gemspec +++ b/vagrant.gemspec @@ -9,7 +9,7 @@ Gem::Specification.new do |s| s.required_rubygems_version = Gem::Requirement.new("> 1.3.1") if s.respond_to? :required_rubygems_version= s.authors = ["Mitchell Hashimoto", "John Bender"] - s.date = %q{2010-05-27} + s.date = %q{2010-05-28} s.default_executable = %q{vagrant} s.description = %q{Vagrant is a tool for building and distributing virtualized development environments.} s.email = ["mitchell.hashimoto@gmail.com", "john.m.bender@gmail.com"]