From 31b8f04bda7234c37bd01589bdb5f0c3d7a10c3c Mon Sep 17 00:00:00 2001
From: Clint Shryock <clint@ctshryock.com>
Date: Fri, 9 Oct 2015 11:49:36 -0500
Subject: [PATCH] provider/aws: Migrate KeyPair to version 1

---
 .../providers/aws/resource_aws_key_pair.go    | 13 +++++
 .../aws/resource_aws_key_pair_migrate.go      | 38 +++++++++++++
 .../aws/resource_aws_key_pair_migrate_test.go | 55 +++++++++++++++++++
 3 files changed, 106 insertions(+)
 create mode 100644 builtin/providers/aws/resource_aws_key_pair_migrate.go
 create mode 100644 builtin/providers/aws/resource_aws_key_pair_migrate_test.go

diff --git a/builtin/providers/aws/resource_aws_key_pair.go b/builtin/providers/aws/resource_aws_key_pair.go
index e747fbfc5..0d6c51fcf 100644
--- a/builtin/providers/aws/resource_aws_key_pair.go
+++ b/builtin/providers/aws/resource_aws_key_pair.go
@@ -2,6 +2,7 @@ package aws
 
 import (
 	"fmt"
+	"strings"
 
 	"github.com/hashicorp/terraform/helper/resource"
 	"github.com/hashicorp/terraform/helper/schema"
@@ -18,6 +19,9 @@ func resourceAwsKeyPair() *schema.Resource {
 		Update: nil,
 		Delete: resourceAwsKeyPairDelete,
 
+		SchemaVersion: 1,
+		MigrateState:  resourceAwsKeyPairMigrateState,
+
 		Schema: map[string]*schema.Schema{
 			"key_name": &schema.Schema{
 				Type:     schema.TypeString,
@@ -29,6 +33,14 @@ func resourceAwsKeyPair() *schema.Resource {
 				Type:     schema.TypeString,
 				Required: true,
 				ForceNew: true,
+				StateFunc: func(v interface{}) string {
+					switch v.(type) {
+					case string:
+						return strings.TrimSpace(v.(string))
+					default:
+						return ""
+					}
+				},
 			},
 			"fingerprint": &schema.Schema{
 				Type:     schema.TypeString,
@@ -45,6 +57,7 @@ func resourceAwsKeyPairCreate(d *schema.ResourceData, meta interface{}) error {
 	if keyName == "" {
 		keyName = resource.UniqueId()
 	}
+
 	publicKey := d.Get("public_key").(string)
 	req := &ec2.ImportKeyPairInput{
 		KeyName:           aws.String(keyName),
diff --git a/builtin/providers/aws/resource_aws_key_pair_migrate.go b/builtin/providers/aws/resource_aws_key_pair_migrate.go
new file mode 100644
index 000000000..0d56123aa
--- /dev/null
+++ b/builtin/providers/aws/resource_aws_key_pair_migrate.go
@@ -0,0 +1,38 @@
+package aws
+
+import (
+	"fmt"
+	"log"
+	"strings"
+
+	"github.com/hashicorp/terraform/terraform"
+)
+
+func resourceAwsKeyPairMigrateState(
+	v int, is *terraform.InstanceState, meta interface{}) (*terraform.InstanceState, error) {
+	switch v {
+	case 0:
+		log.Println("[INFO] Found AWS Key Pair State v0; migrating to v1")
+		return migrateKeyPairStateV0toV1(is)
+	default:
+		return is, fmt.Errorf("Unexpected schema version: %d", v)
+	}
+
+	return is, nil
+}
+
+func migrateKeyPairStateV0toV1(is *terraform.InstanceState) (*terraform.InstanceState, error) {
+	if is.Empty() {
+		log.Println("[DEBUG] Empty InstanceState; nothing to migrate.")
+		return is, nil
+	}
+
+	log.Printf("[DEBUG] Attributes before migration: %#v", is.Attributes)
+
+	// replace public_key with a stripped version, removing `\n` from the end
+	// see https://github.com/hashicorp/terraform/issues/3455
+	is.Attributes["public_key"] = strings.TrimSpace(is.Attributes["public_key"])
+
+	log.Printf("[DEBUG] Attributes after migration: %#v", is.Attributes)
+	return is, nil
+}
diff --git a/builtin/providers/aws/resource_aws_key_pair_migrate_test.go b/builtin/providers/aws/resource_aws_key_pair_migrate_test.go
new file mode 100644
index 000000000..825d3c40f
--- /dev/null
+++ b/builtin/providers/aws/resource_aws_key_pair_migrate_test.go
@@ -0,0 +1,55 @@
+package aws
+
+import (
+	"testing"
+
+	"github.com/hashicorp/terraform/terraform"
+)
+
+func TestAWSKeyPairMigrateState(t *testing.T) {
+	cases := map[string]struct {
+		StateVersion int
+		ID           string
+		Attributes   map[string]string
+		Expected     string
+		Meta         interface{}
+	}{
+		"v0_1": {
+			StateVersion: 0,
+			ID:           "tf-testing-file",
+			Attributes: map[string]string{
+				"fingerprint": "1d:cd:46:31:a9:4a:e0:06:8a:a1:22:cb:3b:bf:8e:42",
+				"key_name":    "tf-testing-file",
+				"public_key":  "ssh-rsa AAAAB3NzaC1yc2EAAAABIwAAAQEA4LBtwcFsQAYWw1cnOwRTZCJCzPSzq0dl3== ctshryock",
+			},
+			Expected: "ssh-rsa AAAAB3NzaC1yc2EAAAABIwAAAQEA4LBtwcFsQAYWw1cnOwRTZCJCzPSzq0dl3== ctshryock",
+		},
+		"v0_2": {
+			StateVersion: 0,
+			ID:           "tf-testing-file",
+			Attributes: map[string]string{
+				"fingerprint": "1d:cd:46:31:a9:4a:e0:06:8a:a1:22:cb:3b:bf:8e:42",
+				"key_name":    "tf-testing-file",
+				"public_key":  "ssh-rsa AAAAB3NzaC1yc2EAAAABIwAAAQEA4LBtwcFsQAYWw1cnOwRTZCJCzPSzq0dl3== ctshryock\n",
+			},
+			Expected: "ssh-rsa AAAAB3NzaC1yc2EAAAABIwAAAQEA4LBtwcFsQAYWw1cnOwRTZCJCzPSzq0dl3== ctshryock",
+		},
+	}
+
+	for tn, tc := range cases {
+		is := &terraform.InstanceState{
+			ID:         tc.ID,
+			Attributes: tc.Attributes,
+		}
+		is, err := resourceAwsKeyPairMigrateState(
+			tc.StateVersion, is, tc.Meta)
+
+		if err != nil {
+			t.Fatalf("bad: %s, err: %#v", tn, err)
+		}
+
+		if is.Attributes["public_key"] != tc.Expected {
+			t.Fatalf("Bad public_key migration: %s\n\n expected: %s", is.Attributes["public_key"], tc.Expected)
+		}
+	}
+}