Add url check during initialization

Author: Sneha-at

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

@@ -21,6 +21,7 @@ import (
 	"flag"
 	"fmt"
 	"math/rand"
+	"net/url"
 	"os"
 	"runtime"
 	"strings"
@@ -40,7 +41,6 @@ import (
 var (
 	cloudConfigFilePath  = flag.String("cloud-config", "", "Path to GCE cloud provider config")
 	endpoint             = flag.String("endpoint", "unix:/tmp/csi.sock", "CSI endpoint")
-	computeEndpoint      = flag.String("compute-endpoint", "", "If set, used as the endpoint for the GCE API.")
 	runControllerService = flag.Bool("run-controller-service", true, "If set to false then the CSI driver does not activate its controller service (default: true)")
 	runNodeService       = flag.Bool("run-node-service", true, "If set to false then the CSI driver does not activate its node service (default: true)")
 	httpEndpoint         = flag.String("http-endpoint", "", "The TCP network address where the prometheus metrics endpoint will listen (example: `:8080`). The default is empty string, which means metrics endpoint is disabled.")
@@ -71,8 +71,9 @@ var (
 	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")
 
 	computeEnvironment        gce.Environment = "production"
+	computeEndpoint           url.URL
 	version                   string
-	allowedComputeEnvironment = []string{"staging", "production"}
+	allowedComputeEnvironment = []gce.Environment{gce.EnvironmentStaging, gce.EnvironmentProduction}
 )
 
 const (
@@ -86,6 +87,7 @@ func init() {
 	// Use V(5) for GCE Cloud Provider Call informational logging
 	// Use V(6) for extra repeated/polling information
 	enumFlag(&computeEnvironment, "compute-environment", allowedComputeEnvironment, "Operating compute environment")
+	urlFlag(&computeEndpoint, "compute-endpoint", "Compute endpoint")
 	klog.InitFlags(flag.CommandLine)
 	flag.Set("logtostderr", "true")
 }
@@ -141,7 +143,7 @@ func handle() {
 	// 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())
 		}
@@ -191,10 +193,10 @@ func handle() {
 	gceDriver.Run(*endpoint, *grpcLogCharCap)
 }
 
-func enumFlag(target *gce.Environment, name string, allowedComputeEnvironment []string, usage string) {
+func enumFlag(target *gce.Environment, name string, allowedComputeEnvironment []gce.Environment, usage string) {
 	flag.Func(name, usage, func(flagValue string) error {
 		for _, allowedValue := range allowedComputeEnvironment {
-			if flagValue == allowedValue {
+			if gce.Environment(flagValue) == allowedValue {
 				*target = gce.Environment(flagValue)
 				return nil
 			}
@@ -204,3 +206,14 @@ func enumFlag(target *gce.Environment, name string, allowedComputeEnvironment []
 	})
 
 }
+
+func urlFlag(target *url.URL, name string, usage string) {
+	flag.Func(name, usage, func(flagValue string) error {
+		computeURL, err := url.ParseRequestURI(flagValue)
+		if err == nil {
+			*target = *computeURL
+			return nil
+		}
+		return err
+	})
+}

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

@@ -54,7 +54,8 @@ const (
 	versionV1                        Version     = "v1"
 	versionBeta                      Version     = "beta"
 	versionAlpha                     Version     = "alpha"
-	environmentStaging               Environment = "staging"
+	EnvironmentStaging               Environment = "staging"
+	EnvironmentProduction            Environment = "production"
 )
 
 type CloudProvider struct {
@@ -79,7 +80,7 @@ type ConfigGlobal struct {
 	Zone      string `gcfg:"zone"`
 }
 
-func CreateCloudProvider(ctx context.Context, vendorVersion string, configPath string, computeEndpoint string, computeEnvironment Environment) (*CloudProvider, error) {
+func CreateCloudProvider(ctx context.Context, vendorVersion string, configPath string, computeEndpoint url.URL, computeEnvironment Environment) (*CloudProvider, error) {
 	configFile, err := readConfig(configPath)
 	if err != nil {
 		return nil, err
@@ -173,7 +174,7 @@ func readConfig(configPath string) (*ConfigFile, error) {
 	return cfg, nil
 }
 
-func createAlphaCloudService(ctx context.Context, vendorVersion string, tokenSource oauth2.TokenSource, computeEndpoint string, computeEnvironment Environment) (*computealpha.Service, error) {
+func createAlphaCloudService(ctx context.Context, vendorVersion string, tokenSource oauth2.TokenSource, computeEndpoint url.URL, 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)
@@ -186,7 +187,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 Environment) (*computebeta.Service, error) {
+func createBetaCloudService(ctx context.Context, vendorVersion string, tokenSource oauth2.TokenSource, computeEndpoint url.URL, 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)
@@ -199,7 +200,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 Environment) (*compute.Service, error) {
+func createCloudService(ctx context.Context, vendorVersion string, tokenSource oauth2.TokenSource, computeEndpoint url.URL, 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)
@@ -212,32 +213,25 @@ func createCloudService(ctx context.Context, vendorVersion string, tokenSource o
 	return service, nil
 }
 
-func getComputeVersion(ctx context.Context, tokenSource oauth2.TokenSource, computeEndpoint string, computeEnvironment Environment, computeVersion Version) ([]option.ClientOption, error) {
+func getComputeVersion(ctx context.Context, tokenSource oauth2.TokenSource, computeEndpoint url.URL, computeEnvironment Environment, computeVersion Version) ([]option.ClientOption, error) {
 	client, err := newOauthClient(ctx, tokenSource)
 	if err != nil {
 		return nil, err
 	}
-	computeEnvironmentSuffix := getPath(computeEnvironment, computeVersion)
 	computeOpts := []option.ClientOption{option.WithHTTPClient(client)}
 
-	if computeEndpoint != "" {
-		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)
-		}
+	if computeEndpoint.String() != "" {
+		computeEnvironmentSuffix := constructComputeEndpointPath(computeEnvironment, computeVersion)
+		computeEndpoint.Path = computeEnvironmentSuffix
+		endpoint := computeEndpoint.String()
 		computeOpts = append(computeOpts, option.WithEndpoint(endpoint))
 	}
 	return computeOpts, nil
 }
 
-func getPath(env Environment, version Version) string {
+func constructComputeEndpointPath(env Environment, version Version) string {
 	prefix := ""
-	if env == environmentStaging {
+	if env == EnvironmentStaging {
 		prefix = fmt.Sprintf("%s_", env)
 	}
 	return fmt.Sprintf("compute/%s%s/", prefix, version)

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

@@ -22,11 +22,13 @@ import (
 	"errors"
 	"fmt"
 	"net/http"
+	"net/url"
 	"testing"
 	"time"
 
 	"golang.org/x/oauth2"
 
+	"google.golang.org/api/compute/v1"
 	"google.golang.org/api/googleapi"
 )
 
@@ -102,7 +104,7 @@ func TestIsGCEError(t *testing.T) {
 func TestGetComputeVersion(t *testing.T) {
 	testCases := []struct {
 		name               string
-		computeEndpoint    string
+		computeEndpoint    url.URL
 		computeEnvironment Environment
 		computeVersion     Version
 		expectedEndpoint   string
@@ -111,30 +113,23 @@ func TestGetComputeVersion(t *testing.T) {
 
 		{
 			name:               "check for production environment",
-			computeEndpoint:    "https://compute.googleapis.com",
-			computeEnvironment: "production",
-			computeVersion:     "v1",
-			expectedEndpoint:   "https://compute.googleapis.com/compute/v1/",
+			computeEndpoint:    convertStringToURL("https://compute.googleapis.com"),
+			computeEnvironment: EnvironmentProduction,
+			computeVersion:     versionBeta,
+			expectedEndpoint:   "https://compute.googleapis.com/compute/beta/",
 			expectError:        false,
 		},
-		{
-			name:               "check for incorrect endpoint",
-			computeEndpoint:    "https://compute.googleapis",
-			computeEnvironment: "prod",
-			computeVersion:     "v1",
-			expectError:        true,
-		},
 		{
 			name:               "check for staging environment",
-			computeEndpoint:    "https://compute.googleapis.com",
-			computeEnvironment: environmentStaging,
-			computeVersion:     "v1",
-			expectedEndpoint:   "compute/staging_v1/",
+			computeEndpoint:    convertStringToURL("https://compute.googleapis.com"),
+			computeEnvironment: EnvironmentStaging,
+			computeVersion:     versionV1,
+			expectedEndpoint:   "https://compute.googleapis.com/compute/staging_v1/",
 			expectError:        false,
 		},
 		{
 			name:               "check for random string as endpoint",
-			computeEndpoint:    "compute-googleapis",
+			computeEndpoint:    url.URL{},
 			computeEnvironment: "prod",
 			computeVersion:     "v1",
 			expectedEndpoint:   "compute/v1/",
@@ -143,10 +138,23 @@ func TestGetComputeVersion(t *testing.T) {
 	}
 	for _, tc := range testCases {
 		ctx := context.Background()
-		_, err := getComputeVersion(ctx, &mockTokenSource{}, tc.computeEndpoint, tc.computeEnvironment, tc.computeVersion)
+		computeOpts, err := getComputeVersion(ctx, &mockTokenSource{}, tc.computeEndpoint, tc.computeEnvironment, tc.computeVersion)
+		service, _ := compute.NewService(ctx, computeOpts...)
+		gotEndpoint := service.BasePath
 		if err != nil && !tc.expectError {
-			t.Fatalf("Got error %v, expected endpoint %s", err, tc.expectedEndpoint)
+			t.Fatalf("Got error %v", err)
+		}
+		if gotEndpoint != tc.expectedEndpoint && !tc.expectError {
+			t.Fatalf("expected endpoint %s, got endpoint %s", tc.expectedEndpoint, gotEndpoint)
 		}
 	}
 
 }
+
+func convertStringToURL(urlString string) url.URL {
+	parsedURL, err := url.ParseRequestURI(urlString)
+	if err != nil {
+		return url.URL{}
+	}
+	return *parsedURL
+}