From 89ca2571f3c175526b5c982fd2cd177718349be0 Mon Sep 17 00:00:00 2001
From: tsturzl <travis@metismachine.com>
Date: Tue, 12 Mar 2019 13:04:33 -0600
Subject: [PATCH] AWS zone on volume IDs Signed-off-by: Travis
 Sturzl<travis@metismachine.com>

---
 changelogs/unreleased/1274-tsturzl        |  1 +
 pkg/cloudprovider/aws/block_store.go      |  9 +++++-
 pkg/cloudprovider/aws/block_store_test.go | 37 +++++++++++++++++++++++
 3 files changed, 46 insertions(+), 1 deletion(-)
 create mode 100644 changelogs/unreleased/1274-tsturzl

diff --git a/changelogs/unreleased/1274-tsturzl b/changelogs/unreleased/1274-tsturzl
new file mode 100644
index 000000000..11f3552a1
--- /dev/null
+++ b/changelogs/unreleased/1274-tsturzl
@@ -0,0 +1 @@
+AWS EBS Volume IDs now contain AZ
diff --git a/pkg/cloudprovider/aws/block_store.go b/pkg/cloudprovider/aws/block_store.go
index 5c3c58df2..985e8044e 100644
--- a/pkg/cloudprovider/aws/block_store.go
+++ b/pkg/cloudprovider/aws/block_store.go
@@ -17,6 +17,7 @@ limitations under the License.
 package aws
 
 import (
+	"fmt"
 	"os"
 	"regexp"
 	"strings"
@@ -282,7 +283,13 @@ func (b *blockStore) SetVolumeID(unstructuredPV runtime.Unstructured, volumeID s
 		return nil, errors.New("spec.awsElasticBlockStore not found")
 	}
 
-	pv.Spec.AWSElasticBlockStore.VolumeID = volumeID
+	pvFailureDomainZone := pv.Labels["failure-domain.beta.kubernetes.io/zone"]
+
+	if len(pvFailureDomainZone) > 0 {
+		pv.Spec.AWSElasticBlockStore.VolumeID = fmt.Sprintf("aws://%s/%s", pvFailureDomainZone, volumeID)
+	} else {
+		pv.Spec.AWSElasticBlockStore.VolumeID = volumeID
+	}
 
 	res, err := runtime.DefaultUnstructuredConverter.ToUnstructured(pv)
 	if err != nil {
diff --git a/pkg/cloudprovider/aws/block_store_test.go b/pkg/cloudprovider/aws/block_store_test.go
index ecfd03120..f88292c72 100644
--- a/pkg/cloudprovider/aws/block_store_test.go
+++ b/pkg/cloudprovider/aws/block_store_test.go
@@ -85,7 +85,44 @@ func TestSetVolumeID(t *testing.T) {
 	pv.Object["spec"] = map[string]interface{}{
 		"awsElasticBlockStore": aws,
 	}
+
+	labels := map[string]interface{}{
+		"failure-domain.beta.kubernetes.io/zone": "us-east-1a",
+	}
+
+	pv.Object["metadata"] = map[string]interface{}{
+		"labels": labels,
+	}
+
 	updatedPV, err = b.SetVolumeID(pv, "vol-updated")
+
+	require.NoError(t, err)
+
+	res := new(v1.PersistentVolume)
+	require.NoError(t, runtime.DefaultUnstructuredConverter.FromUnstructured(updatedPV.UnstructuredContent(), res))
+	require.NotNil(t, res.Spec.AWSElasticBlockStore)
+	assert.Equal(t, "aws://us-east-1a/vol-updated", res.Spec.AWSElasticBlockStore.VolumeID)
+}
+
+func TestSetVolumeIDNoZone(t *testing.T) {
+	b := &blockStore{}
+
+	pv := &unstructured.Unstructured{
+		Object: map[string]interface{}{},
+	}
+
+	// missing spec.awsElasticBlockStore -> error
+	updatedPV, err := b.SetVolumeID(pv, "vol-updated")
+	require.Error(t, err)
+
+	// happy path
+	aws := map[string]interface{}{}
+	pv.Object["spec"] = map[string]interface{}{
+		"awsElasticBlockStore": aws,
+	}
+
+	updatedPV, err = b.SetVolumeID(pv, "vol-updated")
+
 	require.NoError(t, err)
 
 	res := new(v1.PersistentVolume)