From a21997db7febe69292cb981d7f9da2e6d14c0334 Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Sat, 3 Aug 2013 16:21:01 -0700
Subject: [PATCH] builder/amazon/common: refresh instance while connecting to
 SSH [GH-243]

---
 builder/amazon/common/instance.go | 21 +++++++++++++++++++--
 builder/amazon/common/ssh.go      | 22 ++++++++++++++++------
 2 files changed, 35 insertions(+), 8 deletions(-)

diff --git a/builder/amazon/common/instance.go b/builder/amazon/common/instance.go
index 1ea763e54..9a49dbaec 100644
--- a/builder/amazon/common/instance.go
+++ b/builder/amazon/common/instance.go
@@ -9,15 +9,30 @@ import (
 	"time"
 )
 
+// StateRefreshFunc is a function type used for StateChangeConf that is
+// responsible for refreshing the item being watched for a state change.
+//
+// It returns three results. `result` is any object that will be returned
+// as the final object after waiting for state change. This allows you to
+// return the final updated object, for example an EC2 instance after refreshing
+// it.
+//
+// `state` is the latest state of that object. And `err` is any error that
+// may have happened while refreshing the state.
+type StateRefreshFunc func() (result interface{}, state string, err error)
+
+// StateChangeConf is the configuration struct used for `WaitForState`.
 type StateChangeConf struct {
 	Conn      *ec2.EC2
 	Pending   []string
-	Refresh   func() (interface{}, string, error)
+	Refresh   StateRefreshFunc
 	StepState map[string]interface{}
 	Target    string
 }
 
-func InstanceStateRefreshFunc(conn *ec2.EC2, i *ec2.Instance) func() (interface{}, string, error) {
+// InstanceStateRefreshFunc returns a StateRefreshFunc that is used to watch
+// an EC2 instance.
+func InstanceStateRefreshFunc(conn *ec2.EC2, i *ec2.Instance) StateRefreshFunc {
 	return func() (interface{}, string, error) {
 		resp, err := conn.Instances([]string{i.InstanceId}, ec2.NewFilter())
 		if err != nil {
@@ -29,6 +44,8 @@ func InstanceStateRefreshFunc(conn *ec2.EC2, i *ec2.Instance) func() (interface{
 	}
 }
 
+// WaitForState watches an object and waits for it to achieve a certain
+// state.
 func WaitForState(conf *StateChangeConf) (i interface{}, err error) {
 	log.Printf("Waiting for state to become: %s", conf.Target)
 
diff --git a/builder/amazon/common/ssh.go b/builder/amazon/common/ssh.go
index adcbb5391..808a868f6 100644
--- a/builder/amazon/common/ssh.go
+++ b/builder/amazon/common/ssh.go
@@ -10,14 +10,24 @@ import (
 
 // SSHAddress returns a function that can be given to the SSH communicator
 // for determining the SSH address based on the instance DNS name.
-func SSHAddress(port int) func(map[string]interface{}) (string, error) {
+func SSHAddress(e *ec2.EC2, port int) func(map[string]interface{}) (string, error) {
 	return func(state map[string]interface{}) (string, error) {
 		var host string
-		instance := state["instance"].(*ec2.Instance)
-		if instance.DNSName != "" {
-			host = instance.DNSName
-		} else if instance.VpcId == "" {
-			host = instance.PrivateIpAddress
+		i := state["instance"].(*ec2.Instance)
+		r, err := e.Instances([]string{i.InstanceId}, ec2.NewFilter())
+		if err != nil {
+			return "", err
+		}
+
+		if len(r.Reservations) == 0 || len(r.Reservations[0].Instances) == 0 {
+			return "", fmt.Errorf("instance not found: %s", i.InstanceId)
+		}
+
+		i = &r.Reservations[0].Instances[0]
+		if i.DNSName != "" {
+			host = i.DNSName
+		} else if i.VpcId == "" {
+			host = i.PrivateIpAddress
 		} else {
 			return "", errors.New("couldn't determine IP address for instance")
 		}