[KNIFE_EC2-83] Add SSH gateway and custom connection attribute support

Merge branch 'wpeterson/appcloud'

Conflicts:
	lib/chef/knife/ec2_server_create.rb

lib/chef/knife/ec2_server_create.rb

@@ -104,6 +104,13 @@ class Ec2ServerCreate < Knife
         :default => "22",
         :proc => Proc.new { |key| Chef::Config[:knife][:ssh_port] = key }
 
+      option :ssh_gateway,
+        :short => "-w GATEWAY",
+        :long => "--ssh-gateway GATEWAY",
+        :description => "The ssh gateway server",
+        :proc => Proc.new { |key| Chef::Config[:knife][:ssh_gateway] = key }
+
+
       option :identity_file,
         :short => "-i IDENTITY_FILE",
         :long => "--identity-file IDENTITY_FILE",
@@ -184,8 +191,14 @@ class Ec2ServerCreate < Knife
         :proc => lambda { |o| o.split(/[\s,]+/) },
         :default => []
 
-      def tcp_test_ssh(hostname)
-        tcp_socket = TCPSocket.new(hostname, config[:ssh_port])
+      option :server_connect_attribute,
+        :long => "--server-connect-attribute ATTRIBUTE",
+        :short => "-c ATTRIBUTE",
+        :description => "The EC2 server attribute to use for SSH connection",
+        :default => nil
+
+      def tcp_test_ssh(hostname, ssh_port)
+        tcp_socket = TCPSocket.new(hostname, ssh_port)
         readable = IO.select([tcp_socket], nil, nil, 5)
         if readable
           Chef::Log.debug("sshd accepting connections on #{hostname}, banner is #{tcp_socket.gets}")
@@ -194,23 +207,10 @@ def tcp_test_ssh(hostname)
         else
           false
         end
-      rescue SocketError
-        sleep 2
-        false
-      rescue Errno::ETIMEDOUT
-        false
-      rescue Errno::EPERM
-        false
-      rescue Errno::ECONNREFUSED
+      rescue SocketError, Errno::ECONNREFUSED, Errno::EHOSTUNREACH, Errno::ENETUNREACH, IOError
         sleep 2
         false
-      # This happens on EC2 quite often
-      rescue Errno::EHOSTUNREACH
-        sleep 2
-        false
-      # This happens on EC2 sometimes
-      rescue Errno::ENETUNREACH
-        sleep 2
+      rescue Errno::EPERM, Errno::ETIMEDOUT
         false
       ensure
         tcp_socket && tcp_socket.close
@@ -274,14 +274,9 @@ def run
 
         print "\n#{ui.color("Waiting for sshd", :magenta)}"
 
-        fqdn = vpc_mode? ? @server.private_ip_address : @server.dns_name
-
-        print(".") until tcp_test_ssh(fqdn) {
-          sleep @initial_sleep_delay ||= (vpc_mode? ? 40 : 10)
-          puts("done")
-        }
+        wait_for_sshd(ssh_connect_host)
 
-        bootstrap_for_node(@server,fqdn).run
+        bootstrap_for_node(@server,ssh_connect_host).run
 
         puts "\n"
         msg_pair("Instance ID", @server.id)
@@ -323,12 +318,13 @@ def run
         msg_pair("JSON Attributes",config[:json_attributes]) unless !config[:json_attributes] || config[:json_attributes].empty?
       end
 
-      def bootstrap_for_node(server,fqdn)
+      def bootstrap_for_node(server,ssh_host)
         bootstrap = Chef::Knife::Bootstrap.new
-        bootstrap.name_args = [fqdn]
+        bootstrap.name_args = [ssh_host]
         bootstrap.config[:run_list] = locate_config_value(:run_list) || []
         bootstrap.config[:ssh_user] = config[:ssh_user]
         bootstrap.config[:ssh_port] = config[:ssh_port]
+        bootstrap.config[:ssh_gateway] = config[:ssh_gateway]
         bootstrap.config[:identity_file] = config[:identity_file]
         bootstrap.config[:chef_node_name] = locate_config_value(:chef_node_name) || server.id
         bootstrap.config[:prerelease] = config[:prerelease]
