From abeb7ad9f2b9820cbfd1b844dfe9db854c6a5227 Mon Sep 17 00:00:00 2001
From: andyzhangx <xiazhang@microsoft.com>
Date: Fri, 23 Feb 2024 14:17:02 +0000
Subject: [PATCH] feat: allow multiple subnets in updating service endpoints

---
 docs/driver-parameters.md         |  2 +-
 pkg/blob/controllerserver.go      | 17 ++++++++++++-----
 pkg/blob/controllerserver_test.go | 24 ++++++++++++++++++++++++
 3 files changed, 37 insertions(+), 6 deletions(-)

diff --git a/docs/driver-parameters.md b/docs/driver-parameters.md
index 74aadc4e8..d03af01ea 100644
--- a/docs/driver-parameters.md
+++ b/docs/driver-parameters.md
@@ -47,7 +47,7 @@ isHnsEnabled | enable `Hierarchical namespace` for Azure DataLake storage accoun
 mountPermissions | mounted folder permissions. The default is `0777`, if set as `0`, driver will not perform `chmod` after mount | `0777` | No |
 vnetResourceGroup | specify vnet resource group where virtual network is | existing resource group name | No | if empty, driver will use the `vnetResourceGroup` value in azure cloud config file
 vnetName | virtual network name | existing virtual network name | No | if empty, driver will use the `vnetName` value in azure cloud config file
-subnetName | subnet name | existing subnet name of the agent node | No | if empty, driver will use the `subnetName` value in azure cloud config file
+subnetName | subnet name | existing subnet name(s) of the agent node, if you want to update service endpoints on multiple subnets, separate them using a comma (`,`) | No | if empty, driver will use the `subnetName` value in azure cloud config file
 softDeleteBlobs | Enable [soft delete for blobs](https://learn.microsoft.com/en-us/azure/storage/blobs/soft-delete-blob-overview), specify the days to retain deleted blobs | "7" | No | Soft Delete Blobs is disabled if empty
 softDeleteContainers | Enable [soft delete for containers](https://learn.microsoft.com/en-us/azure/storage/blobs/soft-delete-container-overview), specify the days to retain deleted containers | "7" | No | Soft Delete Containers is disabled if empty
 enableBlobVersioning | Enable [blob versioning](https://learn.microsoft.com/en-us/azure/storage/blobs/versioning-overview), can't enabled when `protocol` is `nfs` or `isHnsEnabled` is `true` | `true`,`false` | No | versioning for blobs is disabled if empty
diff --git a/pkg/blob/controllerserver.go b/pkg/blob/controllerserver.go
index 009396878..5dbe00477 100644
--- a/pkg/blob/controllerserver.go
+++ b/pkg/blob/controllerserver.go
@@ -274,6 +274,9 @@ func (d *Driver) CreateVolume(ctx context.Context, req *csi.CreateVolumeRequest)
 
 	enableHTTPSTrafficOnly := true
 	if strings.EqualFold(networkEndpointType, privateEndpoint) {
+		if strings.Contains(subnetName, ",") {
+			return nil, status.Errorf(codes.InvalidArgument, "subnetName(%s) can only contain one subnet for private endpoint", subnetName)
+		}
 		createPrivateEndpoint = pointer.BoolPtr(true)
 	}
 	accountKind := string(armstorage.KindStorageV2)
@@ -284,11 +287,15 @@ func (d *Driver) CreateVolume(ctx context.Context, req *csi.CreateVolumeRequest)
 		storeAccountKey = false
 		if !pointer.BoolDeref(createPrivateEndpoint, false) {
 			// set VirtualNetworkResourceIDs for storage account firewall setting
-			vnetResourceID := d.getSubnetResourceID(vnetResourceGroup, vnetName, subnetName)
-			klog.V(2).Infof("set vnetResourceID(%s) for NFS protocol", vnetResourceID)
-			vnetResourceIDs = []string{vnetResourceID}
-			if err := d.updateSubnetServiceEndpoints(ctx, vnetResourceGroup, vnetName, subnetName); err != nil {
-				return nil, status.Errorf(codes.Internal, "update service endpoints failed with error: %v", err)
+			subnets := strings.Split(subnetName, ",")
+			for _, subnet := range subnets {
+				subnet = strings.TrimSpace(subnet)
+				vnetResourceID := d.getSubnetResourceID(vnetResourceGroup, vnetName, subnet)
+				klog.V(2).Infof("set vnetResourceID(%s) for NFS protocol", vnetResourceID)
+				vnetResourceIDs = []string{vnetResourceID}
+				if err := d.updateSubnetServiceEndpoints(ctx, vnetResourceGroup, vnetName, subnet); err != nil {
+					return nil, status.Errorf(codes.Internal, "update service endpoints failed with error: %v", err)
+				}
 			}
 		}
 	}
diff --git a/pkg/blob/controllerserver_test.go b/pkg/blob/controllerserver_test.go
index d7961cc2e..3009ab856 100644
--- a/pkg/blob/controllerserver_test.go
+++ b/pkg/blob/controllerserver_test.go
@@ -407,6 +407,30 @@ func TestCreateVolume(t *testing.T) {
 				}
 			},
 		},
+		{
+			name: "invalid privateEndpoint and subnetName combination",
+			testFunc: func(t *testing.T) {
+				d := NewFakeDriver()
+				mp := map[string]string{
+					networkEndpointTypeField: "privateendpoint",
+					subnetNameField:          "subnet1,subnet2",
+				}
+				req := &csi.CreateVolumeRequest{
+					Name:               "unit-test",
+					VolumeCapabilities: stdVolumeCapabilities,
+					Parameters:         mp,
+				}
+				d.Cap = []*csi.ControllerServiceCapability{
+					controllerServiceCapability,
+				}
+
+				expectedErr := status.Errorf(codes.InvalidArgument, "subnetName(subnet1,subnet2) can only contain one subnet for private endpoint")
+				_, err := d.CreateVolume(context.Background(), req)
+				if !reflect.DeepEqual(err, expectedErr) {
+					t.Errorf("Unexpected error: %v", err)
+				}
+			},
+		},
 		{
 			name: "NFS not supported by cross subscription",
 			testFunc: func(t *testing.T) {