From 8e4ab129e92cec77c3797f0d3057405c5511aaac Mon Sep 17 00:00:00 2001
From: Pengfei Ni <peni@microsoft.com>
Date: Thu, 23 Aug 2018 10:46:41 +0800
Subject: [PATCH] Fix panic when choosing zone or zones for volume

---
 pkg/volume/util/util.go      | 12 +++++++++++-
 pkg/volume/util/util_test.go | 22 ++++++++++++++++++++++
 2 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/pkg/volume/util/util.go b/pkg/volume/util/util.go
index 874a0b1381..e3987db9fc 100644
--- a/pkg/volume/util/util.go
+++ b/pkg/volume/util/util.go
@@ -520,6 +520,11 @@ func GetPath(mounter volume.Mounter) (string, error) {
 // This means that a StatefulSet's volumes (`claimname-statefulsetname-id`) will spread across available zones,
 // assuming the id values are consecutive.
 func ChooseZoneForVolume(zones sets.String, pvcName string) string {
+	// No zones available, return empty string.
+	if zones.Len() == 0 {
+		return ""
+	}
+
 	// We create the volume in a zone determined by the name
 	// Eventually the scheduler will coordinate placement into an available zone
 	hash, index := getPVCNameHashAndIndexOffset(pvcName)
@@ -541,6 +546,12 @@ func ChooseZoneForVolume(zones sets.String, pvcName string) string {
 
 // ChooseZonesForVolume is identical to ChooseZoneForVolume, but selects a multiple zones, for multi-zone disks.
 func ChooseZonesForVolume(zones sets.String, pvcName string, numZones uint32) sets.String {
+	// No zones available, return empty set.
+	replicaZones := sets.NewString()
+	if zones.Len() == 0 {
+		return replicaZones
+	}
+
 	// We create the volume in a zone determined by the name
 	// Eventually the scheduler will coordinate placement into an available zone
 	hash, index := getPVCNameHashAndIndexOffset(pvcName)
@@ -554,7 +565,6 @@ func ChooseZonesForVolume(zones sets.String, pvcName string, numZones uint32) se
 	// PVC placement (which could also e.g. avoid putting volumes in overloaded or
 	// unhealthy zones)
 	zoneSlice := zones.List()
-	replicaZones := sets.NewString()
 
 	startingIndex := index * numZones
 	for index = startingIndex; index < startingIndex+numZones; index++ {
diff --git a/pkg/volume/util/util_test.go b/pkg/volume/util/util_test.go
index 433207854d..a09d3251ae 100644
--- a/pkg/volume/util/util_test.go
+++ b/pkg/volume/util/util_test.go
@@ -647,6 +647,17 @@ func TestChooseZoneForVolume(t *testing.T) {
 			VolumeName: "medium-henley--4",
 			Expected:   "c", // hash("") + 4 == 2 mod 3
 		},
+		// Test for no zones
+		{
+			Zones:      sets.NewString(),
+			VolumeName: "medium-henley--1",
+			Expected:   "",
+		},
+		{
+			Zones:      nil,
+			VolumeName: "medium-henley--2",
+			Expected:   "",
+		},
 	}
 
 	for _, test := range tests {
@@ -992,6 +1003,17 @@ func TestChooseZonesForVolume(t *testing.T) {
 			NumZones:   3,
 			Expected:   sets.NewString("a" /* hash("henley") == 0 + 3 + 6(startingIndex) */, "b", "c"),
 		},
+		// Test for no zones
+		{
+			Zones:      sets.NewString(),
+			VolumeName: "henley-1",
+			Expected:   sets.NewString(),
+		},
+		{
+			Zones:      nil,
+			VolumeName: "henley-2",
+			Expected:   sets.NewString(),
+		},
 	}
 
 	for _, test := range tests {