@@ -429,6 +425,49 @@ def create_server_def
 
         server_def
       end
+
+      def wait_for_sshd(hostname)
+        config[:ssh_gateway] ? wait_for_tunnelled_sshd(hostname) : wait_for_direct_sshd(hostname, config[:ssh_port])
+      end
+
+      def wait_for_tunnelled_sshd(hostname)
+        print(".")
+        print(".") until tunnel_test_ssh(ssh_connect_host) {
+          sleep @initial_sleep_delay ||= (vpc_mode? ? 40 : 10)
+          puts("done")
+        }
+      end
+
+      def tunnel_test_ssh(hostname, &block)
+        gw_host, gw_user = config[:ssh_gateway].split('@').reverse
+        gw_host, gw_port = gw_host.split(':')
+        gateway = Net::SSH::Gateway.new(gw_host, gw_user, :port => gw_port || 22)
+        status = false
+        gateway.open(hostname, config[:ssh_port]) do |local_tunnel_port|
+          status = tcp_test_ssh('localhost', local_tunnel_port, &block)
+        end
+        status
+      rescue SocketError, Errno::ECONNREFUSED, Errno::EHOSTUNREACH, Errno::ENETUNREACH, IOError
+        sleep 2
+        false
+      rescue Errno::EPERM, Errno::ETIMEDOUT
+        false
+      end
+
+      def wait_for_direct_sshd(hostname, ssh_port)
+        print(".") until tcp_test_ssh(ssh_connect_host, ssh_port) {
+          sleep @initial_sleep_delay ||= (vpc_mode? ? 40 : 10)
+          puts("done")
+        }
+      end
+
+      def ssh_connect_host
+        @ssh_connect_host ||= if config[:server_connect_attribute]
+          server.send(config[:server_connect_attribute])
+        else
+          vpc_mode? ? server.private_ip_address : server.dns_name
+        end
+      end
     end
   end
 end

spec/unit/ec2_server_create_spec.rb

@@ -128,6 +128,7 @@
       @knife_ec2_create.config[:ssh_user] = "ubuntu"
       @knife_ec2_create.config[:identity_file] = "~/.ssh/aws-key.pem"
       @knife_ec2_create.config[:ssh_port] = 22
+      @knife_ec2_create.config[:ssh_gateway] = 'bastion.host.com'
       @knife_ec2_create.config[:chef_node_name] = "blarf"
       @knife_ec2_create.config[:template_file] = '~/.chef/templates/my-bootstrap.sh.erb'
       @knife_ec2_create.config[:distro] = 'ubuntu-10.04-magic-sparkles'
@@ -153,6 +154,10 @@
       @bootstrap.config[:ssh_user].should == 'ubuntu'
     end
 
+    it "configures the bootstrap to use the correct ssh_gateway host" do
+      @bootstrap.config[:ssh_gateway].should == 'bastion.host.com'
+    end
+
     it "configures the bootstrap to use the correct ssh identity file" do
       @bootstrap.config[:identity_file].should == "~/.ssh/aws-key.pem"
     end
@@ -267,4 +272,35 @@
     end
   end
 
+  describe "ssh_connect_host" do
+    before(:each) do
+      @new_ec2_server.stub!(
+        :dns_name => 'public_name',
+        :private_ip_address => 'private_ip',
+        :custom => 'custom'
+      )
+      @knife_ec2_create.stub!(:server => @new_ec2_server)
+    end
+
+    describe "by default" do
+      it 'should use public dns name' do
+        @knife_ec2_create.ssh_connect_host.should == 'public_name'
+      end
+    end
+
+    describe "with vpc_mode?" do
+      it 'should use private ip' do
+        @knife_ec2_create.stub!(:vpc_mode? => true)
+        @knife_ec2_create.ssh_connect_host.should == 'private_ip'
+      end
+
+    end
+
+    describe "with custom server attribute" do
+      it 'should use custom server attribute' do
+        @knife_ec2_create.config[:server_connect_attribute] = 'custom'
+        @knife_ec2_create.ssh_connect_host.should == 'custom'
+      end
+    end
+  end
 end