From f30cdf3eea9faf188b5552b9ddc928df491210fa Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Fri, 28 May 2010 14:19:35 -0700
Subject: [PATCH] New wrapper SSH session class to emit exit status as well.
 Will be used in the future to raise exceptions for bad exit status. For now,
 works as a drop-in replacement.

---
 lib/vagrant/ssh.rb       | 52 ++++++++++++++++++++++++++++++++++++++--
 test/vagrant/ssh_test.rb | 13 +++++++++-
 vagrant.gemspec          |  2 +-
 3 files changed, 63 insertions(+), 4 deletions(-)

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"]