Add checks for compute environment flags and url checks

Author: Sneha-at

cmd/gce-pd-csi-driver/main.go

@@ -17,7 +17,9 @@ package main
 
 import (
 	"context"
+	"errors"
 	"flag"
+	"fmt"
 	"math/rand"
 	"os"
 	"runtime"
@@ -67,13 +69,12 @@ var (
 
 	maxConcurrentFormatAndMount = flag.Int("max-concurrent-format-and-mount", 1, "If set then format and mount operations are serialized on each node. This is stronger than max-concurrent-format as it includes fsck and other mount operations")
 	formatAndMountTimeout       = flag.Duration("format-and-mount-timeout", 1*time.Minute, "The maximum duration of a format and mount operation before another such operation will be started. Used only if --serialize-format-and-mount")
-	computeEnvironment          = flag.String("compute-environment", "prod", "Sets the compute environment")
+	fallbackRequisiteZonesFlag  = flag.String("fallback-requisite-zones", "", "Comma separated list of requisite zones that will be used if there are not sufficient zones present in requisite topologies when provisioning a disk")
 
-	fallbackRequisiteZonesFlag = flag.String("fallback-requisite-zones", "", "Comma separated list of requisite zones that will be used if there are not sufficient zones present in requisite topologies when provisioning a disk")
-
-	enableStoragePoolsFlag = flag.Bool("enable-storage-pools", false, "If set to true, the CSI Driver will allow volumes to be provisioned in Storage Pools")
-
-	version string
+	enableStoragePoolsFlag                    = flag.Bool("enable-storage-pools", false, "If set to true, the CSI Driver will allow volumes to be provisioned in Storage Pools")
+	computeEnvironment        gce.Environment = "production"
+	version                   string
+	allowedComputeEnvironment = []string{"staging", "production"}
 )
 
 const (
@@ -154,10 +155,13 @@ func handle() {
 	// Initilaize requisite zones
 	fallbackRequisiteZones := strings.Split(*fallbackRequisiteZonesFlag, ",")
 
+	enumFlag(&computeEnvironment, "computeEnvironment", allowedComputeEnvironment, "Operating compute environment")
+	flag.Parse()
+
 	// Initialize requirements for the controller service
 	var controllerServer *driver.GCEControllerServer
 	if *runControllerService {
-		cloudProvider, err := gce.CreateCloudProvider(ctx, version, *cloudConfigFilePath, *computeEndpoint, *computeEnvironment)
+		cloudProvider, err := gce.CreateCloudProvider(ctx, version, *cloudConfigFilePath, *computeEndpoint, computeEnvironment)
 		if err != nil {
 			klog.Fatalf("Failed to get cloud provider: %v", err.Error())
 		}
@@ -206,3 +210,17 @@ func handle() {
 
 	gceDriver.Run(*endpoint, *grpcLogCharCap, *enableOtelTracing)
 }
+
+func enumFlag(target *gce.Environment, name string, allowedComputeEnvironment []string, usage string) {
+	flag.Func(name, usage, func(flagValue string) error {
+		for _, allowedValue := range allowedComputeEnvironment {
+			if flagValue == allowedValue {
+				*target = gce.Environment(flagValue)
+				return nil
+			}
+		}
+		errMsg := fmt.Sprintf(`must be one of %v`, allowedComputeEnvironment)
+		return errors.New(errMsg)
+	})
+
+}

pkg/gce-cloud-provider/compute/gce.go

@@ -38,6 +38,9 @@ import (
 	"k8s.io/klog/v2"
 )
 
+type Environment string
+type Version string
+
 const (
 	TokenURL                        = "https://accounts.google.com/o/oauth2/token"
 	diskSourceURITemplateSingleZone = "projects/%s/zones/%s/disks/%s"       // {gce.projectID}/zones/{disk.Zone}/disks/{disk.Name}"
@@ -47,25 +50,23 @@ const (
 
 	regionURITemplate = "projects/%s/regions/%s"
 
-	replicaZoneURITemplateSingleZone = "projects/%s/zones/%s" // {gce.projectID}/zones/{disk.Zone}
-	versionV1                        = "v1"
-	versionBeta                      = "beta"
-	versionAlpha                     = "alpha"
-	googleEnv                        = "googleapis"
+	replicaZoneURITemplateSingleZone             = "projects/%s/zones/%s" // {gce.projectID}/zones/{disk.Zone}
+	versionV1                        Version     = "v1"
+	versionBeta                      Version     = "beta"
+	versionAlpha                     Version     = "alpha"
+	environmentStaging               Environment = "staging"
 )
 
-var computeVersionMap = map[string]map[string]map[string]string{
-	googleEnv: {
-		"prod": {
-			versionV1:    "compute/v1/",
-			versionBeta:  "compute/beta/",
-			versionAlpha: "compute/alpha/",
-		},
-		"staging": {
-			versionV1:    "compute/staging_v1/",
-			versionBeta:  "compute/staging_beta/",
-			versionAlpha: "compute/staging_alpha/",
-		},
+var computeVersionMap = map[Environment]map[Version]string{
+	"prod": {
+		versionV1:    "compute/v1/",
+		versionBeta:  "compute/beta/",
+		versionAlpha: "compute/alpha/",
+	},
+	"staging": {
+		versionV1:    "compute/staging_v1/",
+		versionBeta:  "compute/staging_beta/",
+		versionAlpha: "compute/staging_alpha/",
 	},
 }
 
@@ -92,7 +93,7 @@ type ConfigGlobal struct {
 	Zone      string `gcfg:"zone"`
 }
 
-func CreateCloudProvider(ctx context.Context, vendorVersion string, configPath string, computeEndpoint string, computeEnvironment string) (*CloudProvider, error) {
+func CreateCloudProvider(ctx context.Context, vendorVersion string, configPath string, computeEndpoint string, computeEnvironment Environment) (*CloudProvider, error) {
 	configFile, err := readConfig(configPath)
 	if err != nil {
 		return nil, err
@@ -187,7 +188,7 @@ func readConfig(configPath string) (*ConfigFile, error) {
 	return cfg, nil
 }
 
-func createAlphaCloudService(ctx context.Context, vendorVersion string, tokenSource oauth2.TokenSource, computeEndpoint string, computeEnvironment string) (*computealpha.Service, error) {
+func createAlphaCloudService(ctx context.Context, vendorVersion string, tokenSource oauth2.TokenSource, computeEndpoint string, computeEnvironment Environment) (*computealpha.Service, error) {
 	computeOpts, err := getComputeVersion(ctx, tokenSource, computeEndpoint, computeEnvironment, versionAlpha)
 	if err != nil {
 		klog.Errorf("Failed to get compute endpoint: %s", err)
@@ -200,7 +201,7 @@ func createAlphaCloudService(ctx context.Context, vendorVersion string, tokenSou
 	return service, nil
 }
 
-func createBetaCloudService(ctx context.Context, vendorVersion string, tokenSource oauth2.TokenSource, computeEndpoint string, computeEnvironment string) (*computebeta.Service, error) {
+func createBetaCloudService(ctx context.Context, vendorVersion string, tokenSource oauth2.TokenSource, computeEndpoint string, computeEnvironment Environment) (*computebeta.Service, error) {
 	computeOpts, err := getComputeVersion(ctx, tokenSource, computeEndpoint, computeEnvironment, versionBeta)
 	if err != nil {
 		klog.Errorf("Failed to get compute endpoint: %s", err)
@@ -213,7 +214,7 @@ func createBetaCloudService(ctx context.Context, vendorVersion string, tokenSour
 	return service, nil
 }
 
-func createCloudService(ctx context.Context, vendorVersion string, tokenSource oauth2.TokenSource, computeEndpoint string, computeEnvironment string) (*compute.Service, error) {
+func createCloudService(ctx context.Context, vendorVersion string, tokenSource oauth2.TokenSource, computeEndpoint string, computeEnvironment Environment) (*compute.Service, error) {
 	computeOpts, err := getComputeVersion(ctx, tokenSource, computeEndpoint, computeEnvironment, versionV1)
 	if err != nil {
 		klog.Errorf("Failed to get compute endpoint: %s", err)
@@ -226,20 +227,21 @@ func createCloudService(ctx context.Context, vendorVersion string, tokenSource o
 	return service, nil
 }
 
-func getComputeVersion(ctx context.Context, tokenSource oauth2.TokenSource, computeEndpoint string, computeEnvironment string, computeVersion string) ([]option.ClientOption, error) {
+func getComputeVersion(ctx context.Context, tokenSource oauth2.TokenSource, computeEndpoint string, computeEnvironment Environment, computeVersion Version) ([]option.ClientOption, error) {
 	client, err := newOauthClient(ctx, tokenSource)
 	if err != nil {
 		return nil, err
 	}
-	computeEnvironmentSuffix, ok := computeVersionMap[googleEnv][computeEnvironment][computeVersion]
-	if !ok {
-		return nil, errors.New("Unable to fetch compute endpoint")
-	}
+	computeEnvironmentSuffix := getPath(computeEnvironment, computeVersion)
 	computeOpts := []option.ClientOption{option.WithHTTPClient(client)}
+
 	if computeEndpoint != "" {
-		endpoint := fmt.Sprintf("%s%s", computeEndpoint, computeEnvironmentSuffix)
-		klog.Infof("Got compute endpoint %s", endpoint)
-		_, err := url.ParseRequestURI(endpoint)
+		computeURL, err := url.ParseRequestURI(computeEndpoint)
+		if err != nil {
+			return nil, err
+		}
+		endpoint := computeURL.JoinPath(computeEnvironmentSuffix).String()
+		_, err = url.ParseRequestURI(endpoint)
 		if err != nil {
 			klog.Fatalf("Error parsing compute endpoint %s", endpoint)
 		}
@@ -248,6 +250,14 @@ func getComputeVersion(ctx context.Context, tokenSource oauth2.TokenSource, comp
 	return computeOpts, nil
 }
 
+func getPath(env Environment, version Version) string {
+	prefix := ""
+	if env == environmentStaging {
+		prefix = fmt.Sprintf("%s_", env)
+	}
+	return fmt.Sprintf("compute/%s%s/", prefix, version)
+}
+
 func newOauthClient(ctx context.Context, tokenSource oauth2.TokenSource) (*http.Client, error) {
 	if err := wait.PollImmediate(5*time.Second, 30*time.Second, func() (bool, error) {
 		if _, err := tokenSource.Token(); err != nil {

test/e2e/tests/single_zone_e2e_test.go

@@ -1296,13 +1296,12 @@ var _ = Describe("GCE PD CSI Driver", func() {
 		klog.Infof("Creating new driver and client for node %s\n", i.GetName())
 
 		// Create new driver and client w/ invalid endpoint
-		tcInvalid, err := testutils.GCEClientAndDriverSetup(i, "invalid-string")
-		if err != nil {
-			klog.Fatalf("Failed to set up Test Context for instance %v: %v", i.GetName(), err)
+		computeEndpoint := "invalid-string"
+		tcInvalid, err := testutils.GCEClientAndDriverSetup(i, computeEndpoint)
+		if tcInvalid != nil {
+			klog.Fatalf("Driver setup with incorrect compute %v: %v", i.GetName(), computeEndpoint)
 		}
-
-		_, err = tcInvalid.Client.ListVolumes()
-		Expect(err.Error()).To(ContainSubstring("no such host"), "expected error when passed invalid compute url")
+		Expect(err.Error()).To(ContainSubstring("failed start driver"), "expected error when passed invalid compute url")
 
 		// Create new driver and client w/ valid, passed-in endpoint
 		tcValid, err := testutils.GCEClientAndDriverSetup(i, "https://compute.googleapis.com")